{-# LANGUAGE TypeFamilies #-}

module Tahoe.Download.Internal.Capability where

import Control.Exception (SomeException, throwIO, try)
import Control.Monad.IO.Class
import Data.Bifunctor (Bifunctor (..))
import Data.Binary (Word8, decodeOrFail)
import Data.Binary.Get (ByteOffset)
import qualified Data.ByteString.Lazy as LB
import Data.Foldable (foldlM)
import qualified Data.Set as Set
import Network.HTTP.Types (Status (statusCode))
import Servant.Client (ClientError (FailureResponse), ResponseF (..))
import qualified Tahoe.CHK
import qualified Tahoe.CHK.Capability as CHK
import qualified Tahoe.CHK.Encrypt
import Tahoe.CHK.Server
import qualified Tahoe.CHK.Share
import Tahoe.CHK.Types
import Tahoe.Download.Internal.Client
import qualified Tahoe.SDMF as SDMF
import qualified Tahoe.SDMF.Internal.Keys as SDMF.Keys

-- | A capability which confers the ability to locate and verify some stored data.
class Verifiable v where
    -- | Represent the type of share to operate on.
    type ShareT v

    -- | 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).

    -- SDMF can fail to figure this out in lots of ways so consider switching
    -- to Either or something?
    getRequiredTotal :: MonadIO m => v -> StorageServer -> m (Maybe (Int, Int))

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

    -- | Deserialize some bytes representing some kind of share to the kind of
    -- share associated with this capability type, if possible.
    deserializeShare ::
        -- | A type witness revealing what type of share to decode to.
        v ->
        -- | The bytes of the serialized share.
        LB.ByteString ->
        Either (LB.ByteString, ByteOffset, String) (ShareT v)

{- | A capability which confers the ability to recover plaintext from
 ciphertext.
-}
class Readable r where
    -- | Represent the type of a Verifiable associated with the Readable.
    type Verifier r

    -- | Attentuate the capability.
    getVerifiable :: r -> Verifier r

    -- | 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, ShareT r)] -> LB.ByteString
    --
    -- cipherTextToPlainText :: r -> LB.ByteString -> LB.ByteString
    decodeShare :: MonadIO m => r -> [(Int, ShareT (Verifier r))] -> m (Either DownloadError LB.ByteString)

instance Verifiable CHK.Verifier where
    type ShareT CHK.Verifier = Tahoe.CHK.Share.Share

    getShareNumbers :: Verifier -> StorageServer -> m (Set ShareNum)
getShareNumbers Verifier
v StorageServer
s = IO (Set ShareNum) -> m (Set ShareNum)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Set ShareNum) -> m (Set ShareNum))
-> IO (Set ShareNum) -> m (Set ShareNum)
forall a b. (a -> b) -> a -> b
$ StorageServer -> StorageIndex -> IO (Set ShareNum)
storageServerGetBuckets StorageServer
s (Verifier -> StorageIndex
CHK.storageIndex Verifier
v)
    getStorageIndex :: Verifier -> StorageIndex
getStorageIndex CHK.Verifier{StorageIndex
storageIndex :: StorageIndex
storageIndex :: Verifier -> StorageIndex
storageIndex} = StorageIndex
storageIndex

    -- CHK is pure, we don't have to ask the StorageServer
    getRequiredTotal :: Verifier -> StorageServer -> m (Maybe (Int, Int))
getRequiredTotal CHK.Verifier{Word16
required :: Verifier -> Word16
required :: Word16
required, Word16
total :: Verifier -> Word16
total :: Word16
total} StorageServer
_ = Maybe (Int, Int) -> m (Maybe (Int, Int))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Int, Int) -> m (Maybe (Int, Int)))
-> Maybe (Int, Int) -> m (Maybe (Int, Int))
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> Maybe (Int, Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
required, Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
total)

    deserializeShare :: Verifier
