Spec.hs

module Main where

import Control.Exception (Exception, throwIO)
import Control.Monad (replicateM, zipWithM_)
import Control.Monad.IO.Class (liftIO)
import Crypto.Cipher.Types (nullIV)
import Crypto.Classes (buildKey)
import qualified Data.Binary as Binary
import qualified Data.ByteString as B
import Data.ByteString.Base32 (encodeBase32Unpadded)
import qualified Data.ByteString.Lazy as BL
import Data.Default.Class (Default (def))
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 Data.Word (Word16)
import qualified Data.Yaml as Yaml
import Generators (genAnnouncements, genParameters, genRSAKeys)
import Hedgehog (MonadGen, diff, forAll, property, tripping)
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import System.IO (hSetEncoding, stderr, stdout, utf8)
import Tahoe.Announcement (
    Announcements,
    StorageServerAnnouncement (..),
    StorageServerID,
    URI (..),
    URIAuth (..),
    parseURI',
 )
import qualified Tahoe.CHK
import Tahoe.CHK.Capability (Reader (..), Verifier (..))
import qualified Tahoe.CHK.Encrypt
import Tahoe.CHK.Server (StorageServer (..))
import Tahoe.CHK.Types (Parameters (..))
import Tahoe.CHK.Upload (getConvergentKey)
import Tahoe.Download (
    DiscoverError (..),
    DownloadError (..),
    LookupError (..),
    LookupServer,
    download,
 )
import qualified Tahoe.SDMF as SDMF
import qualified Tahoe.SDMF.Keys as SDMF.Keys
import Tahoe.Server (memoryStorageServer)
import Test.Tasty (TestTree, defaultMain, testGroup)
import Test.Tasty.HUnit (assertEqual, testCase)
import Test.Tasty.Hedgehog (testProperty)

data PlacementError = RanOutOfPlacementInfo | RanOutOfServers deriving (Eq, Show)
instance Exception PlacementError

{- | Return a new StorageServer like the given one but with a get-buckets
 interface that always throws an IO exception.
-}
breakGetBuckets :: Exception e => e -> StorageServer -> StorageServer
breakGetBuckets exc ss =
    ss
        { storageServerGetBuckets = const $ throwIO exc
        }

{- | Return a new StorageServer like the given one but with a read-share
 interface that always throws an IO exception.
-}
breakRead :: Exception e => e -> StorageServer -> StorageServer
breakRead exc ss =
    ss
        { storageServerRead = \_ _ -> throwIO exc
        }

{- | A completely arbitrary exception that the download implementation can't
 know anything specific about.
-}
data BespokeFailure = BespokeFailure deriving (Show)

instance Exception BespokeFailure

-- | Make an announcement that's real enough to convince a test.
simpleAnnouncement :: T.Text -> T.Text -> (T.Text, StorageServerAnnouncement)
simpleAnnouncement nick furl =
    ( T.concat ["v0-", nick]
    , def
        { storageServerAnnouncementFURL = Just furl
        , storageServerAnnouncementNick = Just nick
        }
    )

{- | Build a lookup function that can look up any server in the given list
 from its announcement.
-}
simpleLookup :: Applicative f => [(T.Text, b)] -> StorageServerAnnouncement -> f (Either LookupError b)
simpleLookup [] _ = pure . Left $ AnnouncementStructureUnmatched
simpleLookup ((furl, server) : ss) ann@StorageServerAnnouncement{storageServerAnnouncementFURL} =
    if Just furl == storageServerAnnouncementFURL
        then pure . pure $ server
        else simpleLookup ss ann

