Download.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FunctionalDependencies #-}

{- | A high-level interface to downloading share data as bytes from storage
 servers.
-}
module Tahoe.Download (
    LookupServer,
    DownloadError (..),
    LookupError (..),
    DiscoverError (..),
    discoverShares,
    download,
    announcementToImmutableStorageServer,
    getShareNumbers,
) where

import Control.Exception (Exception (displayException), SomeException, try)
import Control.Monad.IO.Class (MonadIO (liftIO))
import Data.Bifunctor (Bifunctor (first, second))
import Data.Binary (Word16, decodeOrFail)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as LB
import Data.Either (partitionEithers, rights)
import Data.List (foldl')
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import Tahoe.Announcement (StorageServerAnnouncement)
import qualified Tahoe.CHK
import Tahoe.CHK.Capability (Reader (..), Verifier (..))
import qualified Tahoe.CHK.Capability as CHK
import qualified Tahoe.CHK.Encrypt
import Tahoe.CHK.Server (StorageServer (..), StorageServerID)
import qualified Tahoe.CHK.Share
import Tahoe.CHK.Types (ShareNum, StorageIndex)
import Tahoe.Download.Internal.Client
import Tahoe.Download.Internal.Immutable

print' :: MonadIO m => String -> m ()
-- print' = liftIO . print
print' = const $ pure ()

-- | Partially describe one share download.
type DownloadTask = (ShareNum, StorageServer)

-- | A downloaded share
type Share = (ShareNum, LB.ByteString)

{- | Recover the application data associated with a given capability from the
 given servers, if possible.
-}
download ::
    (MonadIO m, Verifiable v, Readable r v) =>
    -- | Information about the servers from which to consider downloading shares
    -- representing the application data.
    Map.Map StorageServerID StorageServerAnnouncement ->
    -- | The read capability for the application data.
    r ->
    -- | Get functions for interacting with a server given its URL.
    LookupServer m ->
    -- | Either a description of how the recovery failed or the recovered
    -- application data.
    m (Either DownloadError LB.ByteString)
download servers cap lookupServer = do
    print' ("Going to download: " <> show (getStorageIndex $ getVerifiable cap))
    let verifier = getVerifiable cap
    let storageIndex = getStorageIndex verifier
    -- TODO: If getRequiredTotal fails on the first storage server, we may
    -- need to try more.  If it fails for all of them, we need to represent
    -- the failure coherently.
    ss <- firstStorageServer (Map.elems servers) lookupServer
    (required, _) <- getRequiredTotal verifier ss
    locationE <- locateShares servers lookupServer storageIndex (fromIntegral required)
    print' "Finished locating shares"
    case locationE of
        Left err -> do
            print' "Got an error locating shares"
            pure $ Left err
        Right discovered -> do
            print' "Found some shares, fetching them"
            -- XXX note shares can contain failures
            shares <- executeDownloadTasks storageIndex (makeDownloadTasks =<< discovered)
            print' "Fetched the shares, decoding them"
            s <- decodeShares cap shares required
            print' "Decoded them"
            pure s

-- We also need "first successful share"!
firstStorageServer :: Monad m => [StorageServerAnnouncement] -> LookupServer m -> m StorageServer
firstStorageServer servers finder = do
    responses <- mapM finder servers
    pure $ head $ take 1 $ rights responses -- XXX don't do this at home kids, head isn't safe

-- | A capability which confers the ability to locate and verify some stored data.
class Verifiable v where
    -- | Ask a storage server which share numbers related to this capability it
    -- is holding.  This is an unverified result and the storage server could
    -- present incorrect information.  Even if it correctly reports that it
    -- holds a share, it could decline to give it out when asked.
    getShareNumbers :: MonadIO m => v -> StorageServer -> m (Set.Set ShareNum)

    -- | Get the encoding parameters used for the shares of this capability.
    -- The information is presented as a tuple of (required, total).
    getRequiredTotal :: MonadIO m => v -> StorageServer -> m (Int, Int)

    -- | Get the location information for shares of this capability.
    getStorageIndex :: v -> StorageIndex

instance Verifiable CHK.Verifier where
    getShareNumbers v s = liftIO $ storageServerGetBuckets s (storageIndex v)
    getStorageIndex Verifier{storageIndex} = storageIndex

    -- CHK is pure, we don't have to ask the StorageServer
    getRequiredTotal Verifier{required, total} _ = pure (fromIntegral required, fromIntegral total)

{- | A capability which confers the ability to interpret some stored data to
 recover the original plaintext.  Additionally, it can be attentuated to a
 Verifiable.
-}
class (Verifiable v) => Readable r v | r -> v where
    -- | Attentuate the capability.
    getVerifiable :: r -> v

    -- | Interpret the required number of shares to recover the plaintext.
    --
    -- Note: might want to split the two functions below out of decodeShare
    --
    -- shareToCipherText :: r -> [(Int, WhichShare)] -> LB.ByteString
    --
    -- cipherTextToPlainText :: r -> LB.ByteString -> LB.ByteString
    decodeShare :: MonadIO m => r -> [(Int, WhichShare)] -> m (Either DownloadError LB.ByteString)

instance Readable CHK.Reader CHK.Verifier where
    getVerifiable = verifier
    decodeShare r shareList = do
        cipherText <- liftIO $ Tahoe.CHK.decode r (second unWhich <$> shareList)
        case cipherText of
            Nothing -> pure $ Left ShareDecodingFailed
            Just ct ->
                pure . Right $ Tahoe.CHK.Encrypt.decrypt (readKey r) ct

{- | Represent the kind of share to operate on.  This forms a closed world of
 share types.  It might eventually be interesting to make an open world
 variation instead.
-}
newtype WhichShare = CHK {unWhich :: Tahoe.CHK.Share.Share} -- \| SDMF SDMF.Share

{- | Deserialize some bytes representing some kind of share to that kind of
 share, if possible.
-}
bytesToShare :: LB.ByteString -> Either DeserializeError WhichShare
bytesToShare bytes = do
    case decodeOrFail bytes of
        Left _ -> Left UnknownDeserializeError
        Right (_, _, r) -> Right $ CHK r

{- | Execute each download task sequentially and return only the successful
 results.
-}
executeDownloadTasks ::
    MonadIO m =>
    -- | The storage index of the shares to download.
    StorageIndex ->
    -- | The downloads to attempt.
    [DownloadTask] ->
    -- | The results of all successful downloads.
    m [Share]
executeDownloadTasks storageIndex tasks = do
    downloadResults <- mapM (downloadShare storageIndex) tasks
    pure . rights $ inject <$> downloadResults
  where
    inject (a, b) = (a,) <$> b

-- | Find out which servers claim to have shares related to a given storage index.
locateShares ::
    MonadIO m =>
    -- | Information about the servers from which to consider downloading shares
    -- representing the application data.
    Map.Map StorageServerID StorageServerAnnouncement ->
    -- | Get functions for interacting with a server given its URL.
    LookupServer m ->
    -- | The storage index about which to retrieve information.
    B.ByteString ->
    -- | The number of shares we need to locate.  If we cannot find at least
    -- this many shares the result will be an error.
    Word16 ->
    -- | Either an error or a guide to where shares are placed.
    m (Either DownloadError [(StorageServer, Set.Set ShareNum)])
locateShares servers lookupServer storageIndex required =
    case Map.toList servers of
        [] -> pure . Left $ NoConfiguredServers
        serverList -> do
            print' "Discovering shares"
            -- Ask each server for all shares it has.
            ( problems :: [DiscoverError]
                , discovered :: [(StorageServer, Set.Set ShareNum)]
                ) <-
                partitionEithers <$> mapM (discoverShares lookupServer storageIndex) serverList
            if null discovered
                then pure . Left . NoReachableServers $ problems
                else
                    if (fromIntegral required >) . countDistinctShares $ discovered
                        then pure $ Left NotEnoughShares{notEnoughSharesNeeded = fromIntegral required, notEnoughSharesFound = countDistinctShares discovered}
                        else pure $ Right discovered

{- | Given the results of downloading shares related to a given capability,
 decode them and decrypt the contents of possible.
-}
decodeShares ::
    (MonadIO m, Readable r v) =>
    -- | The read capability which allows the contents to be decrypted.
    r ->
    -- | The results of downloading the shares.
    [Share] ->
    Int ->
    m (Either DownloadError LB.ByteString)
decodeShares r shares required = do
    -- Filter down to shares we actually got.
    let fewerShares :: [(ShareNum, Either DeserializeError WhichShare)] = second bytesToShare <$> shares
        onlyDecoded = rights $ (\(a, b) -> (fromIntegral a,) <$> b) <$> fewerShares
    if length onlyDecoded < required
        then pure $ Left NotEnoughDecodedShares{notEnoughDecodedSharesNeeded = fromIntegral required, notEnoughDecodedSharesFound = length onlyDecoded}
        else do
            decodeShare r onlyDecoded

{- | Figure the total number of distinct shares reported by all of the servers
 we asked.
-}
countDistinctShares :: Ord b => [(a, Set.Set b)] -> Int
countDistinctShares = Set.size . foldl' Set.union mempty . map snd

{- | Ask one server which shares it has related to the storage index in
 question.
-}
discoverShares ::
    MonadIO m =>
    LookupServer m ->
    StorageIndex ->
    (StorageServerID, StorageServerAnnouncement) ->
    m (Either DiscoverError (StorageServer, Set.Set ShareNum))
discoverShares lookupServer storageIndex (_sid, sann) = do
    print' "Looking up server from announcement"
    server <- lookupServer sann
    print' "Looked it up"
    case server of
        Left e -> pure . Left . StorageServerUnreachable $ e
        Right ss@StorageServer{storageServerGetBuckets} -> do
            print' $ "Getting buckets for " <> show storageIndex
            buckets <- liftIO $ try (storageServerGetBuckets storageIndex)
            let massaged = first (StorageServerCommunicationError . (displayException :: SomeException -> String)) buckets
            print' $ "Got them " <> show massaged
            pure $ (ss,) <$> massaged

{- | Expand a one-to-many mapping into a list of pairs with each of the "many"
   values as the first element and the corresponding "one" value as the second
   element.
-}
makeDownloadTasks :: Ord k => (v, Set.Set k) -> [(k, v)]
makeDownloadTasks (v, ks) = zip (Set.toList ks) (repeat v)

-- | Download the bytes of a share from one (or more!) of the given servers.
downloadShare ::
    MonadIO m =>
    -- | The storage index of the share to download.
    StorageIndex ->
    -- | Addressing information about the share to download.
    DownloadTask ->
    -- | The bytes of the share or some error that was encountered during
    -- download.
    m (ShareNum, Either DownloadError LB.ByteString)
downloadShare storageIndex (shareNum, s) = do
    print' $ "Going to download " <> show storageIndex <> " " <> show shareNum
    shareBytes <- liftIO $ try (storageServerRead s storageIndex shareNum)
    let massaged = first (ShareDownloadError . (displayException :: SomeException -> String)) shareBytes
    print' "Downloaded it"
    pure (shareNum, LB.fromStrict <$> massaged)