xftp: add quota param to server CLI, restrict chunk sizes (#659)

* xftp: add quota param to server CLI

* only allow certain file sizes, fix tests

src/Simplex/FileTransfer/Client.hs

@@ -19,6 +19,7 @@ import Data.List.NonEmpty (NonEmpty (..))
 import Data.Word (Word32)
 import qualified Network.HTTP.Types as N
 import qualified Network.HTTP2.Client as H
+import Simplex.FileTransfer.Description (mb)
 import Simplex.FileTransfer.Protocol
 import Simplex.FileTransfer.Transport
 import Simplex.Messaging.Client
@@ -122,7 +123,6 @@ sendXFTPCommand XFTPClient {config, http2Client = http2@HTTP2Client {sessionId}}
       _ -> pure (r, body)
     Left e -> throwError $ PCEResponseError e
   where
-    mb = 1024 * 1024
     streamBody :: ByteString -> (Builder -> IO ()) -> IO () -> IO ()
     streamBody t send done = do
       send $ byteString t

src/Simplex/FileTransfer/Client/Main.hs

@@ -52,7 +52,7 @@ xftpClientVersion :: String
 xftpClientVersion = "0.1.0"
 
 defaultChunkSize :: Word32
-defaultChunkSize = 8 * mb
+defaultChunkSize = 4 * mb
 
 smallChunkSize :: Word32
 smallChunkSize = 1 * mb
@@ -63,9 +63,6 @@ fileSizeLen = 8
 authTagSize :: Int64
 authTagSize = fromIntegral C.authTagSize
 
-mb :: Num a => a
-mb = 1024 * 1024
-
 newtype CLIError = CLIError String
   deriving (Eq, Show, Exception)
 

src/Simplex/FileTransfer/Description.hs

@@ -9,7 +9,6 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
 
 module Simplex.FileTransfer.Description
   ( FileDescription (..),
@@ -28,6 +27,9 @@ module Simplex.FileTransfer.Description
     validateFileDescription,
     groupReplicasByServer,
     replicaServer,
+    kb,
+    mb,
+    gb,
   )
 where
 
@@ -46,6 +48,7 @@ import Data.List (foldl', groupBy, sortOn)
 import Data.Map (Map)
 import qualified Data.Map as M
 import Data.Maybe (fromMaybe)
+import Data.String
 import Data.Word (Word32)
 import qualified Data.Yaml as Y
 import GHC.Generics (Generic)
@@ -195,13 +198,13 @@ newtype FileSize a = FileSize {unFileSize :: a}
 instance (Integral a, Show a) => StrEncoding (FileSize a) where
   strEncode (FileSize b)
     | b' /= 0 = bshow b
-    | kb' /= 0 = bshow kb <> "kb"
-    | mb' /= 0 = bshow mb <> "mb"
-    | otherwise = bshow gb <> "gb"
+    | ks' /= 0 = bshow ks <> "kb"
+    | ms' /= 0 = bshow ms <> "mb"
+    | otherwise = bshow gs <> "gb"
     where
-      (kb, b') = b `divMod` 1024
-      (mb, kb') = kb `divMod` 1024
-      (gb, mb') = mb `divMod` 1024
+      (ks, b') = b `divMod` 1024
+      (ms, ks') = ks `divMod` 1024
+      (gs, ms') = ms `divMod` 1024
   strP =
     FileSize
       <$> A.choice
@@ -210,10 +213,18 @@ instance (Integral a, Show a) => StrEncoding (FileSize a) where
           (kb *) <$> A.decimal <* "kb",
           A.decimal
         ]
-    where
-      kb = 1024
-      mb = 1024 * kb
-      gb = 1024 * mb
+
+kb :: Integral a => a
+kb = 1024
+
+mb :: Integral a => a
+mb = 1024 * kb
+
+gb :: Integral a => a
+gb = 1024 * mb
+
+instance (Integral a, Show a) => IsString (FileSize a) where
+  fromString = either error id . strDecode . B.pack
 
 groupReplicasByServer :: FileSize Word32 -> [FileChunk] -> [[FileServerReplica]]
 groupReplicasByServer defChunkSize =

src/Simplex/FileTransfer/Server.hs

@@ -238,8 +238,10 @@ processXFTPRequest HTTP2Body {bodyPart} = \case
     noFile resp = pure (resp, Nothing)
     createFile :: FileStore -> FileInfo -> NonEmpty RcvPublicVerifyKey -> M FileResponse
     createFile st file rks = do
-      ts <- liftIO getSystemTime
       r <- runExceptT $ do
+        sizes <- asks $ allowedChunkSizes . config
+        unless (size file `elem` sizes) $ throwError SIZE
+        ts <- liftIO getSystemTime
         -- TODO validate body empty
         sId <- ExceptT $ addFileRetry 3 ts
         rcps <- mapM (ExceptT . addRecipientRetry 3 sId) rks

src/Simplex/FileTransfer/Server/Env.hs

@@ -15,6 +15,7 @@ import Crypto.Random
 import Data.Int (Int64)
 import Data.List.NonEmpty (NonEmpty)
 import Data.Time.Clock (getCurrentTime)
+import Data.Word (Word32)
 import Data.X509.Validation (Fingerprint (..))
 import Network.Socket
 import qualified Network.TLS as T
@@ -37,6 +38,8 @@ data XFTPServerConfig = XFTPServerConfig
     filesPath :: FilePath,
     -- | server storage quota
     fileSizeQuota :: Maybe Int64,
+    -- | allowed file chunk sizes
+    allowedChunkSizes :: [Word32],
     -- | set to False to prohibit creating new files
     allowNewFiles :: Bool,
     -- | simple password that the clients need to pass in handshake to be able to create new files

src/Simplex/FileTransfer/Server/Main.hs

@@ -7,17 +7,20 @@
 
 module Simplex.FileTransfer.Server.Main where
 
+import qualified Data.ByteString.Char8 as B
 import Data.Either (fromRight)
 import Data.Functor (($>))
 import Data.Ini (lookupValue, readIniFile)
+import Data.Int (Int64)
 import Data.Maybe (fromMaybe)
 import qualified Data.Text as T
 import Network.Socket (HostName)
 import Options.Applicative
-import Simplex.FileTransfer.Description (FileSize (..))
+import Simplex.FileTransfer.Description (FileSize (..), kb, mb)
 import Simplex.FileTransfer.Server (runXFTPServer)
 import Simplex.FileTransfer.Server.Env (XFTPServerConfig (..), defaultFileExpiration)
 import qualified Simplex.Messaging.Crypto as C
+import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Protocol (ProtoServerWithAuth (..), pattern XFTPServer)
 import Simplex.Messaging.Server.CLI
 import Simplex.Messaging.Transport.Client (TransportHost (..))
@@ -51,7 +54,7 @@ xftpServerCLI cfgPath logPath = do
     defaultServerPort = "443"
     executableName = "file-server"
     storeLogFilePath = combine logPath "file-server-store.log"
-    initializeServer InitOptions {enableStoreLog, signAlgorithm, ip, fqdn, filesPath} = do
+    initializeServer InitOptions {enableStoreLog, signAlgorithm, ip, fqdn, filesPath, fileSizeQuota} = do
       clearDirIfExists cfgPath
       clearDirIfExists logPath
       createDirectoryIfMissing True cfgPath
@@ -95,7 +98,7 @@ xftpServerCLI cfgPath logPath = do
                \\n\
                \[FILES]\n"
             <> ("path: " <> filesPath <> "\n")
-            <> "# storage_quota: 100gb\n"
+            <> ("storage_quota: " <> B.unpack (strEncode fileSizeQuota) <> "\n")
     runServer ini = do
       hSetBuffering stdout LineBuffering
       hSetBuffering stderr LineBuffering
@@ -128,6 +131,7 @@ xftpServerCLI cfgPath logPath = do
               storeLogFile = enableStoreLog $> storeLogFilePath,
               filesPath = T.unpack $ strictIni "FILES" "path" ini,
               fileSizeQuota = either error unFileSize <$> strDecodeIni "FILES" "storage_quota" ini,
+              allowedChunkSizes = [256 * kb, 1 * mb, 4 * mb],
               allowNewFiles = fromMaybe True $ iniOnOff "AUTH" "new_files" ini,
               newFileBasicAuth = either error id <$> strDecodeIni "AUTH" "create_password" ini,
               fileExpiration = Just defaultFileExpiration,
@@ -151,7 +155,8 @@ data InitOptions = InitOptions
     signAlgorithm :: SignAlgorithm,
     ip :: HostName,
     fqdn :: Maybe HostName,
-    filesPath :: FilePath
+    filesPath :: FilePath,
+    fileSizeQuota :: FileSize Int64
   }
   deriving (Show)
 
@@ -201,3 +206,9 @@ cliCommandP cfgPath logPath iniFile =
               <> help "Path to the directory to store files"
               <> metavar "PATH"
           )
+        <*> strOption
+          ( long "quota"
+              <> short 'q'
+              <> help "File storage quota (e.g. 100gb)"
+              <> metavar "QUOTA"
+          )

src/Simplex/FileTransfer/Transport.hs

@@ -25,7 +25,7 @@ import qualified Data.ByteString.Char8 as B
 import qualified Data.ByteString.Lazy.Char8 as LB
 import Data.Word (Word32)
 import GHC.IO.Handle.Internals (ioe_EOF)
-import Simplex.FileTransfer.Protocol (XFTPErrorType (..), xftpBlockSize)
+import Simplex.FileTransfer.Protocol (XFTPErrorType (..))
 import qualified Simplex.Messaging.Crypto as C
 import qualified Simplex.Messaging.Crypto.Lazy as LC
 import Simplex.Messaging.Version

tests/CLITests.hs

@@ -91,7 +91,7 @@ ntfServerTest storeLog = do
 
 xftpServerTest :: Bool -> IO ()
 xftpServerTest storeLog = do
-  capture_ (withArgs (["init", "-p tests/tmp"] <> ["-l" | storeLog]) $ xftpServerCLI fileCfgPath fileLogPath)
+  capture_ (withArgs (["init", "-p", "tests/tmp", "-q", "10gb"] <> ["-l" | storeLog]) $ xftpServerCLI fileCfgPath fileLogPath)
     >>= (`shouldSatisfy` (("Server initialized, you can modify configuration in " <> fileCfgPath <> "/file-server.ini") `isPrefixOf`))
   Right ini <- readIniFile $ fileCfgPath <> "/file-server.ini"
   lookupValue "STORE_LOG" "enable" ini `shouldBe` Right (if storeLog then "on" else "off")

tests/FileDescriptionTests.hs

@@ -30,9 +30,6 @@ fileDescPath = "tests/fixtures/file_description.yaml"
 tmpFileDescPath :: FilePath
 tmpFileDescPath = "tests/tmp/file_description.yaml"
 
-mb :: Num a => a
-mb = 1024 * 1024
-
 testSbKey :: C.SbKey
 testSbKey = either error id $ strDecode "00n8p1tJq5E-SGnHcYTOrS4A9I07gTA_WFD6MTFFFOY="
 

tests/XFTPCLI.hs

@@ -4,6 +4,7 @@ import Control.Exception (bracket_)
 import qualified Data.ByteString as LB
 import Data.List (isInfixOf, isPrefixOf, isSuffixOf)
 import Simplex.FileTransfer.Client.Main (xftpClientCLI)
+import Simplex.FileTransfer.Description (mb)
 import System.Directory (createDirectoryIfMissing, getFileSize, listDirectory, removeDirectoryRecursive)
 import System.Environment (withArgs)
 import System.FilePath ((</>))
@@ -31,15 +32,12 @@ senderFiles = "tests/tmp/xftp-sender-files"
 recipientFiles :: FilePath
 recipientFiles = "tests/tmp/xftp-recipient-files"
 
-mb :: Num a => a
-mb = 1024 * 1024
-
 testXFTPCLISendReceive :: IO ()
 testXFTPCLISendReceive = withXFTPServer $ do
   let filePath = senderFiles </> "testfile"
-  xftp ["rand", filePath, "19mb"] `shouldReturn` ["File created: " <> filePath]
+  xftp ["rand", filePath, "17mb"] `shouldReturn` ["File created: " <> filePath]
   file <- LB.readFile filePath
-  getFileSize filePath `shouldReturn` 19 * mb
+  getFileSize filePath `shouldReturn` 17 * mb
   let fdRcv1 = filePath <> ".xftp" </> "rcv1.xftp"
       fdRcv2 = filePath <> ".xftp" </> "rcv2.xftp"
       fdSnd = filePath <> ".xftp" </> "snd.xftp.private"
@@ -63,7 +61,7 @@ testXFTPCLISendReceive = withXFTPServer $ do
     xftp params = lines <$> capture_ (withArgs params xftpClientCLI)
     testInfoFile fd party = do
       xftp ["info", fd]
-        `shouldReturn` [party <> " file description", "File download size: 20mb", "File server(s):", testXFTPServerStr <> ": 20mb"]
+        `shouldReturn` [party <> " file description", "File download size: 18mb", "File server(s):", testXFTPServerStr <> ": 18mb"]
     testReceiveFile fd fileName file = do
       progress : recvResult <- xftp ["recv", fd, recipientFiles, "--tmp=tests/tmp"]
       progress `shouldSatisfy` downloadProgress fileName
@@ -73,9 +71,9 @@ testXFTPCLISendReceive = withXFTPServer $ do
 testXFTPCLISendReceive2servers :: IO ()
 testXFTPCLISendReceive2servers = withXFTPServer . withXFTPServer2 $ do
   let filePath = senderFiles </> "testfile"
-  xftp ["rand", filePath, "19mb"] `shouldReturn` ["File created: " <> filePath]
+  xftp ["rand", filePath, "17mb"] `shouldReturn` ["File created: " <> filePath]
   file <- LB.readFile filePath
-  getFileSize filePath `shouldReturn` 19 * mb
+  getFileSize filePath `shouldReturn` 17 * mb
   let fdRcv1 = filePath <> ".xftp" </> "rcv1.xftp"
       fdRcv2 = filePath <> ".xftp" </> "rcv2.xftp"
       fdSnd = filePath <> ".xftp" </> "snd.xftp.private"
@@ -95,7 +93,7 @@ testXFTPCLISendReceive2servers = withXFTPServer . withXFTPServer2 $ do
     testReceiveFile fd fileName file = do
       partyStr : sizeStr : srvStr : srvs <- xftp ["info", fd]
       partyStr `shouldContain` "Recipient file description"
-      sizeStr `shouldBe` "File download size: 20mb"
+      sizeStr `shouldBe` "File download size: 18mb"
       srvStr `shouldBe` "File server(s):"
       case srvs of
         [srv1] -> any (`isInfixOf` srv1) [testXFTPServerStr, testXFTPServerStr2] `shouldBe` True
@@ -111,9 +109,9 @@ testXFTPCLISendReceive2servers = withXFTPServer . withXFTPServer2 $ do
 testXFTPCLIDelete :: IO ()
 testXFTPCLIDelete = withXFTPServer . withXFTPServer2 $ do
   let filePath = senderFiles </> "testfile"
-  xftp ["rand", filePath, "19mb"] `shouldReturn` ["File created: " <> filePath]
+  xftp ["rand", filePath, "17mb"] `shouldReturn` ["File created: " <> filePath]
   file <- LB.readFile filePath
-  getFileSize filePath `shouldReturn` 19 * mb
+  getFileSize filePath `shouldReturn` 17 * mb
   let fdRcv1 = filePath <> ".xftp" </> "rcv1.xftp"
       fdRcv2 = filePath <> ".xftp" </> "rcv2.xftp"
       fdSnd = filePath <> ".xftp" </> "snd.xftp.private"

tests/XFTPClient.hs

@@ -10,6 +10,7 @@ import Data.String (fromString)
 import Network.Socket (ServiceName)
 import SMPClient (serverBracket)
 import Simplex.FileTransfer.Client
+import Simplex.FileTransfer.Description
 import Simplex.FileTransfer.Server (runXFTPServerBlocking)
 import Simplex.FileTransfer.Server.Env (XFTPServerConfig (..), defaultFileExpiration)
 import Simplex.Messaging.Protocol (XFTPServer)
@@ -87,6 +88,7 @@ testXFTPServerConfig =
       storeLogFile = Nothing,
       filesPath = xftpServerFiles,
       fileSizeQuota = Nothing,
+      allowedChunkSizes = [128 * kb, 256 * kb, 1 * mb, 4 * mb],
       allowNewFiles = True,
       newFileBasicAuth = Nothing,
       fileExpiration = Just defaultFileExpiration,

tests/XFTPServerTests.hs

@@ -19,6 +19,7 @@ import qualified Data.ByteString.Lazy.Char8 as LB
 import Data.List (isInfixOf)
 import ServerTests (logSize)
 import Simplex.FileTransfer.Client
+import Simplex.FileTransfer.Description (kb)
 import Simplex.FileTransfer.Protocol (FileInfo (..), XFTPErrorType (..))
 import Simplex.FileTransfer.Server.Env (XFTPServerConfig (..))
 import Simplex.FileTransfer.Transport (XFTPRcvChunkSpec (..))
@@ -44,6 +45,7 @@ xftpServerTests =
       it "should delete file chunk (2 clients)" testFileChunkDelete2
       it "should acknowledge file chunk reception (1 client)" testFileChunkAck
       it "should acknowledge file chunk reception (2 clients)" testFileChunkAck2
+      it "should not allow chunks of wrong size" testWrongChunkSize
       it "should expire chunks after set interval" testFileChunkExpiration
       it "should not allow uploading chunks after specified storage quota" testFileStorageQuota
       it "should store file records to log and restore them after server restart" testFileLog
@@ -55,8 +57,8 @@ xftpServerTests =
         it "allowed with correct basic auth" $ testFileBasicAuth True (Just "pwd") (Just "pwd") True
         it "allowed with auth on server without auth" $ testFileBasicAuth True Nothing (Just "any") True
 
-chSize :: Num n => n
-chSize = 128 * 1024
+chSize :: Integral a => a
+chSize = 128 * kb
 
 testChunkPath :: FilePath
 testChunkPath = "tests/tmp/chunk1"
@@ -149,6 +151,16 @@ runTestFileChunkAck s r = do
   ackXFTPChunk r rpKey rId
     `catchError` (liftIO . (`shouldBe` PCEProtocolError AUTH))
 
+testWrongChunkSize :: Expectation
+testWrongChunkSize = xftpTest $ \c -> runRight_ $ do
+  (sndKey, spKey) <- liftIO $ C.generateSignatureKeyPair C.SEd25519
+  (rcvKey, _rpKey) <- liftIO $ C.generateSignatureKeyPair C.SEd25519
+  liftIO $ B.writeFile testChunkPath =<< getRandomBytes (96 * kb)
+  digest <- liftIO $ LC.sha512Hash <$> LB.readFile testChunkPath
+  let file = FileInfo {sndKey, size = 96 * kb, digest}
+  void (createXFTPChunk c spKey file [rcvKey] Nothing)
+    `catchError` (liftIO . (`shouldBe` PCEProtocolError SIZE))
+
 testFileChunkExpiration :: Expectation
 testFileChunkExpiration = withXFTPServerCfg testXFTPServerConfig {fileExpiration} $
   \_ -> testXFTPClient $ \c -> runRight_ $ do