Files
simplexmq/src/Simplex/FileTransfer/Client/Main.hs
Evgeny Poberezkin acdf4f41c5 XFTP: CLI tests, output command results, info command (#642)
* XFTP: CLI tests, output command results

* info command, test

* fix test
2023-02-17 11:38:43 +00:00

449 lines
20 KiB
Haskell

{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Simplex.FileTransfer.Client.Main (xftpClientCLI) where
import Control.Concurrent.STM (stateTVar)
import Control.Monad
import Control.Monad.Except
import Crypto.Random (getRandomBytes)
import qualified Data.Attoparsec.ByteString.Char8 as A
import Data.Bifunctor (first)
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString.Lazy.Char8 as LB
import Data.Int (Int64)
import Data.List (foldl', sortOn)
import Data.List.NonEmpty (NonEmpty (..), nonEmpty)
import qualified Data.List.NonEmpty as L
import Data.Map (Map)
import qualified Data.Map as M
import Data.Maybe (fromMaybe)
import Data.Word (Word32)
import Options.Applicative
import Simplex.FileTransfer.Client
import Simplex.FileTransfer.Client.Agent
import Simplex.FileTransfer.Description
import Simplex.FileTransfer.Description (FileSize (unFileSize))
import Simplex.FileTransfer.Protocol (FileInfo (..))
import Simplex.Messaging.Agent.Lock
import qualified Simplex.Messaging.Crypto as C
import qualified Simplex.Messaging.Crypto.Lazy as LC
import Simplex.Messaging.Encoding
import Simplex.Messaging.Encoding.String (StrEncoding (..))
import Simplex.Messaging.Parsers (parseAll)
import Simplex.Messaging.Protocol (SenderId, SndPrivateSignKey, SndPublicVerifyKey, XFTPServer)
import Simplex.Messaging.Server.CLI (getCliCommand')
import Simplex.Messaging.Util (ifM, whenM)
import System.Exit (exitFailure)
import System.FilePath (splitExtensions, splitFileName, (</>))
import System.IO.Temp (getCanonicalTemporaryDirectory)
import System.Random (StdGen, newStdGen, randomR)
import UnliftIO
import UnliftIO.Directory
xftpClientVersion :: String
xftpClientVersion = "0.1.0"
defaultChunkSize :: Word32
defaultChunkSize = 8 * mb
smallChunkSize :: Word32
smallChunkSize = 1 * mb
fileSizeLen :: Int64
fileSizeLen = 8
cbAuthTagLen :: Int64
cbAuthTagLen = fromIntegral C.cbAuthTagSize
mb :: Num a => a
mb = 1024 * 1024
newtype CLIError = CLIError String
deriving (Eq, Show, Exception)
data CliCommand
= SendFile SendOptions
| ReceiveFile ReceiveOptions
| RandomFile RandomFileOptions
| FileDescrInfo InfoOptions
data SendOptions = SendOptions
{ filePath :: FilePath,
outputDir :: Maybe FilePath,
numRecipients :: Int,
xftpServers :: [XFTPServer],
retryCount :: Int,
tempPath :: Maybe FilePath
}
deriving (Show)
data ReceiveOptions = ReceiveOptions
{ fileDescription :: FilePath,
filePath :: Maybe FilePath,
retryCount :: Int,
tempPath :: Maybe FilePath
}
deriving (Show)
newtype InfoOptions = InfoOptions
{ fileDescription :: FilePath
}
deriving (Show)
data RandomFileOptions = RandomFileOptions
{ filePath :: FilePath,
fileSize :: FileSize Int
}
deriving (Show)
defaultRetryCount :: Int
defaultRetryCount = 3
defaultXFTPServers :: NonEmpty XFTPServer
defaultXFTPServers = L.fromList ["xftp://vr0bXzm4iKkLvleRMxLznTS-lHjXEyXunxn_7VJckk4=@localhost:443"]
cliCommandP :: Parser CliCommand
cliCommandP =
hsubparser
( command "send" (info (SendFile <$> sendP) (progDesc "Send file"))
<> command "recv" (info (ReceiveFile <$> receiveP) (progDesc "Receive file"))
<> command "info" (info (FileDescrInfo <$> infoP) (progDesc "Show file description"))
<> command "rand" (info (RandomFile <$> randomP) (progDesc "Generate a random file of a given size"))
)
where
sendP :: Parser SendOptions
sendP =
SendOptions
<$> argument str (metavar "FILE" <> help "File to send")
<*> optional (argument str $ metavar "DIR" <> help "Directory to save file descriptions (default: current directory)")
<*> option auto (short 'n' <> metavar "COUNT" <> help "Number of recipients" <> value 1 <> showDefault)
<*> xftpServers
<*> retries
<*> temp
receiveP :: Parser ReceiveOptions
receiveP =
ReceiveOptions
<$> fileDescrArg
<*> optional (argument str $ metavar "DIR" <> help "Directory to save file (default: system Downloads directory)")
<*> retries
<*> temp
infoP :: Parser InfoOptions
infoP = InfoOptions <$> fileDescrArg
randomP :: Parser RandomFileOptions
randomP =
RandomFileOptions
<$> argument str (metavar "FILE" <> help "Path to save file")
<*> argument strDec (metavar "SIZE" <> help "File size (bytes/kb/mb)")
strDec = eitherReader $ strDecode . B.pack
fileDescrArg = argument str (metavar "FILE" <> help "File description file")
retries = option auto (long "retry" <> short 'r' <> metavar "RETRY" <> help "Number of network retries" <> value defaultRetryCount <> showDefault)
temp = optional (strOption $ long "tmp" <> metavar "TMP" <> help "Directory for temporary encrypted file (default: system temp directory)")
xftpServers =
option
parseXFTPServers
( long "servers"
<> short 's'
<> metavar "SERVER"
<> help "Semicolon-separated list of XFTP server(s) to use (each server can have more than one hostname)"
<> value []
)
parseXFTPServers = eitherReader $ parseAll xftpServersP . B.pack
xftpServersP = strP `A.sepBy1` A.char ';'
data SentFileChunk = SentFileChunk
{ chunkNo :: Int,
sndId :: SenderId,
sndPrivateKey :: SndPrivateSignKey,
chunkSize :: FileSize Word32,
digest :: FileDigest,
replicas :: [SentFileChunkReplica]
}
deriving (Eq, Show)
data SentFileChunkReplica = SentFileChunkReplica
{ server :: XFTPServer,
recipients :: [(ChunkReplicaId, C.APrivateSignKey)]
}
deriving (Eq, Show)
data SentRecipientReplica = SentRecipientReplica
{ chunkNo :: Int,
server :: XFTPServer,
rcvNo :: Int,
rcvId :: ChunkReplicaId,
rcvKey :: C.APrivateSignKey,
digest :: FileDigest,
chunkSize :: FileSize Word32
}
xftpClientCLI :: IO ()
xftpClientCLI =
getCliCommand' cliCommandP clientVersion >>= \case
SendFile opts -> runE $ cliSendFile opts
ReceiveFile opts -> runE $ cliReceiveFile opts
FileDescrInfo opts -> runE $ cliFileDescrInfo opts
RandomFile opts -> cliRandomFile opts
where
clientVersion = "SimpleX XFTP client v" <> xftpClientVersion
runE :: ExceptT CLIError IO () -> IO ()
runE a =
runExceptT a >>= \case
Left (CLIError e) -> putStrLn e >> exitFailure
_ -> pure ()
-- fileExtra is added to allow header extension in future versions
data FileHeader = FileHeader
{ fileName :: String,
fileExtra :: Maybe String
}
deriving (Eq, Show)
instance Encoding FileHeader where
smpEncode FileHeader {fileName, fileExtra} = smpEncode (fileName, fileExtra)
smpP = do
(fileName, fileExtra) <- smpP
pure FileHeader {fileName, fileExtra}
cliSendFile :: SendOptions -> ExceptT CLIError IO ()
cliSendFile SendOptions {filePath, outputDir, numRecipients, xftpServers, retryCount, tempPath} = do
let (_, fileName) = splitFileName filePath
(encPath, fd, chunkSpecs) <- encryptFile fileName
sentChunks <- uploadFile chunkSpecs
whenM (doesFileExist encPath) $ removeFile encPath
-- TODO if only small chunks, use different default size
liftIO $ do
fds <- writeFileDescriptions fileName $ createFileDescriptions fd sentChunks
putStrLn "File uploaded!\nPass file descriptions to the recipient(s):"
forM_ fds putStrLn
where
encryptFile :: String -> ExceptT CLIError IO (FilePath, FileDescription, [XFTPChunkSpec])
encryptFile fileName = do
encPath <- getEncPath tempPath "xftp"
key <- liftIO C.randomSbKey
nonce <- liftIO C.randomCbNonce
fileSize <- fromInteger <$> getFileSize filePath
let fileHdr = smpEncode FileHeader {fileName, fileExtra = Nothing}
fileSize' = fromIntegral (B.length fileHdr) + fileSize
chunkSizes = prepareChunkSizes $ fileSize' + fileSizeLen + cbAuthTagLen
paddedSize = fromIntegral $ sum chunkSizes
encrypt fileHdr key nonce fileSize' paddedSize encPath
digest <- liftIO $ LC.sha512Hash <$> LB.readFile encPath
let chunkSpecs = prepareChunkSpecs encPath chunkSizes
fd = FileDescription {size = FileSize paddedSize, digest = FileDigest digest, key, nonce, chunkSize = FileSize defaultChunkSize, chunks = []}
pure (encPath, fd, chunkSpecs)
where
encrypt :: ByteString -> C.SbKey -> C.CbNonce -> Int64 -> Int64 -> FilePath -> ExceptT CLIError IO ()
encrypt fileHdr key nonce fileSize' paddedSize encFile = do
f <- liftIO $ LB.readFile filePath
let f' = LB.fromStrict fileHdr <> f
c <- liftEither $ first (CLIError . show) $ LC.sbEncrypt key nonce f' fileSize' $ paddedSize - cbAuthTagLen
liftIO $ LB.writeFile encFile c
-- let padSize = paddedSize - fileSize - fromIntegral (B.length fileHdr)
-- when (padSize > 0) . LB.hPut h $ LB.replicate padSize '#'
uploadFile :: [XFTPChunkSpec] -> ExceptT CLIError IO [SentFileChunk]
uploadFile chunks = do
a <- atomically $ newXFTPAgent defaultXFTPClientAgentConfig
gen <- newTVarIO =<< liftIO newStdGen
let xftpSrvs = fromMaybe defaultXFTPServers (nonEmpty xftpServers)
-- TODO shuffle chunks
sentChunks <- pooledForConcurrentlyN 32 (zip [1 ..] chunks) $ uploadFileChunk a gen xftpSrvs
-- TODO unshuffle chunks
pure $ map snd sentChunks
where
retries :: Show e => ExceptT e IO a -> ExceptT CLIError IO a
retries = withRetry retryCount . withExceptT (CLIError . show)
uploadFileChunk :: XFTPClientAgent -> TVar StdGen -> NonEmpty XFTPServer -> (Int, XFTPChunkSpec) -> ExceptT CLIError IO (Int, SentFileChunk)
uploadFileChunk a gen srvs (chunkNo, chunkSpec@XFTPChunkSpec {chunkSize}) = do
(sndKey, spKey) <- liftIO $ C.generateSignatureKeyPair C.SEd25519
rKeys <- liftIO $ L.fromList <$> replicateM numRecipients (C.generateSignatureKeyPair C.SEd25519)
chInfo@FileInfo {digest} <- liftIO $ getChunkInfo sndKey chunkSpec
xftpServer <- liftIO $ getXFTPServer gen srvs
c <- retries $ getXFTPServerClient a xftpServer
(sndId, rIds) <- retries $ createXFTPChunk c spKey chInfo $ L.map fst rKeys
retries $ uploadXFTPChunk c spKey sndId chunkSpec
let recipients = L.toList $ L.map ChunkReplicaId rIds `L.zip` L.map snd rKeys
replicas = [SentFileChunkReplica {server = xftpServer, recipients}]
pure (chunkNo, SentFileChunk {chunkNo, sndId, sndPrivateKey = spKey, chunkSize = FileSize $ fromIntegral chunkSize, digest = FileDigest digest, replicas})
getChunkInfo :: SndPublicVerifyKey -> XFTPChunkSpec -> IO FileInfo
getChunkInfo sndKey XFTPChunkSpec {filePath = chunkPath, chunkOffset, chunkSize} =
withFile chunkPath ReadMode $ \h -> do
hSeek h AbsoluteSeek $ fromIntegral chunkOffset
digest <- LC.sha512Hash <$> LB.hGet h (fromIntegral chunkSize)
pure FileInfo {sndKey, size = fromIntegral chunkSize, digest}
getXFTPServer :: TVar StdGen -> NonEmpty XFTPServer -> IO XFTPServer
getXFTPServer gen = \case
srv :| [] -> pure srv
servers -> do
atomically $ (servers L.!!) <$> stateTVar gen (randomR (0, L.length servers - 1))
-- M chunks, R replicas, N recipients
-- rcvReplicas: M[SentFileChunk] -> M * R * N [SentRecipientReplica]
-- rcvChunks: M * R * N [SentRecipientReplica] -> N[ M[FileChunk] ]
createFileDescriptions :: FileDescription -> [SentFileChunk] -> [FileDescription]
createFileDescriptions fd sentChunks = map (\chunks -> (fd :: FileDescription) {chunks}) rcvChunks
where
rcvReplicas :: [SentRecipientReplica]
rcvReplicas =
concatMap
( \SentFileChunk {chunkNo, digest, chunkSize, replicas} ->
concatMap
( \SentFileChunkReplica {server, recipients} ->
zipWith (\rcvNo (rcvId, rcvKey) -> SentRecipientReplica {chunkNo, server, rcvNo, rcvId, rcvKey, digest, chunkSize}) [1 ..] recipients
)
replicas
)
sentChunks
rcvChunks :: [[FileChunk]]
rcvChunks = map (sortChunks . M.elems) $ M.elems $ foldl' addRcvChunk M.empty rcvReplicas
sortChunks :: [FileChunk] -> [FileChunk]
sortChunks = map reverseReplicas . sortOn (chunkNo :: FileChunk -> Int)
reverseReplicas ch@FileChunk {replicas} = (ch :: FileChunk) {replicas = reverse replicas}
addRcvChunk :: Map Int (Map Int FileChunk) -> SentRecipientReplica -> Map Int (Map Int FileChunk)
addRcvChunk m SentRecipientReplica {chunkNo, server, rcvNo, rcvId, rcvKey, digest, chunkSize} =
M.alter (Just . addOrChangeRecipient) rcvNo m
where
addOrChangeRecipient :: Maybe (Map Int FileChunk) -> Map Int FileChunk
addOrChangeRecipient = \case
Just m' -> M.alter (Just . addOrChangeChunk) chunkNo m'
_ -> M.singleton chunkNo $ FileChunk {chunkNo, digest, chunkSize, replicas = [replica]}
addOrChangeChunk :: Maybe FileChunk -> FileChunk
addOrChangeChunk = \case
Just ch@FileChunk {replicas} -> ch {replicas = replica : replicas}
_ -> FileChunk {chunkNo, digest, chunkSize, replicas = [replica]}
replica = FileChunkReplica {server, rcvId, rcvKey}
writeFileDescriptions :: String -> [FileDescription] -> IO [FilePath]
writeFileDescriptions fileName fds = do
outDir <- uniqueCombine (fromMaybe "." outputDir) (fileName <> ".xftp")
createDirectoryIfMissing True outDir
forM (zip [1 ..] fds) $ \(i :: Int, fd) -> do
let fdPath = outDir </> ("rcv" <> show i <> ".xftp")
B.writeFile fdPath $ strEncode fd
pure fdPath
cliReceiveFile :: ReceiveOptions -> ExceptT CLIError IO ()
cliReceiveFile ReceiveOptions {fileDescription, filePath, retryCount, tempPath} = do
ValidFileDescription FileDescription {size, digest, key, nonce, chunks} <- getFileDescription fileDescription
encPath <- getEncPath tempPath "xftp"
-- withFile encPath WriteMode $ \h -> do
-- liftIO $ LB.hPut h $ LB.replicate (unFileSize size) '#'
a <- atomically $ newXFTPAgent defaultXFTPClientAgentConfig
writeLock <- atomically createLock
let chunkSizes = prepareChunkSizes $ unFileSize size
chunkSpecs = prepareChunkSpecs encPath chunkSizes
-- chunks have to be ordered because of AppendMode
forM_ (zip chunkSpecs chunks) $ \(chunkSpec, chunk) -> do
downloadFileChunk a writeLock chunk chunkSpec
encDigest <- liftIO $ LC.sha512Hash <$> LB.readFile encPath
when (encDigest /= unFileDigest digest) $ throwError $ CLIError "File digest mismatch"
path <- decryptFile encPath key nonce
whenM (doesFileExist encPath) $ removeFile encPath
liftIO $ putStrLn $ "File received: " <> path
where
retries :: Show e => ExceptT e IO a -> ExceptT CLIError IO a
retries = withRetry retryCount . withExceptT (CLIError . show)
downloadFileChunk :: XFTPClientAgent -> Lock -> FileChunk -> XFTPChunkSpec -> ExceptT CLIError IO ()
downloadFileChunk a writeLock FileChunk {replicas = replica : _} chunkSpec = do
let FileChunkReplica {server, rcvId, rcvKey} = replica
c <- retries $ getXFTPServerClient a server
(rKey, rpKey) <- liftIO C.generateKeyPair'
(sKey, body) <- retries $ downloadXFTPChunk c rcvKey (unChunkReplicaId rcvId) rKey
-- download and decrypt (DH) chunk from server using XFTPClient
-- verify chunk digest - in the client
-- save to correct location in file - also in the client
retries $ withLock writeLock "save" $ receiveXFTPChunk body chunkSpec
downloadFileChunk _ _ _ _ = pure ()
decryptFile :: FilePath -> C.SbKey -> C.CbNonce -> ExceptT CLIError IO FilePath
decryptFile encPath key nonce = do
f <- liftIO $ LB.readFile encPath
f' <- liftEither $ first (CLIError . show) $ LC.sbDecrypt key nonce f
let (fileHdr, f'') = LB.splitAt 1024 f'
-- withFile encPath ReadMode $ \r -> do
-- fileHdr <- liftIO $ B.hGet r 1024
case A.parse smpP $ LB.toStrict fileHdr of
A.Fail _ _ e -> throwError $ CLIError $ "Invalid file header: " <> e
A.Partial _ -> throwError $ CLIError "Invalid file header"
A.Done rest FileHeader {fileName} -> do
path <- getFilePath fileName
liftIO $ LB.writeFile path $ LB.fromStrict rest <> f''
pure path
getFilePath :: String -> ExceptT CLIError IO FilePath
getFilePath name =
case filePath of
Just path ->
ifM (doesDirectoryExist path) (uniqueCombine path name) $
ifM (doesFileExist path) (throwError $ CLIError "File already exists") (pure path)
_ -> (`uniqueCombine` name) . (</> "Downloads") =<< getHomeDirectory
cliFileDescrInfo :: InfoOptions -> ExceptT CLIError IO ()
cliFileDescrInfo InfoOptions {fileDescription} = do
ValidFileDescription FileDescription {size, chunkSize, chunks} <- getFileDescription fileDescription
let replicas = groupReplicasByServer chunkSize chunks
liftIO $ do
putStrLn $ "File download size: " <> strEnc size
putStrLn "File server(s):"
forM_ replicas $ \srvReplicas -> do
let srv = replicaServer $ head srvReplicas
chSizes = map (\FileServerReplica {chunkSize = chSize_} -> unFileSize $ fromMaybe chunkSize chSize_) srvReplicas
putStrLn $ strEnc srv <> ": " <> strEnc (FileSize $ sum chSizes)
strEnc :: StrEncoding a => a -> String
strEnc = B.unpack . strEncode
getFileDescription :: FilePath -> ExceptT CLIError IO ValidFileDescription
getFileDescription path = do
fd <- ExceptT $ first (CLIError . ("Failed to parse file description: " <>)) . strDecode <$> B.readFile path
liftEither . first CLIError $ validateFileDescription fd
prepareChunkSizes :: Int64 -> [Word32]
prepareChunkSizes 0 = []
prepareChunkSizes size
| size >= defSz = replicate (fromIntegral n1) defaultChunkSize <> prepareChunkSizes remSz
| size > defSz `div` 2 = [defaultChunkSize]
| otherwise = replicate (fromIntegral n2') smallChunkSize
where
(n1, remSz) = size `divMod` defSz
n2' = let (n2, remSz2) = (size `divMod` fromIntegral smallChunkSize) in if remSz2 == 0 then n2 else n2 + 1
defSz = fromIntegral defaultChunkSize :: Int64
prepareChunkSpecs :: FilePath -> [Word32] -> [XFTPChunkSpec]
prepareChunkSpecs filePath chunkSizes = reverse . snd $ foldl' addSpec (0, []) chunkSizes
where
addSpec :: (Int64, [XFTPChunkSpec]) -> Word32 -> (Int64, [XFTPChunkSpec])
addSpec (chunkOffset, specs) sz =
let spec = XFTPChunkSpec {filePath, chunkOffset, chunkSize = fromIntegral sz}
in (chunkOffset + fromIntegral sz, spec : specs)
getEncPath :: MonadIO m => Maybe FilePath -> String -> m FilePath
getEncPath path name = (`uniqueCombine` (name <> ".encrypted")) =<< maybe (liftIO getCanonicalTemporaryDirectory) pure path
uniqueCombine :: MonadIO m => FilePath -> String -> m FilePath
uniqueCombine filePath fileName = tryCombine (0 :: Int)
where
tryCombine n =
let (name, ext) = splitExtensions fileName
suffix = if n == 0 then "" else "_" <> show n
f = filePath </> (name <> suffix <> ext)
in ifM (doesPathExist f) (tryCombine $ n + 1) (pure f)
withRetry :: Int -> ExceptT CLIError IO a -> ExceptT CLIError IO a
withRetry 0 _ = throwError $ CLIError "internal: no retry attempts"
withRetry 1 a = a
withRetry n a = a `catchError` \_ -> withRetry (n - 1) a
cliRandomFile :: RandomFileOptions -> IO ()
cliRandomFile RandomFileOptions {filePath, fileSize = FileSize size} = do
withFile filePath WriteMode (`saveRandomFile` size)
putStrLn $ "File created: " <> filePath
where
saveRandomFile h sz = do
bytes <- getRandomBytes $ min mb sz
B.hPut h bytes
when (sz > mb) $ saveRandomFile h (sz - mb)