From faec32dc5eae3901fcd56f1e0ff671cd0c9affdc Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com>
Date: Wed, 1 Mar 2023 13:50:08 +0000
Subject: [PATCH] xftp: constrain supported file size to 1gb

---
 src/Simplex/FileTransfer/Client/Main.hs | 11 ++++++++++-
 tests/XFTPCLI.hs                        |  9 +++++++--
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/Simplex/FileTransfer/Client/Main.hs b/src/Simplex/FileTransfer/Client/Main.hs
index e1f07ceb5..d0c68c3be 100644
--- a/src/Simplex/FileTransfer/Client/Main.hs
+++ b/src/Simplex/FileTransfer/Client/Main.hs
@@ -70,6 +70,12 @@ chunkSize2 = mb 1
 chunkSize3 :: Word32
 chunkSize3 = mb 4
 
+maxFileSize :: Int64
+maxFileSize = gb 1
+
+maxFileSizeStr :: String
+maxFileSizeStr = B.unpack . strEncode $ FileSize maxFileSize
+
 fileSizeLen :: Int64
 fileSizeLen = 8
 
@@ -267,6 +273,7 @@ instance Encoding FileHeader where
 cliSendFile :: SendOptions -> ExceptT CLIError IO ()
 cliSendFile SendOptions {filePath, outputDir, numRecipients, xftpServers, retryCount, tempPath, verbose} = do
   let (_, fileName) = splitFileName filePath
+  liftIO $ printNoNewLine "Encrypting file..."
   (encPath, fdRcv, fdSnd, chunkSpecs, encSize) <- encryptFile fileName
   liftIO $ printNoNewLine "Uploading file..."
   uploadedChunks <- newTVarIO []
@@ -284,10 +291,11 @@ cliSendFile SendOptions {filePath, outputDir, numRecipients, xftpServers, retryC
   where
     encryptFile :: String -> ExceptT CLIError IO (FilePath, FileDescription 'FPRecipient, FileDescription 'FPSender, [XFTPChunkSpec], Int64)
     encryptFile fileName = do
+      fileSize <- fromInteger <$> getFileSize filePath
+      when (fileSize > maxFileSize) $ throwError $ CLIError $ "Files bigger than " <> maxFileSizeStr <> " are not supported"
       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 + authTagSize
@@ -432,6 +440,7 @@ cliReceiveFile ReceiveOptions {fileDescription, filePath, retryCount, tempPath,
       when (encDigest /= unFileDigest digest) $ throwError $ CLIError "File digest mismatch"
       encSize <- liftIO $ foldM (\s path -> (s +) . fromIntegral <$> getFileSize path) 0 chunkPaths
       when (FileSize encSize /= size) $ throwError $ CLIError "File size mismatch"
+      liftIO $ printNoNewLine "Decrypting file..."
       path <- decryptFile encSize chunkPaths key nonce
       forM_ chunks $ acknowledgeFileChunk a
       whenM (doesPathExist encPath) $ removeDirectoryRecursive encPath
diff --git a/tests/XFTPCLI.hs b/tests/XFTPCLI.hs
index 2789bbcc5..6cbbb9090 100644
--- a/tests/XFTPCLI.hs
+++ b/tests/XFTPCLI.hs
@@ -158,8 +158,13 @@ testPrepareChunkSizes = do
     r3 = replicate 3
 
 uploadProgress :: String -> Bool
-uploadProgress s = "Uploading file..." `isPrefixOf` s && "File uploaded!" `isInfixOf` s
+uploadProgress s =
+  "Encrypting file..." `isPrefixOf` s
+    && "Uploading file..." `isInfixOf` s
+    && "File uploaded!" `isInfixOf` s
 
 downloadProgress :: FilePath -> String -> Bool
 downloadProgress fileName s =
-  "Downloading file..." `isPrefixOf` s && ("File downloaded: " <> recipientFiles </> fileName <> "\r") `isSuffixOf` s
+  "Downloading file..." `isPrefixOf` s
+    && "Decrypting file..." `isInfixOf` s
+    && ("File downloaded: " <> recipientFiles </> fileName <> "\r") `isSuffixOf` s