-> ByteString
-> Either (ByteString, ByteOffset, String) (ShareT Verifier)
deserializeShare Verifier
_ = ((ByteString, ByteOffset, Share) -> Share)
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, Share)
-> Either (ByteString, ByteOffset, String) Share
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(ByteString
_, ByteOffset
_, Share
c) -> Share
c) (Either
   (ByteString, ByteOffset, String) (ByteString, ByteOffset, Share)
 -> Either (ByteString, ByteOffset, String) Share)
-> (ByteString
    -> Either
         (ByteString, ByteOffset, String) (ByteString, ByteOffset, Share))
-> ByteString
-> Either (ByteString, ByteOffset, String) Share
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, Share)
forall a.
Binary a =>
ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
decodeOrFail

{- | A capability which confers the ability to interpret some stored data to
 recover the original plaintext.  Additionally, it can be attentuated to a
 Verifiable.
-}
instance Readable CHK.Reader where
    type Verifier CHK.Reader = CHK.Verifier

    getVerifiable :: Reader -> Verifier Reader
getVerifiable = Reader -> Verifier
Reader -> Verifier Reader
CHK.verifier
    decodeShare :: Reader
-> [(Int, ShareT (Verifier Reader))]
-> m (Either DownloadError ByteString)
decodeShare Reader
r [(Int, ShareT (Verifier Reader))]
shareList = do
        Maybe ByteString
cipherText <- IO (Maybe ByteString) -> m (Maybe ByteString)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe ByteString) -> m (Maybe ByteString))
-> IO (Maybe ByteString) -> m (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ Reader -> [(Int, Share)] -> IO (Maybe ByteString)
Tahoe.CHK.decode Reader
r [(Int, Share)]
[(Int, ShareT (Verifier Reader))]
shareList
        case Maybe ByteString
cipherText of
            Maybe ByteString
Nothing -> Either DownloadError ByteString
-> m (Either DownloadError ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either DownloadError ByteString
 -> m (Either DownloadError ByteString))
-> Either DownloadError ByteString
-> m (Either DownloadError ByteString)
forall a b. (a -> b) -> a -> b
$ DownloadError -> Either DownloadError ByteString
forall a b. a -> Either a b
Left DownloadError
ShareDecodingFailed
            Just ByteString
ct ->
                Either DownloadError ByteString
-> m (Either DownloadError ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either DownloadError ByteString
 -> m (Either DownloadError ByteString))
-> (ByteString -> Either DownloadError ByteString)
-> ByteString
-> m (Either DownloadError ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Either DownloadError ByteString
forall a b. b -> Either a b
Right (ByteString -> m (Either DownloadError ByteString))
-> ByteString -> m (Either DownloadError ByteString)
forall a b. (a -> b) -> a -> b
$ AESKey128 -> ByteString -> ByteString
Tahoe.CHK.Encrypt.decrypt (Reader -> AESKey128
CHK.readKey Reader
r) ByteString
ct

firstJustsM :: (Monad m, Foldable f) => f (m (Maybe a)) -> m (Maybe a)
firstJustsM :: f (m (Maybe a)) -> m (Maybe a)
firstJustsM = (Maybe a -> m (Maybe a) -> m (Maybe a))
-> Maybe a -> f (m (Maybe a)) -> m (Maybe a)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM Maybe a -> m (Maybe a) -> m (Maybe a)
forall (m :: * -> *) a.
Monad m =>
Maybe a -> m (Maybe a) -> m (Maybe a)
go Maybe a
forall a. Maybe a
Nothing
  where
    go :: Monad m => Maybe a -> m (Maybe a) -> m (Maybe a)
    go :: Maybe a -> m (Maybe a) -> m (Maybe a)
go Maybe a
Nothing m (Maybe a)
action = m (Maybe a)
action
    go result :: Maybe a
result@(Just a
_) m (Maybe a)
_action = Maybe a -> m (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
result

instance Verifiable SDMF.Verifier where
    type ShareT SDMF.Verifier = SDMF.Share

    getShareNumbers :: Verifier -> StorageServer -> m (Set ShareNum)
getShareNumbers Verifier
v StorageServer
s = IO (Set ShareNum) -> m (Set ShareNum)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Set ShareNum) -> m (Set ShareNum))
-> IO (Set ShareNum) -> m (Set ShareNum)
forall a b. (a -> b) -> a -> b
$ StorageServer -> StorageIndex -> IO (Set ShareNum)
storageServerGetBuckets StorageServer
s (StorageIndex -> StorageIndex
SDMF.Keys.unStorageIndex (StorageIndex -> StorageIndex) -> StorageIndex -> StorageIndex
forall a b. (a -> b) -> a -> b
$ Verifier -> StorageIndex
SDMF.verifierStorageIndex Verifier
v)
    getStorageIndex :: Verifier -> StorageIndex
getStorageIndex = StorageIndex -> StorageIndex
SDMF.Keys.unStorageIndex (StorageIndex -> StorageIndex)
-> (Verifier -> StorageIndex) -> Verifier -> StorageIndex
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Verifier -> StorageIndex
SDMF.verifierStorageIndex
    getRequiredTotal :: Verifier -> StorageServer -> m (Maybe (Int, Int))
getRequiredTotal SDMF.Verifier{Digest SHA256
StorageIndex
verifierVerificationKeyHash :: Verifier -> Digest SHA256
verifierVerificationKeyHash :: Digest SHA256
verifierStorageIndex :: StorageIndex
verifierStorageIndex :: Verifier -> StorageIndex
..} StorageServer
ss = IO (Maybe (Int, Int)) -> m (Maybe (Int, Int))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe (Int, Int)) -> m (Maybe (Int, Int)))
-> IO (Maybe (Int, Int)) -> m (Maybe (Int, Int))
forall a b. (a -> b) -> a -> b
$ do
        -- Find out what shares it has.  Any share will do but we need to tell
        -- it which we want.
        Either SomeException (Set ShareNum)