tests :: TestTree
tests =
    testGroup
        "All tests"
        [ testCase "Tahoe-LAFS fURLs can be parsed to a structured representation" $
            let tubid = "gnuer2axzoq3ggnn7gjoybmfqsjvaow3"
                swissnum = "sxytycucj5eeunlx6modfazq5byp2hpb"
             in assertEqual
                    "The result is as expected"
                    ( Just
                        URI
                            { uriScheme = "pb:"
                            , uriAuthority =
                                Just
                                    URIAuth
                                        { uriUserInfo = tubid <> "@"
                                        , uriRegName = "localhost"
                                        , uriPort = ":46185"
                                        }
                            , uriPath = "/" <> swissnum
                            , uriQuery = ""
                            , uriFragment = ""
                            }
                    )
                    (parseURI' $ T.pack $ "pb://" <> tubid <> "@tcp:localhost:46185/" <> swissnum)
        , testProperty "Announcements round-trip through YAML encoding/decoding" $
            property $ do
                announcements <- forAll $ genAnnouncements (Range.linear 0 3)
                tripping announcements Yaml.encode (Yaml.decodeThrow :: B.ByteString -> Maybe Announcements)
        , testCase
            "no configured servers"
            $ do
                -- If there are no servers then we can't possibly get enough
                -- shares to recover the application data.
                result <- liftIO $ download mempty (trivialCap 1 1) noServers
                assertEqual
                    "download should fail with no servers"
                    (Left NoConfiguredServers)
                    result
        , testCase "no reachable servers" $ do
            -- If we can't contact any configured server then we can't
            -- possibly get enough shares to recover the application data.
            let ann = def{storageServerAnnouncementNick = Just "unreachable"}
                anns =
                    Map.fromList
                        [ ("v0-abc123", ann)
                        ]

            result <- liftIO $ download anns (trivialCap 1 1) noServers
            assertEqual
                "download should fail with no reachable servers"
                (Left $ NoReachableServers [StorageServerUnreachable (URIParseError ann)])
                result
        , testCase "not enough shares" $ do
            -- If we can't recover enough shares from the configured servers
            -- then we can't possibly get enough shares to recover the
            -- application data.
            let anns = Map.fromList [simpleAnnouncement "abc123" "somewhere"]
                cap = trivialCap 3 3

            -- Two shares exist.
            server <- memoryStorageServer
            storageServerWrite server (storageIndex . verifier $ cap) 0 0 "Hello world"
            storageServerWrite server (storageIndex . verifier $ cap) 1 0 "Hello world"

            -- Make the server reachable.
            let openServer = simpleLookup [("somewhere", server)]

            -- Try to download the cap which requires three shares to reconstruct.
            result <- liftIO $ download anns cap openServer
            assertEqual
                "download should fail with not enough shares"
                (Left NotEnoughShares{notEnoughSharesNeeded = 3, notEnoughSharesFound = 2})
                result
        , testCase "not enough distinct shares" $ do
            -- If we can't recover enough *distinct* shares from the
            -- configured servers then we can't possibly get enough shares to
            -- recover the application data.  Duplicate shares do us no good.
            let anns =
                    Map.fromList
                        [ simpleAnnouncement "abc123" "somewhere"
                        , simpleAnnouncement "abc456" "elsewhere"
                        ]
                cap = trivialCap 3 3

            -- Three shares exist
            somewhere <- memoryStorageServer
            let idx = storageIndex . verifier $ cap
                offset = 0
            storageServerWrite somewhere idx 0 offset "Hello world"
            storageServerWrite somewhere idx 1 offset "Hello world"
            -- But this one is just a duplicate of share 0 on the other
            -- server.
            elsewhere <- memoryStorageServer
            storageServerWrite elsewhere idx 0 offset "Hello world"

            -- Make the server reachable.
            let openServer = simpleLookup [("somewhere", somewhere), ("elsewhere", elsewhere)]

            -- Try to download the cap which requires three shares to reconstruct.
            result <- liftIO $ download anns cap openServer
            assertEqual
                "download should fail with not enough shares"
                (Left NotEnoughShares{notEnoughSharesNeeded = 3, notEnoughSharesFound = 2})
                result
        , testCase "IO exceptions from storageServerGetBuckets are handled" $ do
            -- An announcement for our server
            let anns = Map.fromList [simpleAnnouncement "abc123" "somewhere"]
            -- A broken interface to the server
            server <- breakGetBuckets BespokeFailure <$> memoryStorageServer

            -- Make the server reachable.
            let openServer = simpleLookup [("somewhere", server)]

            -- Something to pretend to try to download
            let cap = trivialCap 3 13

            -- Try to download the cap which requires three shares to reconstruct.
            result <- liftIO $ download anns cap openServer
            assertEqual
                "download should fail with details about unreachable server"
                (Left (NoReachableServers [StorageServerCommunicationError "BespokeFailure"]))
                result
        , testCase "IO exceptions from storageServerRead are handled" $ do
            -- An announcement for our server
            let anns = Map.fromList [simpleAnnouncement "abc123" "somewhere"]

            -- A broken interface to the server
            server <- breakRead BespokeFailure <$> memoryStorageServer

            -- Something to pretend to try to download
            let cap = trivialCap 3 13

            -- Three shares exist
            let idx = storageIndex . verifier $ cap
                offset = 0
            storageServerWrite server idx 0 offset "Hello world"
            storageServerWrite server idx 1 offset "Hello world"
            storageServerWrite server idx 2 offset "Hello world"

            -- Make the server reachable.
            let openServer = simpleLookup [("somewhere", server)]

            -- Try to download the cap which requires three shares to reconstruct.

            result <- liftIO $ download anns cap openServer
            assertEqual
                "download should fail with details about unreachable server"
                (Left (NotEnoughDecodedShares{notEnoughDecodedSharesNeeded = 3, notEnoughDecodedSharesFound = 0}))
                result
        , testProperty "chk success" $
            property $ do
                -- If we can recover enough distinct, decodeable shares from the
                -- configured servers then we can recover the application data.

                -- Generates configurations where it should be possible to recover
                -- the data (have all the shares, have enough of the shares,
                -- spread them across many servers, concentrate them on one or a
                -- few, etc)

                secret <- forAll $ Gen.bytes (Range.singleton 32)
                plaintext <- forAll $ BL.fromStrict <$> Gen.bytes (Range.exponential 56 1024)
                params@Parameters{paramTotalShares} <- forAll genParameters

                -- Since multiple shares can be placed on a single server, as long
                -- as we have one server we have a valid case.  Since some shares
                -- might be placed non-optimally it is also nice to allow for some
                -- empty servers so allow for that as well.
                let numServers = Range.linear 1 (fromIntegral paramTotalShares + 1)
                    serverIDs = Gen.text (Range.singleton 2) Gen.ascii
                serverIDs' <- forAll $ Gen.set numServers serverIDs

                -- Choose a share distribution.  Each element of the resulting
                -- list tells us how many shares to place on the next server, for
                -- some arbitrary (stable) server ordering.
                perServerShareCount <-
                    forAll $
                        genListWithSum (length serverIDs') (fromIntegral paramTotalShares)

                -- Make the servers.
                servers <- liftIO $ replicateM (length serverIDs') memoryStorageServer

                -- Encrypt and encode the data into shares.
                let key = getConvergentKey secret params plaintext
                    ciphertext = Tahoe.CHK.Encrypt.encrypt key plaintext
                (shares, cap) <- liftIO $ Tahoe.CHK.encode key params ciphertext

                -- Distribute the shares.
                liftIO $ placeShares (storageIndex . verifier $ cap) (Binary.encode <$> shares) perServerShareCount servers 0

                let serverMap = Map.fromList $ zip (Set.toList serverIDs') servers
                    lookupServer = someServers serverMap
                    serverAnnouncements = Map.fromSet makeAnn serverIDs'

                -- Recover the plaintext from the servers.
                result <- liftIO $ download serverAnnouncements cap lookupServer
                diff (Right plaintext) (==) result
        , testProperty "ssk success" $
            property $ do
                -- Like "chk success" above, but for SDMF (a case of SSK).
                plaintext <- forAll $ BL.fromStrict <$> Gen.bytes (Range.exponential 56 1024)
                sequenceNumber <- forAll $ Gen.integral (Range.exponential 1 10000)
                keypair <- SDMF.Keys.KeyPair <$> forAll genRSAKeys
                Parameters{paramRequiredShares = required, paramTotalShares = total} <- forAll genParameters

                -- Since multiple shares can be placed on a single server, as long
                -- as we have one server we have a valid case.  Since some shares
                -- might be placed non-optimally it is also nice to allow for some
                -- empty servers so allow for that as well.
                let numServers = Range.linear 1 (fromIntegral total + 1)
                    serverIDs = Gen.text (Range.singleton 2) Gen.ascii
                serverIDs' <- forAll $ Gen.set numServers serverIDs

                perServerShareCount <-
                    forAll $
                        genListWithSum (length serverIDs') (fromIntegral total)

                -- Make the servers.
                servers <- liftIO $ replicateM (length serverIDs') memoryStorageServer

                -- Derive the keys, encode the data.
                let -- Not a very good IV choice in reality but it's okay for
                    -- this test where confidentiality and key secrecy is not
                    -- particularly a concern.
                    iv = SDMF.Keys.SDMF_IV nullIV
                    signatureKey = SDMF.Keys.toSignatureKey keypair
                    (Just writeKey) = SDMF.Keys.deriveWriteKey signatureKey
                    (Just readKey) = SDMF.Keys.deriveReadKey writeKey
                    (Just dataKey) = SDMF.Keys.deriveDataKey iv readKey
                    ciphertext = SDMF.encrypt dataKey plaintext
                (shares, writeCap) <- liftIO $ SDMF.encode keypair sequenceNumber required total ciphertext
                let readCap = SDMF.writerReader writeCap
                -- Distribute the shares.
                liftIO $ placeShares (SDMF.Keys.unStorageIndex . SDMF.verifierStorageIndex . SDMF.readerVerifier . SDMF.writerReader $ writeCap) (Binary.encode <$> shares) perServerShareCount servers 0

                let serverMap = Map.fromList $ zip (Set.toList serverIDs') servers
                    lookupServer = someServers serverMap
                    serverAnnouncements = Map.fromSet makeAnn serverIDs'

                -- Recover the plaintext from the servers.
                result <- liftIO $ download serverAnnouncements readCap lookupServer
                diff (Right plaintext) (==) result
        , testCase "immutable upload/download to using Great Black Swamp" $ do
            pure ()
            -- Consider moving these tests to another module, they're pretty
            -- different and there's quite a handful of them.
            --
            -- ERROR CASES
            -- Server presents incorrect TLS certificate
            --   * See https://whetstone.private.storage/privatestorage/tahoe-great-black-swamp/-/issues/27
            -- Server returns error response to our request
            --   * https://whetstone.private.storage/privatestorage/gbs-downloader/-/issues/4
            -- Server returns tampered share data
            --   * https://whetstone.private.storage/privatestorage/gbs-downloader/-/issues/5
        ]
  where
    -- A server lookup function that always fails.
    noServers = pure . Left . URIParseError

    -- A server lookup function that finds servers already present in a Map.
    someServers :: Applicative m => Map.Map StorageServerID StorageServer -> LookupServer m
    someServers servers ann =
        pure $ case result of
            Nothing -> Left AnnouncementStructureUnmatched
            Just ss -> Right ss
      where
        result = do
            furl <- storageServerAnnouncementFURL ann
            let serverId = parseURL furl
            Map.lookup serverId servers

        -- Exactly match the nonsense makeAnn spits out
        parseURL = T.take 2 . T.drop 5

    --- PHILOSOFY
    -- We wish that share numbers were an opaque type instead of a
    -- numeric/integral type.  This is not the place to argue the point
    -- though.
    placeShares :: B.ByteString -> [BL.ByteString] -> [Int] -> [StorageServer] -> Int -> IO ()
    -- Out of shares, done.
    placeShares _ [] _ _ _ = pure ()
    -- Out of placement info but not out of shares is a programming error.
    placeShares _ _ [] _ _ = throwIO RanOutOfPlacementInfo
    -- Out of servers but not out of shares is a programming error.
    placeShares _ _ _ [] _ = throwIO RanOutOfServers
    -- Having some of all three means we can make progress.
    placeShares si shares (n : ns) (s : ss) sharesSoFar = do
        -- write the right number of shares to this server
        zipWithM_
            (\shnum share -> storageServerWrite s si shnum 0 share)
            [fromIntegral sharesSoFar ..]
            (BL.toStrict <$> take n shares)
        -- recurse to write the rest
        placeShares si (drop n shares) ns ss (sharesSoFar + n)

    -- Make up a distinct (but nonsense) announcement for a given storage
    -- server identifier.
    makeAnn :: StorageServerID -> StorageServerAnnouncement
    makeAnn sid =
        def
            { storageServerAnnouncementFURL = Just $ "pb://" <> sid <> "/" <> sid
            , storageServerAnnouncementNick = Just . encodeBase32Unpadded . encodeUtf8 $ sid
            }

    -- Generate lists of ints that sum to a given total.
    genListWithSum :: MonadGen m => Int -> Int -> m [Int]
    -- We hit the target.
    genListWithSum _ 0 = pure []
    -- We only have room for one more element.
    genListWithSum 1 t = pure [t]
    -- Use up some of what's left on one element and recurse.
    genListWithSum maxLength t = do
        v <- Gen.int (Range.linear 0 t)
        (v :) <$> genListWithSum (maxLength - 1) (t - v)

trivialCap :: Word16 -> Word16 -> Reader
trivialCap required total = Reader{..}
  where
    Just readKey = buildKey $ B.replicate 32 0x00
    storageIndex = B.replicate 32 0x00
    fingerprint = B.replicate 32 0x00
    size = 1234
    verifier = Verifier{..}

main :: IO ()
main = do
    -- Hedgehog writes some non-ASCII and the whole test process will die if
    -- it can't be encoded.  Increase the chances that all of the output can
    -- be encoded by forcing the use of UTF-8 (overriding the LANG-based
    -- choice normally made).
    hSetEncoding stdout utf8
    hSetEncoding stderr utf8

    defaultMain tests