{-# LANGUAGE RecordWildCards #-}

module Tahoe.CHK.Server where

import Data.Aeson (
    FromJSON (..),
    ToJSON (..),
    object,
    withObject,
    (.:),
    (.:?),
    (.=),
 )
import qualified Data.ByteString as B
import Data.ByteString.Base32 (decodeBase32Unpadded, encodeBase32Unpadded)
import Data.List (sortOn)
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import qualified Data.Text as T
import Data.Text.Encoding (encodeUtf8)
import Tahoe.CHK.Crypto (sha1, sha256)
import Tahoe.CHK.Types (Offset, ShareNum, StorageIndex)

-- Where can a server be found
type URL = T.Text

-- The unique identifier for a particular storage server, conventionally the
-- lowercase base32 encoding of some public key controlled by the server.
type StorageServerID = T.Text

-- | An announcement from a storage server about its storage service.
data StorageServerAnnouncement = StorageServerAnnouncement
    { StorageServerAnnouncement -> Maybe URL
storageServerAnnouncementFURL :: Maybe URL
    , StorageServerAnnouncement -> Maybe URL
storageServerAnnouncementNick :: Maybe T.Text
    , StorageServerAnnouncement -> Maybe ByteString
storageServerAnnouncementPermutationSeed :: Maybe B.ByteString
    }
    deriving (StorageServerAnnouncement -> StorageServerAnnouncement -> Bool
(StorageServerAnnouncement -> StorageServerAnnouncement -> Bool)
-> (StorageServerAnnouncement -> StorageServerAnnouncement -> Bool)
-> Eq StorageServerAnnouncement
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: StorageServerAnnouncement -> StorageServerAnnouncement -> Bool
$c/= :: StorageServerAnnouncement -> StorageServerAnnouncement -> Bool
== :: StorageServerAnnouncement -> StorageServerAnnouncement -> Bool
$c== :: StorageServerAnnouncement -> StorageServerAnnouncement -> Bool
Eq, Eq StorageServerAnnouncement
Eq StorageServerAnnouncement
-> (StorageServerAnnouncement
    -> StorageServerAnnouncement -> Ordering)
-> (StorageServerAnnouncement -> StorageServerAnnouncement -> Bool)
-> (StorageServerAnnouncement -> StorageServerAnnouncement -> Bool)
-> (StorageServerAnnouncement -> StorageServerAnnouncement -> Bool)
-> (StorageServerAnnouncement -> StorageServerAnnouncement -> Bool)
-> (StorageServerAnnouncement
    -> StorageServerAnnouncement -> StorageServerAnnouncement)
-> (StorageServerAnnouncement
    -> StorageServerAnnouncement -> StorageServerAnnouncement)
-> Ord StorageServerAnnouncement
StorageServerAnnouncement -> StorageServerAnnouncement -> Bool
StorageServerAnnouncement -> StorageServerAnnouncement -> Ordering
StorageServerAnnouncement
-> StorageServerAnnouncement -> StorageServerAnnouncement
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: StorageServerAnnouncement
-> StorageServerAnnouncement -> StorageServerAnnouncement
$cmin :: StorageServerAnnouncement
-> StorageServerAnnouncement -> StorageServerAnnouncement
max :: StorageServerAnnouncement
-> StorageServerAnnouncement -> StorageServerAnnouncement
$cmax :: StorageServerAnnouncement
-> StorageServerAnnouncement -> StorageServerAnnouncement
>= :: StorageServerAnnouncement -> StorageServerAnnouncement -> Bool
$c>= :: StorageServerAnnouncement -> StorageServerAnnouncement -> Bool
> :: StorageServerAnnouncement -> StorageServerAnnouncement -> Bool
$c> :: StorageServerAnnouncement -> StorageServerAnnouncement -> Bool
<= :: StorageServerAnnouncement -> StorageServerAnnouncement -> Bool
$c<= :: StorageServerAnnouncement -> StorageServerAnnouncement -> Bool
< :: StorageServerAnnouncement -> StorageServerAnnouncement -> Bool
$c< :: StorageServerAnnouncement -> StorageServerAnnouncement -> Bool
compare :: StorageServerAnnouncement -> StorageServerAnnouncement -> Ordering
$ccompare :: StorageServerAnnouncement -> StorageServerAnnouncement -> Ordering
$cp1Ord :: Eq StorageServerAnnouncement
Ord, Int -> StorageServerAnnouncement -> ShowS
[StorageServerAnnouncement] -> ShowS
StorageServerAnnouncement -> String
(Int -> StorageServerAnnouncement -> ShowS)
-> (StorageServerAnnouncement -> String)
-> ([StorageServerAnnouncement] -> ShowS)
-> Show StorageServerAnnouncement
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [StorageServerAnnouncement] -> ShowS
$cshowList :: [StorageServerAnnouncement] -> ShowS
show :: StorageServerAnnouncement -> String
$cshow :: StorageServerAnnouncement -> String
showsPrec :: Int -> StorageServerAnnouncement -> ShowS
$cshowsPrec :: Int -> StorageServerAnnouncement -> ShowS
Show)

-- A server that can have some data uploaded to it.
data StorageServer = StorageServer
    { StorageServer -> URL
storageServerID :: StorageServerID
    , -- TODO Strict byte strings here are unfortunate.  They will force whole
      -- chunks of data into memory at once.
      StorageServer -> ByteString -> Int -> Offset -> ByteString -> IO ()
storageServerWrite :: StorageIndex -> ShareNum -> Offset -> B.ByteString -> IO ()
    , StorageServer -> ByteString -> Int -> IO ByteString
storageServerRead :: StorageIndex -> ShareNum -> IO B.ByteString
    , StorageServer -> ByteString -> IO (Set Int)
storageServerGetBuckets :: StorageIndex -> IO (Set.Set ShareNum)
    }

instance Eq StorageServer where
    StorageServer
a == :: StorageServer -> StorageServer -> Bool
== StorageServer
b = StorageServer -> URL
storageServerID StorageServer
a URL -> URL -> Bool
forall a. Eq a => a -> a -> Bool
== StorageServer -> URL
storageServerID StorageServer
b

instance Ord StorageServer where
    StorageServer
a <= :: StorageServer -> StorageServer -> Bool
<= StorageServer
b = StorageServer -> URL
storageServerID StorageServer
a URL -> URL -> Bool
forall a. Ord a => a -> a -> Bool
<= StorageServer -> URL
storageServerID StorageServer
b

instance Show StorageServer where
    show :: StorageServer -> String
show StorageServer
ss = URL -> String
forall a. Show a => a -> String
show (URL -> String) -> URL -> String
forall a b. (a -> b) -> a -> b
$ StorageServer -> URL
storageServerID StorageServer
ss

type ShareMap = Map.Map ShareNum (Set.Set StorageServer)

instance FromJSON StorageServerAnnouncement where
    parseJSON :: Value -> Parser StorageServerAnnouncement
parseJSON = String
-> (Object -> Parser StorageServerAnnouncement)
-> Value
-> Parser StorageServerAnnouncement
forall a. String -> (Object -> Parser a) -> Value -> Parser a
withObject String
"StorageServerAnnouncement" ((Object -> Parser StorageServerAnnouncement)
 -> Value -> Parser StorageServerAnnouncement)
-> (Object -> Parser StorageServerAnnouncement)
-> Value
-> Parser StorageServerAnnouncement
forall a b. (a -> b) -> a -> b
$ \Object
ann -> do
        Object
v <- Object
ann Object -> Key -> Parser Object
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"ann"
        Maybe URL
storageServerAnnouncementFURL <- Object
v Object -> Key -> Parser (Maybe URL)
forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
.:? Key
"anonymous-storage-FURL"
        Maybe URL
storageServerAnnouncementNick <- Object
v Object -> Key -> Parser (Maybe URL)
forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
.:? Key
"nickname"
        Maybe URL
permutationSeed <- Object
v Object -> Key -> Parser (Maybe URL)
forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
.:? Key
"permutation-seed-base32"
        let storageServerAnnouncementPermutationSeed :: Maybe ByteString
storageServerAnnouncementPermutationSeed =
                case Maybe URL
permutationSeed of
                    Maybe URL
Nothing -> Maybe ByteString
forall a. Maybe a
Nothing
                    Just URL
txt -> case ByteString -> Either URL ByteString
decodeBase32Unpadded (ByteString -> Either URL ByteString)
-> (URL -> ByteString) -> URL -> Either URL ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. URL -> ByteString
encodeUtf8 (URL -> Either URL ByteString) -> URL -> Either URL ByteString
forall a b. (a -> b) -> a -> b
$ URL
txt of
                        Left URL
_ -> Maybe ByteString
forall a. Maybe a
Nothing
                        Right ByteString
ps -> ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
ps

        StorageServerAnnouncement -> Parser StorageServerAnnouncement
forall (f :: * -> *) a. Applicative f => a -> f a
pure StorageServerAnnouncement :: Maybe URL
-> Maybe URL -> Maybe ByteString -> StorageServerAnnouncement
StorageServerAnnouncement{Maybe ByteString
Maybe URL
storageServerAnnouncementPermutationSeed :: Maybe ByteString
storageServerAnnouncementNick :: Maybe URL
storageServerAnnouncementFURL :: Maybe URL
storageServerAnnouncementPermutationSeed :: Maybe ByteString
storageServerAnnouncementNick :: Maybe URL
storageServerAnnouncementFURL :: Maybe URL
..}

instance ToJSON StorageServerAnnouncement where
    toJSON :: StorageServerAnnouncement -> Value
toJSON StorageServerAnnouncement{Maybe ByteString
Maybe URL
storageServerAnnouncementPermutationSeed :: Maybe ByteString
storageServerAnnouncementNick :: Maybe URL
storageServerAnnouncementFURL :: Maybe URL
storageServerAnnouncementPermutationSeed :: StorageServerAnnouncement -> Maybe ByteString
storageServerAnnouncementNick :: StorageServerAnnouncement -> Maybe URL
storageServerAnnouncementFURL :: StorageServerAnnouncement -> Maybe URL
..} =
        [Pair] -> Value
object
            [ Key
"ann"
                Key -> Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= [Pair] -> Value
object
                    [ Key
"anonymous-storage-FURL" Key -> Maybe URL -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Maybe URL
storageServerAnnouncementFURL
                    , Key
"nickname" Key -> Maybe URL -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Maybe URL
storageServerAnnouncementNick
                    , Key
"permutation-seed-base32"
                        Key -> Maybe URL -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= (ByteString -> URL
encodeBase32Unpadded (ByteString -> URL) -> Maybe ByteString -> Maybe URL
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe ByteString
storageServerAnnouncementPermutationSeed)
                    ]
            ]

{- | Find the preferred order of servers for an object with the given index.

 This is like allmydata.storage_client.StorageFarmBroker.get_servers_for_psi
-}
preferredServers :: StorageIndex -> Map.Map T.Text StorageServerAnnouncement -> [(StorageServerID, StorageServerAnnouncement)]
preferredServers :: ByteString
-> Map URL StorageServerAnnouncement
-> [(URL, StorageServerAnnouncement)]
preferredServers ByteString
storageIndex = ((URL, StorageServerAnnouncement) -> ByteString)
-> [(URL, StorageServerAnnouncement)]
-> [(URL, StorageServerAnnouncement)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (URL, StorageServerAnnouncement) -> ByteString
permutedServerHash ([(URL, StorageServerAnnouncement)]
 -> [(URL, StorageServerAnnouncement)])
-> (Map URL StorageServerAnnouncement
    -> [(URL, StorageServerAnnouncement)])
-> Map URL StorageServerAnnouncement
-> [(URL, StorageServerAnnouncement)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map URL StorageServerAnnouncement
-> [(URL, StorageServerAnnouncement)]
forall k a. Map k a -> [(k, a)]
Map.toList
  where
    permutedServerHash :: (URL, StorageServerAnnouncement) -> ByteString
permutedServerHash =
        -- allmydata.util.hashutil.permute_server_hash
        ByteString -> ByteString
sha1 (ByteString -> ByteString)
-> ((URL, StorageServerAnnouncement) -> ByteString)
-> (URL, StorageServerAnnouncement)
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString
storageIndex ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<>) (ByteString -> ByteString)
-> ((URL, StorageServerAnnouncement) -> ByteString)
-> (URL, StorageServerAnnouncement)
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (URL -> StorageServerAnnouncement -> ByteString)
-> (URL, StorageServerAnnouncement) -> ByteString
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry URL -> StorageServerAnnouncement -> ByteString
storageServerPermutationSeed

{- | Compute a sort key for a storage server given its identifier and storage
 service announcement.

 This is like pieces of allmydata.storage_client._parse_announcement
-}
storageServerPermutationSeed :: StorageServerID -> StorageServerAnnouncement -> B.ByteString
storageServerPermutationSeed :: URL -> StorageServerAnnouncement -> ByteString
storageServerPermutationSeed URL
serverId StorageServerAnnouncement
ann =
    case StorageServerAnnouncement -> Maybe ByteString
storageServerAnnouncementPermutationSeed StorageServerAnnouncement
ann of
        Just ByteString
bs -> ByteString
bs
        Maybe ByteString
Nothing ->
            case ByteString -> Either URL ByteString
decodeBase32Unpadded (ByteString -> Either URL ByteString)
-> (URL -> ByteString) -> URL -> Either URL ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. URL -> ByteString
encodeUtf8 (URL -> ByteString) -> (URL -> URL) -> URL -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> URL -> URL
T.drop Int
3 (URL -> Either URL ByteString) -> URL -> Either URL ByteString
forall a b. (a -> b) -> a -> b
$ URL
serverId of
                Right ByteString
bs -> ByteString
bs
                Left URL
_ -> ByteString -> ByteString
sha256 (ByteString -> ByteString)
-> (URL -> ByteString) -> URL -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. URL -> ByteString
encodeUtf8 (URL -> ByteString) -> URL -> ByteString
forall a b. (a -> b) -> a -> b
$ URL
serverId