errorOrShareNums <- IO (Set ShareNum) -> IO (Either SomeException (Set ShareNum))
forall e a. Exception e => IO a -> IO (Either e a)
try (IO (Set ShareNum) -> IO (Either SomeException (Set ShareNum)))
-> IO (Set ShareNum) -> IO (Either SomeException (Set ShareNum))
forall a b. (a -> b) -> a -> b
$ StorageServer -> StorageIndex -> IO (Set ShareNum)
storageServerGetBuckets StorageServer
ss StorageIndex
storageIndex
        case Set ShareNum -> [ShareNum]
forall a. Set a -> [a]
Set.toList (Set ShareNum -> [ShareNum])
-> Either SomeException (Set ShareNum)
-> Either SomeException [ShareNum]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either SomeException (Set ShareNum)
errorOrShareNums of
            -- Literally anything could go wrong with that...
            Left (SomeException
e :: SomeException) -> SomeException -> IO (Maybe (Int, Int))
forall e a. Exception e => e -> IO a
throwIO SomeException
e
            -- Or the server may have no shares for this storage index.
            Right [] -> Maybe (Int, Int) -> IO (Maybe (Int, Int))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Int, Int)
forall a. Maybe a
Nothing
            -- Or it might have at least one.  Check each in turn, stopping as
            -- soon as we get a result.
            Right [ShareNum]
shareNums -> [IO (Maybe (Int, Int))] -> IO (Maybe (Int, Int))
forall (m :: * -> *) (f :: * -> *) a.
(Monad m, Foldable f) =>
f (m (Maybe a)) -> m (Maybe a)
firstJustsM (ShareNum -> IO (Maybe (Int, Int))
forall (m :: * -> *). MonadIO m => ShareNum -> m (Maybe (Int, Int))
getParams (ShareNum -> IO (Maybe (Int, Int)))
-> [ShareNum] -> [IO (Maybe (Int, Int))]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [ShareNum]
shareNums)
      where
        -- Get the Required, Total parameters for one share number, if
        -- possible.
        getParams :: MonadIO m => Word8 -> m (Maybe (Int, Int))
        getParams :: ShareNum -> m (Maybe (Int, Int))
