xftp server: release storage of interrupted uploads

src/Simplex/FileTransfer/Server.hs

@@ -573,7 +573,7 @@ processXFTPRequest ent HTTP2Body {bodyPart} = \case
     receiveServerFile FileRec {senderId, fileInfo = FileInfo {size, digest}, filePath} = case bodyPart of
       Nothing -> pure $ FRErr SIZE
       -- TODO validate body size from request before downloading, once it's populated
-      Just getBody -> skipCommitted $ ifM reserve receive (pure $ FRErr QUOTA)
+      Just getBody -> skipCommitted $ mask $ \restore -> ifM reserve (receive restore) (pure $ FRErr QUOTA)
         where
           -- having a filePath means the file is already uploaded and committed, must not change anything
           skipCommitted = ifM (isJust <$> readTVarIO filePath) (liftIO $ drain $ fromIntegral size)
@@ -591,10 +591,10 @@ processXFTPRequest ent HTTP2Body {bodyPart} = \case
             quota <- asks $ fromMaybe maxBound . fileSizeQuota . config
             atomically . stateTVar us $
               \used -> let used' = used + fromIntegral size in if used' <= quota then (True, used') else (False, used)
-          receive = do
+          receive restore = do
             path <- asks $ filesPath . config
             let fPath = path </> serverFileName senderId
-            receiveChunk (XFTPRcvChunkSpec fPath size digest) >>= \case
+            (restore (receiveChunk $ XFTPRcvChunkSpec fPath size digest) `onException` release fPath) >>= \case
               Right () -> do
                 stats <- asks serverStats
                 st <- asks fileStore
@@ -605,16 +605,12 @@ processXFTPRequest ent HTTP2Body {bodyPart} = \case
                     incFileStat filesCount
                     liftIO $ atomicModifyIORef'_ (filesSize stats) (+ fromIntegral size)
                     pure FROk
-                  Left _e -> do
-                    us <- asks usedStorage
-                    atomically $ modifyTVar' us $ subtract (fromIntegral size)
-                    liftIO $ whenM (doesFileExist fPath) (removeFile fPath) `catch` logFileError
-                    pure $ FRErr AUTH
-              Left e -> do
-                us <- asks usedStorage
-                atomically $ modifyTVar' us $ subtract (fromIntegral size)
-                liftIO $ whenM (doesFileExist fPath) (removeFile fPath) `catch` logFileError
-                pure $ FRErr e
+                  Left _e -> FRErr AUTH <$ release fPath
+              Left e -> FRErr e <$ release fPath
+          release fPath = do
+            us <- asks usedStorage
+            atomically $ modifyTVar' us $ subtract (fromIntegral size)
+            liftIO $ whenM (doesFileExist fPath) (removeFile fPath) `catch` logFileError
           receiveChunk spec = do
             t <- asks $ fileTimeout . config
             liftIO $ fromMaybe (Left TIMEOUT) <$> timeout t (runExceptT $ receiveFile getBody spec)
@@ -664,7 +660,7 @@ blockServerFile fr@FileRec {senderId} info = do
   deleteOrBlockServerFile_ fr filesBlocked $ \st -> blockFile st senderId info True
 
 deleteOrBlockServerFile_ :: FileStoreClass s => FileRec -> (FileServerStats -> IORef Int) -> (s -> IO (Either XFTPErrorType ())) -> M s (Either XFTPErrorType ())
-deleteOrBlockServerFile_ FileRec {filePath, fileInfo} stat storeAction = runExceptT $ do
+deleteOrBlockServerFile_ FileRec {filePath, fileInfo} stat storeAction = mask_ $ runExceptT $ do
   path <- readTVarIO filePath
   stats <- asks serverStats
   ExceptT $ first (\(_ :: SomeException) -> FILE_IO) <$> try (forM_ path $ \p -> whenM (doesFileExist p) (removeFile p >> deletedStats stats))

tests/XFTPServerTests.hs

@@ -9,7 +9,7 @@
 
 module XFTPServerTests (xftpServerTests) where
 
-import AgentTests.FunctionalAPITests (runRight_)
+import AgentTests.FunctionalAPITests (runRight, runRight_)
 import Control.Concurrent (threadDelay)
 import Control.Exception (SomeException, bracket)
 import Control.Monad
@@ -35,7 +35,7 @@ import ServerTests (logSize)
 import Simplex.FileTransfer.Client hiding (createXFTPChunk)
 import qualified Simplex.FileTransfer.Client as A
 import Simplex.FileTransfer.Description (kb)
-import Simplex.FileTransfer.Protocol (FileInfo (..), XFTPFileId, xftpBlockSize)
+import Simplex.FileTransfer.Protocol (FileCmd (..), FileCommand (..), FileInfo (..), SFileParty (..), XFTPFileId, xftpBlockSize, xftpEncodeAuthTransmission)
 import Simplex.FileTransfer.Server.Env (AFStoreType, XFTPServerConfig (..))
 import Simplex.FileTransfer.Transport (XFTPClientHandshake (..), XFTPClientHello (..), XFTPErrorType (..), XFTPRcvChunkSpec (..), XFTPServerHandshake (..), pattern VersionXFTP)
 import Simplex.Messaging.Client (NetworkConfig (..), NetworkTimeout (..), ProtocolClientError (..), defaultNetworkConfig)
@@ -50,10 +50,10 @@ import Simplex.Messaging.Transport.HTTP2 (HTTP2Body (..))
 import qualified Simplex.Messaging.Transport.HTTP2.Client as HC
 import Simplex.Messaging.Transport.Server (loadFileFingerprint)
 import Simplex.Messaging.Transport.Shared (ChainCertificates (..), chainIdCaCerts)
-import System.Directory (removeFile)
+import System.Directory (listDirectory, removeFile)
 import System.FilePath ((</>))
 import Test.Hspec hiding (fit, it)
-import UnliftIO.Async (async, cancel)
+import UnliftIO.Async (async, cancel, pooledReplicateConcurrentlyN, race_)
 import UnliftIO.STM
 import UnliftIO.Timeout (timeout)
 import Util
@@ -72,6 +72,7 @@ xftpServerTests = do
       it "should disconnect inactive clients" testInactiveClientExpiration
       it "should close connection when connecting client is cancelled" $ \_ -> testCancelledClientConnect
       it "should not allow uploading chunks after specified storage quota" testFileStorageQuota
+      it "should release storage of interrupted uploads" testInterruptedUploadStorage
       it "should store file records to log and restore them after server restart" testFileLog
       describe "XFTP basic auth" $ do
         --                                               allow FNEW | server auth | clnt auth | success
@@ -278,6 +279,31 @@ testFileStorageQuota fsType = withXFTPServerConfigOn (updateXFTPCfg (cfgFS fsTyp
     uploadXFTPChunk c spKey sId3 chunkSpec
     download rId3
 
+testInterruptedUploadStorage :: HasTestEnv => AFStoreType -> Expectation
+testInterruptedUploadStorage fsType = withXFTPServerConfigOn (updateXFTPCfg (cfgFS fsType) $ \c -> c {fileSizeQuota = Just $ chSize * 10}) $
+  \_ -> do
+    g <- C.newRandom
+    (sndKey, spKey) <- atomically $ C.generateAuthKeyPair C.SEd25519 g
+    (rcvKey, _) <- atomically $ C.generateAuthKeyPair C.SEd25519 g
+    void $ createTestChunk testChunkPath
+    digest <- LC.sha256Hash <$> LB.readFile testChunkPath
+    let file = FileInfo {sndKey, size = chSize, digest}
+        chunkSpec = XFTPChunkSpec {filePath = testChunkPath, chunkOffset = 0, chunkSize = chSize}
+        interruptUpload = testXFTPClient $ \c@XFTPClient {http2Client, thParams} -> do
+          (sId, _) <- runRight $ createXFTPChunk c spKey file [rcvKey] Nothing
+          t <- either (error . show) pure $ xftpEncodeAuthTransmission thParams spKey ("", sId, FileCmd SFSender FPUT)
+          let req = H2.requestStreaming "POST" "/" [] $ \send flush -> do
+                send $ byteString t
+                send $ byteString $ B.replicate 1024 'x'
+                flush
+                threadDelay maxBound
+          race_ (HC.sendRequest http2Client req Nothing) (threadDelay 20000)
+          sId <$ closeXFTPClient c
+    sIds <- pooledReplicateConcurrentlyN 10 300 interruptUpload
+    threadDelay 1000000
+    listDirectory xftpServerFiles `shouldReturn` []
+    testXFTPClient $ \c -> runRight_ $ forM_ (take 10 sIds) $ \sId -> uploadXFTPChunk c spKey sId chunkSpec
+
 testFileLog :: HasTestEnv => AFStoreType -> Expectation
 testFileLog _ = do
   g <- C.newRandom