xftp: validate digests of file chunks in client and server

src/Simplex/FileTransfer/Client.hs

@@ -14,7 +14,7 @@ import Data.Bifunctor (first)
 import Data.ByteString.Builder (Builder, byteString)
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
-import Data.Functor (($>))
+import qualified Data.ByteString.Lazy.Char8 as LB
 import Data.Int (Int64)
 import Data.List.NonEmpty (NonEmpty (..))
 import Data.Word (Word32)
@@ -32,8 +32,10 @@ import Simplex.Messaging.Client
     transportClientConfig,
   )
 import qualified Simplex.Messaging.Crypto as C
+import qualified Simplex.Messaging.Crypto.Lazy as LC
 import Simplex.Messaging.Protocol
-  ( ProtocolServer (ProtocolServer),
+  ( Protocol (..),
+    ProtocolServer (..),
     RcvPublicDhKey,
     RecipientId,
     SenderId,
@@ -43,9 +45,8 @@ import Simplex.Messaging.Transport.Client (TransportClientConfig)
 import Simplex.Messaging.Transport.HTTP2
 import Simplex.Messaging.Transport.HTTP2.Client
 import Simplex.Messaging.Util (bshow, liftEitherError, whenM)
-import System.Directory (doesFileExist, removeFile)
-import System.IO (IOMode (..), SeekMode (..))
-import UnliftIO.IO (hSeek, withFile)
+import UnliftIO.Directory
+import UnliftIO.IO
 
 data XFTPClient = XFTPClient
   { http2Client :: HTTP2Client,
@@ -110,7 +111,9 @@ sendXFTPCommand XFTPClient {http2Client = http2@HTTP2Client {sessionId}} pKey fI
   -- TODO validate that the file ID is the same as in the request?
   (_, _, (_, _fId, respOrErr)) <- liftEither . first PCEResponseError $ xftpDecodeTransmission sessionId bodyHead
   case respOrErr of
-    Right r -> pure (r, body)
+    Right r -> case protocolError r of
+      Just e -> throwError $ PCEProtocolError e
+      _ -> pure (r, body)
     Left e -> throwError $ PCEResponseError e
   where
     streamBody :: ByteString -> (Builder -> IO ()) -> IO () -> IO ()
@@ -152,13 +155,17 @@ downloadXFTPChunk c rpKey fId rKey =
       _ -> throwError $ PCEResponseError NO_FILE
     (r, _) -> throwError . PCEUnexpectedResponse $ bshow r
 
-receiveXFTPChunk :: XFTPChunkBody -> FilePath -> Word32 -> ExceptT XFTPClientError IO ()
-receiveXFTPChunk XFTPChunkBody {chunkPart} filePath chunkSize = do
-  withExceptT PCEResponseError . ExceptT $ do
+receiveXFTPChunk :: XFTPChunkBody -> FilePath -> Word32 -> ByteString -> ExceptT XFTPClientError IO ()
+receiveXFTPChunk XFTPChunkBody {chunkPart} filePath chunkSize chunkDigest = do
+  withExceptT PCEResponseError $ do
     -- TODO chunk decryption
-    withFile filePath WriteMode (\h -> receiveFile h chunkPart chunkSize) >>= \case
-      Right () -> pure $ Right ()
-      Left e -> whenM (doesFileExist filePath) (removeFile filePath) $> Left e
+    receiveChunk `catchError` \e ->
+      whenM (doesFileExist filePath) (removeFile filePath) >> throwError e
+  where
+    receiveChunk = do
+      ExceptT . withFile filePath WriteMode $ \h -> receiveFile h chunkPart chunkSize
+      digest' <- liftIO $ LC.sha512Hash <$> LB.readFile filePath
+      when (digest' /= chunkDigest) $ throwError DIGEST
 
 -- FADD :: NonEmpty RcvPublicVerifyKey -> FileCommand Sender
 -- FDEL :: FileCommand Sender

src/Simplex/FileTransfer/Client/Main.hs

@@ -342,7 +342,7 @@ cliReceiveFile ReceiveOptions {fileDescription, filePath, retryCount, tempPath}
     retries :: Show e => ExceptT e IO a -> ExceptT CLIError IO a
     retries = withRetry retryCount . withExceptT (CLIError . show)
     downloadFileChunk :: XFTPClientAgent -> FilePath -> FileChunk -> ExceptT CLIError IO FilePath
-    downloadFileChunk a encPath FileChunk {chunkNo, chunkSize, replicas = replica : _} = do
+    downloadFileChunk a encPath FileChunk {chunkNo, chunkSize, digest, replicas = replica : _} = do
       let FileChunkReplica {server, rcvId, rcvKey} = replica
       chunkPath <- uniqueCombine encPath $ show chunkNo
       c <- retries $ getXFTPServerClient a server
@@ -350,7 +350,7 @@ cliReceiveFile ReceiveOptions {fileDescription, filePath, retryCount, tempPath}
       (sKey, body) <- retries $ downloadXFTPChunk c rcvKey (unChunkReplicaId rcvId) rKey
       -- download and decrypt (DH) chunk from server using XFTPClient
       -- verify chunk digest - in the client
-      retries $ receiveXFTPChunk body chunkPath $ unFileSize chunkSize
+      retries $ receiveXFTPChunk body chunkPath (unFileSize chunkSize) (unFileDigest digest)
       pure chunkPath
     downloadFileChunk _ _ _ = throwError $ CLIError "chunk has no replicas"
     decryptFile :: [FilePath] -> C.SbKey -> C.CbNonce -> ExceptT CLIError IO FilePath
@@ -368,6 +368,7 @@ cliReceiveFile ReceiveOptions {fileDescription, filePath, retryCount, tempPath}
           pure path
     readChunks :: [FilePath] -> IO LB.ByteString
     readChunks = foldM (\s path -> (s <>) <$> LB.readFile path) LB.empty
+    {-# NOINLINE readChunks #-}
     getFilePath :: String -> ExceptT CLIError IO FilePath
     getFilePath name =
       case filePath of

src/Simplex/FileTransfer/Server.hs

@@ -21,6 +21,7 @@ import qualified Data.ByteString.Base64.URL as B64
 import Data.ByteString.Builder (byteString)
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
+import qualified Data.ByteString.Lazy.Char8 as LB
 import Data.Functor (($>))
 import Data.List (intercalate)
 import qualified Data.List.NonEmpty as L
@@ -36,6 +37,7 @@ import Simplex.FileTransfer.Server.Stats
 import Simplex.FileTransfer.Server.Store
 import Simplex.FileTransfer.Transport (receiveFile, sendFile)
 import qualified Simplex.Messaging.Crypto as C
+import qualified Simplex.Messaging.Crypto.Lazy as LC
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Protocol (CorrId, RcvPublicDhKey)
 import Simplex.Messaging.Server (dummyVerifyCmd, verifyCmdSignature)
@@ -147,7 +149,6 @@ processRequest HTTP2Request {sessionId, reqBody = body@HTTP2Body {bodyHead}, sen
   where
     sendXFTPResponse :: (CorrId, XFTPFileId, FileResponse) -> Maybe ServerFile -> M ()
     sendXFTPResponse (corrId, fId, resp) serverFile_ = do
-      -- liftIO . sendResponse . H.responseBuilder N.ok200 [] . byteString $
       let t_ = xftpEncodeTransmission sessionId Nothing (corrId, fId, resp)
       liftIO $ sendResponse $ H.responseStreaming N.ok200 [] $ streamBody t_
       where
@@ -157,7 +158,7 @@ processRequest HTTP2Request {sessionId, reqBody = body@HTTP2Body {bodyHead}, sen
             Right t -> do
               send $ byteString t
               -- TODO chunk encryption
-              forM_ serverFile_ $ \ServerFile {filePath, fileSize, fileDhSecret} -> do
+              forM_ serverFile_ $ \ServerFile {filePath, fileSize, fileDhSecret} ->
                 withFile filePath ReadMode $ \h -> sendFile h send $ fromIntegral fileSize
           done
 
@@ -207,17 +208,22 @@ processXFTPRequest HTTP2Body {bodyPart} = \case
     noFile resp = pure (resp, Nothing)
     receiveServerFile :: FileRec -> M FileResponse
     receiveServerFile FileRec {senderId, fileInfo, filePath} = case bodyPart of
+      -- TODO do not allow repeated file upload
       Nothing -> pure $ FRErr SIZE
       Just getBody -> do
         -- TODO validate body size before downloading, once it's populated
         path <- asks $ filesPath . config
         let fPath = path </> B.unpack (B64.encode senderId)
-            FileInfo {size, digest} = fileInfo
-        -- TODO check digest
         liftIO $
-          withFile fPath WriteMode (\h -> receiveFile h getBody size) >>= \case
+          runExceptT (receiveChunk fPath fileInfo) >>= \case
             Right () -> atomically $ writeTVar filePath (Just fPath) $> FROk
             Left e -> whenM (doesFileExist fPath) (removeFile fPath) $> FRErr e
+        where
+          receiveChunk fPath FileInfo {size, digest} = do
+            ExceptT . withFile fPath WriteMode $ \h -> receiveFile h getBody size
+            digest' <- liftIO $ LC.sha512Hash <$> LB.readFile fPath
+            when (digest' /= digest) $ throwError DIGEST
+
     sendServerFile :: FileRec -> RcvPublicDhKey -> M (FileResponse, Maybe ServerFile)
     sendServerFile FileRec {filePath, fileInfo = FileInfo {size}} rKey = do
       readTVarIO filePath >>= \case

src/Simplex/FileTransfer/Transport.hs

@@ -11,7 +11,6 @@ module Simplex.FileTransfer.Transport
   )
 where
 
-import Control.Monad.Except
 import Data.ByteString.Builder (Builder, byteString)
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
@@ -25,14 +24,16 @@ supportedFileServerVRange :: VersionRange
 supportedFileServerVRange = mkVersionRange 1 1
 
 sendFile :: Handle -> (Builder -> IO ()) -> Word32 -> IO ()
-sendFile _ _ 0 = pure ()
-sendFile h send sz = do
-  B.hGet h xftpBlockSize >>= \case
-    "" -> when (sz /= 0) ioe_EOF
-    ch -> do
-      let ch' = B.take (fromIntegral sz) ch -- sz >= xftpBlockSize
-      send (byteString ch')
-      sendFile h send $ sz - fromIntegral (B.length ch')
+sendFile h send = go
+  where
+    go 0 = pure ()
+    go sz =
+      B.hGet h xftpBlockSize >>= \case
+        "" -> ioe_EOF
+        ch -> do
+          let ch' = B.take (fromIntegral sz) ch -- sz >= xftpBlockSize
+          send $ byteString ch'
+          go $ sz - fromIntegral (B.length ch')
 
 receiveFile :: Handle -> (Int -> IO ByteString) -> Word32 -> IO (Either XFTPErrorType ())
 receiveFile h receive = go

tests/XFTPServerTests.hs

@@ -6,14 +6,17 @@
 module XFTPServerTests where
 
 import AgentTests.FunctionalAPITests (runRight_)
-import Control.Monad.IO.Class (liftIO)
+import Control.Monad.Except
 import Crypto.Random (getRandomBytes)
 import qualified Data.ByteString.Base64.URL as B64
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
+import qualified Data.ByteString.Lazy.Char8 as LB
 import Simplex.FileTransfer.Client
-import Simplex.FileTransfer.Protocol (FileInfo (..))
+import Simplex.FileTransfer.Protocol (FileInfo (..), XFTPErrorType (..))
+import Simplex.Messaging.Client (ProtocolClientError (..))
 import qualified Simplex.Messaging.Crypto as C
+import qualified Simplex.Messaging.Crypto.Lazy as LC
 import Simplex.Messaging.Protocol (SenderId)
 import System.Directory (createDirectoryIfMissing, removeDirectoryRecursive)
 import System.FilePath ((</>))
@@ -30,6 +33,9 @@ xftpServerTests =
 chSize :: Num n => n
 chSize = 256 * 1024
 
+testChunkPath :: FilePath
+testChunkPath = "tests/tmp/chunk1"
+
 createTestChunk :: FilePath -> IO ByteString
 createTestChunk fp = do
   bytes <- getRandomBytes chSize
@@ -45,13 +51,21 @@ testFileChunkDelivery =
     (sndKey, spKey) <- C.generateSignatureKeyPair C.SEd25519
     (rcvKey, rpKey) <- C.generateSignatureKeyPair C.SEd25519
     (rDhKey, _rpDhKey) <- C.generateKeyPair'
-    bytes <- createTestChunk "tests/tmp/chunk1"
+    bytes <- createTestChunk testChunkPath
     xftpTest $ \c -> runRight_ $ do
-      let file = FileInfo {sndKey, size = chSize, digest = "abc="}
+      digest <- liftIO $ LC.sha512Hash <$> LB.readFile testChunkPath
+      let file = FileInfo {sndKey, size = chSize, digest}
+          chunkSpec = XFTPChunkSpec {filePath = testChunkPath, chunkOffset = 0, chunkSize = chSize}
       (sId, [rId]) <- createXFTPChunk c spKey file [rcvKey]
-      uploadXFTPChunk c spKey sId $ XFTPChunkSpec {filePath = "tests/tmp/chunk1", chunkOffset = 0, chunkSize = chSize}
+      uploadXFTPChunk c spKey sId chunkSpec
+      (sId', _) <- createXFTPChunk c spKey file {digest = digest <> "_wrong"} [rcvKey]
+      uploadXFTPChunk c spKey sId' chunkSpec
+        `catchError` (liftIO . (`shouldBe` PCEProtocolError DIGEST))
       liftIO $ readChunk sId `shouldReturn` bytes
       (_sDhKey, chunkBody) <- downloadXFTPChunk c rpKey rId rDhKey
-      receiveXFTPChunk chunkBody "tests/tmp/received_chunk1" chSize
+      receiveXFTPChunk chunkBody "tests/tmp/received_chunk1" chSize (digest <> "_wrong")
+        `catchError` (liftIO . (`shouldBe` PCEResponseError DIGEST))
+      (_sDhKey, chunkBody') <- downloadXFTPChunk c rpKey rId rDhKey
+      receiveXFTPChunk chunkBody' "tests/tmp/received_chunk1" chSize digest
       liftIO $ B.readFile "tests/tmp/received_chunk1" `shouldReturn` bytes
       pure ()