getParams ShareNum
shareNum = IO (Maybe (Int, Int)) -> m (Maybe (Int, Int))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe (Int, Int)) -> m (Maybe (Int, Int)))
-> IO (Maybe (Int, Int)) -> m (Maybe (Int, Int))
forall a b. (a -> b) -> a -> b
$ do
            Either ClientError StorageIndex
errorOrShareBytes <- IO StorageIndex -> IO (Either ClientError StorageIndex)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO StorageIndex -> IO (Either ClientError StorageIndex))
-> IO StorageIndex -> IO (Either ClientError StorageIndex)
forall a b. (a -> b) -> a -> b
$ StorageServer -> StorageIndex -> ShareNum -> IO StorageIndex
storageServerRead StorageServer
ss StorageIndex
storageIndex ShareNum
shareNum
            case Either ClientError StorageIndex
errorOrShareBytes of
                Left e :: ClientError
e@(FailureResponse RequestF () (BaseUrl, StorageIndex)
_ Response
response) ->
                    -- It should not be very surprising for the requested share to
                    -- be missing from the server (you can never be sure what a
                    -- server will have).  Other issues should probably be kept
                    -- visible.
                    if Int -> Response -> Bool
forall a. Int -> ResponseF a -> Bool
isStatusCode Int
404 Response
response
                        then Maybe (Int, Int) -> IO (Maybe (Int, Int))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Int, Int)
forall a. Maybe a
Nothing
                        else ClientError -> IO (Maybe (Int, Int))
forall e a. Exception e => e -> IO a
throwIO ClientError
e
                Left ClientError
e -> ClientError -> IO (Maybe (Int, Int))
forall e a. Exception e => e -> IO a
throwIO ClientError
e
                Right StorageIndex
shareBytes ->
                    case ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, Share)
forall a.
Binary a =>
ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
decodeOrFail (StorageIndex -> ByteString
LB.fromStrict StorageIndex
shareBytes) of
                        Left (ByteString, ByteOffset, String)
_ -> Maybe (Int, Int) -> IO (Maybe (Int, Int))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Int, Int)
forall a. Maybe a
Nothing
                        Right (ByteString
_, ByteOffset
_, Share
sh) -> Maybe (Int, Int) -> IO (Maybe (Int, Int))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Int, Int) -> IO (Maybe (Int, Int)))
-> Maybe (Int, Int) -> IO (Maybe (Int, Int))
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> Maybe (Int, Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ShareNum -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ShareNum -> Int) -> ShareNum -> Int
forall a b. (a -> b) -> a -> b
$ Share -> ShareNum
SDMF.shareRequiredShares Share
sh, ShareNum -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ShareNum -> Int) -> ShareNum -> Int
forall a b. (a -> b) -> a -> b
$ Share -> ShareNum
SDMF.shareTotalShares Share
sh)

        storageIndex :: StorageIndex
storageIndex = StorageIndex -> StorageIndex
SDMF.Keys.unStorageIndex StorageIndex
verifierStorageIndex

    deserializeShare :: Verifier
-> ByteString
-> Either (ByteString, ByteOffset, String) (ShareT Verifier)
deserializeShare Verifier
_ = ((ByteString, ByteOffset, Share) -> Share)
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, Share)
-> Either (ByteString, ByteOffset, String) Share
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(ByteString
_, ByteOffset
_, Share
c) -> Share
c) (Either
   (ByteString, ByteOffset, String) (ByteString, ByteOffset, Share)
 -> Either (ByteString, ByteOffset, String) Share)
-> (ByteString
    -> Either
         (ByteString, ByteOffset, String) (ByteString, ByteOffset, Share))
-> ByteString
-> Either (ByteString, ByteOffset, String) Share
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, Share)
forall a.
Binary a =>
ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
decodeOrFail

-- | Test the status code of a response for equality against a given value.
isStatusCode :: Int -> ResponseF a -> Bool
isStatusCode :: Int -> ResponseF a -> Bool
isStatusCode Int
expected = (Int
expected Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==) (Int -> Bool) -> (ResponseF a -> Int) -> ResponseF a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Status -> Int
statusCode (Status -> Int) -> (ResponseF a -> Status) -> ResponseF a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ResponseF a -> Status
forall a. ResponseF a -> Status
responseStatusCode

instance Readable SDMF.Reader where
    type Verifier SDMF.Reader = SDMF.Verifier
    getVerifiable :: Reader -> Verifier Reader
getVerifiable = Reader -> Verifier
Reader -> Verifier Reader
SDMF.readerVerifier
    decodeShare :: Reader
-> [(Int, ShareT (Verifier Reader))]
-> m (Either DownloadError ByteString)
decodeShare Reader
r [(Int, ShareT (Verifier Reader))]
shareList = do
        Either Any ByteString
cipherText <- ByteString -> Either Any ByteString
forall a b. b -> Either a b
Right (ByteString -> Either Any ByteString)
-> m ByteString -> m (Either Any ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO ByteString -> m ByteString
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Reader -> [(Word16, Share)] -> IO ByteString
forall (m :: * -> *).
(MonadFail m, MonadIO m) =>
Reader -> [(Word16, Share)] -> m ByteString
SDMF.decode Reader
r ((Int -> Word16) -> (Int, Share) -> (Word16, Share)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Int -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Int, Share) -> (Word16, Share))
-> [(Int, Share)] -> [(Word16, Share)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Int, Share)]
[(Int, ShareT (Verifier Reader))]
shareList))
        case Either Any ByteString
cipherText of
            Left Any
_ -> Either DownloadError ByteString
-> m (Either DownloadError ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either DownloadError ByteString
 -> m (Either DownloadError ByteString))
-> Either DownloadError ByteString
-> m (Either DownloadError ByteString)
forall a b. (a -> b) -> a -> b
$ DownloadError -> Either DownloadError ByteString
forall a b. a -> Either a b
Left DownloadError
ShareDecodingFailed
            Right ByteString
ct -> do
                String -> m ()
forall (m :: * -> *). MonadIO m => String -> m ()
print' (String
"Got some ciphertext: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ByteString -> String
forall a. Show a => a -> String
show ByteString
ct)
                String -> m ()
forall (m :: * -> *). MonadIO m => String -> m ()
print' (String
"Decrypting with iv: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> SDMF_IV -> String
forall a. Show a => a -> String
show SDMF_IV
iv)
                Either DownloadError ByteString
-> m (Either DownloadError ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either DownloadError ByteString
 -> m (Either DownloadError ByteString))
-> (ByteString -> Either DownloadError ByteString)
-> ByteString
-> m (Either DownloadError ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Either DownloadError ByteString
forall a b. b -> Either a b
Right (ByteString -> m (Either DownloadError ByteString))
-> ByteString -> m (Either DownloadError ByteString)
forall a b. (a -> b) -> a -> b
$ Read -> SDMF_IV -> ByteString -> ByteString
SDMF.decrypt Read
readKey SDMF_IV
iv ByteString
ct
              where
                readKey :: Read
readKey = Reader -> Read
SDMF.readerReadKey Reader
r
                iv :: SDMF_IV
iv = Share -> SDMF_IV
SDMF.shareIV ((Int, Share) -> Share
forall a b. (a, b) -> b
snd ((Int, Share) -> Share)
-> ([(Int, Share)] -> (Int, Share)) -> [(Int, Share)] -> Share
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, Share)] -> (Int, Share)
forall a. [a] -> a
head ([(Int, Share)] -> Share) -> [(Int, Share)] -> Share
forall a b. (a -> b) -> a -> b
$ [(Int, Share)]
[(Int, ShareT (Verifier Reader))]
shareList)

print' :: MonadIO m => String -> m ()
-- print' = liftIO . putStrLn
print' :: String -> m ()
print' = m () -> String -> m ()
forall a b. a -> b -> a
const (m () -> String -> m ()) -> m () -> String -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()