From 1473d2fc72c2dcf26869134d212c8fecddf99e94 Mon Sep 17 00:00:00 2001 From: Moritz Angermann Date: Sat, 17 Oct 2026 08:25:02 +0900 Subject: [PATCH] Remove XFTP receive chunks before reporting completion --- src/Simplex/FileTransfer/Agent.hs | 4 +- tests/XFTPAgent.hs | 65 ++++++++++++++++++++++++++++++- 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/src/Simplex/FileTransfer/Agent.hs b/src/Simplex/FileTransfer/Agent.hs index f31922117..9f41c208d 100644 --- a/src/Simplex/FileTransfer/Agent.hs +++ b/src/Simplex/FileTransfer/Agent.hs @@ -313,8 +313,10 @@ runXFTPRcvLocalWorker c Worker {doWork} = do void $ liftError (FILE . FILE_IO . show) $ decryptChunks encSize chunkPaths key nonce $ \_ -> pure destFile case redirect of Nothing -> do - notify c rcvFileEntityId $ RFDONE fsSavePath + -- RFDONE lets the consumer delete the shared receive prefix. + -- Finish chunk cleanup before publishing completion. lift $ forM_ tmpPath (removePath <=< toFSFilePath) + notify c rcvFileEntityId $ RFDONE fsSavePath liftIO $ waitUntilForeground c withStore' c (`updateRcvFileComplete` rcvFileId) Just RcvFileRedirect {redirectFileInfo, redirectDbId} -> do diff --git a/tests/XFTPAgent.hs b/tests/XFTPAgent.hs index e55b9483c..4df3a4fa0 100644 --- a/tests/XFTPAgent.hs +++ b/tests/XFTPAgent.hs @@ -13,6 +13,7 @@ module XFTPAgent where import AgentTests.FunctionalAPITests (get, rfGet, runRight, runRight_, sfGet, withAgent, testServerInformation) +import Control.Concurrent.STM (retry) import Control.Logger.Simple import Control.Monad import Control.Monad.Except @@ -25,6 +26,7 @@ import qualified Data.Map.Strict as M import Data.Maybe (fromJust) import Data.Time.Clock (addUTCTime, getCurrentTime, nominalDay) import Data.Time.Clock.System (getSystemTime, systemSeconds) +import GHC.Conc (BlockReason (BlockedOnSTM), ThreadStatus (ThreadBlocked), threadStatus) import SMPAgentClient (agentCfg, initAgentServers, testDB, testDB2, testDB3) import SMPClient (xit'') import Simplex.FileTransfer.Client (XFTPClientConfig (..)) @@ -37,6 +39,7 @@ import Simplex.FileTransfer.Types (RcvFileId, SndFileId) import Simplex.Messaging.Agent (AgentClient, testProtocolServer, xftpDeleteRcvFile, xftpDeleteSndFileInternal, xftpDeleteSndFileRemote, xftpPrepareReceiveFile, xftpPrepareSendFile, xftpReceiveFile, xftpSendDescription, xftpStartReceiveFile, xftpStartSendFile, xftpStartWorkers) import qualified Simplex.Messaging.Agent as XA import Simplex.Messaging.Agent.Client (ProtocolTestFailure (..), ProtocolTestStep (..)) +import qualified Simplex.Messaging.Agent.Client as AC import Simplex.Messaging.Agent.Env.SQLite (AgentConfig, xftpCfg) import qualified Simplex.Messaging.Agent.Env.SQLite as AEnv import Simplex.Messaging.Agent.Protocol hiding (SFDONE) @@ -52,9 +55,11 @@ import Simplex.Messaging.Protocol (BasicAuth, NetworkError (..), ProtoServerWith import Simplex.Messaging.Server.Expiration (ExpirationConfig (..)) import Simplex.Messaging.Server.Information (ServerPublicInfo) import Simplex.Messaging.Transport (EntitlementConfig (..)) -import Simplex.Messaging.Util (tshow) +import qualified Simplex.Messaging.TMap as TM +import Simplex.Messaging.Util (ifM, tshow) import System.Directory (doesDirectoryExist, doesFileExist, getFileSize, listDirectory, removeFile) import System.FilePath (()) +import System.Mem.Weak (deRefWeak) import Test.Hspec hiding (fit, it, xit) import UnliftIO import UnliftIO.Concurrent @@ -75,6 +80,7 @@ xftpAgentTests = it "should send and receive with encrypted local files" testXFTPAgentSendReceiveEncrypted it "should send and receive large file with a redirect" testXFTPAgentSendReceiveRedirect it "should send and receive small file without a redirect" testXFTPAgentSendReceiveNoRedirect + it "should remove received chunks before publishing completion" testXFTPAgentCompletionCleanup it "should receive prepared file only after it is started" testXFTPAgentPrepareReceive it "should send prepared file only after it is started" testXFTPAgentPrepareSend it "should extend storage time with an entitlement proof and report the granted expiry" $ \_ -> testXFTPAgentEntitlement @@ -277,6 +283,63 @@ testXFTPAgentSendReceiveNoRedirect = withXFTPServer $ do inBytes <- B.readFile filePathIn B.readFile out `shouldReturn` inBytes +-- A full event queue is a barrier at RFDONE, after real download/decryption. +-- Holding the not-yet-started local worker lets us install that barrier after +-- draining RFPROG, without sleeps or changing the transfer implementation. +testXFTPAgentCompletionCleanup :: (HasCallStack, HasTestEnv) => AFStoreType -> IO () +testXFTPAgentCompletionCleanup = withXFTPServer $ do + filePathIn <- createRandomFile_ (kb 64) "completion-order" + sentBytes <- B.readFile filePathIn + withAgent 1 agentCfg initAgentServers testDB $ \sndr -> do + runRight_ $ xftpStartWorkers sndr (Just senderFiles) + _ <- runRight $ xftpSendFile sndr 1 (CryptoFile filePathIn Nothing) 1 + sfProgress sndr (kb 128) + (_, _, SFDONE _ (description : _)) <- sfGet sndr + withAgent 2 agentCfg {AEnv.tbqSize = 1} initAgentServers testDB2 $ \rcp -> do + let AEnv.XFTPAgent {AEnv.xftpWorkDir, AEnv.xftpRcvWorkers} = AEnv.xftpAgent $ AC.agentEnv rcp + AEnv.Worker {AEnv.action, AEnv.doWork, AEnv.restarts} <- atomically $ do + writeTVar xftpWorkDir $ Just recipientFiles + workerId <- stateTVar (AC.workerSeq rcp) $ \next -> (next, next + 1) + w <- AEnv.Worker workerId <$> newTMVar () <*> newEmptyTMVar <*> newTVar (AEnv.RestartCount 0 0) + TM.insert Nothing w xftpRcvWorkers + pure w + let releaseWorker = atomically $ void $ tryPutTMVar action Nothing + flip finally releaseWorker $ do + rfId <- runRight $ xftpReceiveFile rcp 1 description Nothing True + rfProgress rcp (kb 128) + [prefixDir] <- listDirectory recipientFiles + let prefixPath = recipientFiles prefixDir + tmpPath = prefixPath "xftp.encrypted" + outputPath = prefixPath "xftp.decrypted" + doesDirectoryExist tmpPath `shouldReturn` True + -- No producer can publish completion until this sentinel is consumed. + atomically $ writeTBQueue (AC.subQ rcp) ("barrier", rfId, AEvt SAERcvFile $ RFPROG 0 (kb 128)) + releaseWorker + weak <- atomically $ readTMVar action >>= maybe retry pure + Just tid <- deRefWeak weak + let awaitPublish = do + blocked <- (== ThreadBlocked BlockedOnSTM) <$> threadStatus tid + decrypted <- + if blocked + then ifM (doesFileExist outputPath) ((== sentBytes) <$> B.readFile outputPath) (pure False) + else pure False + unless (blocked && decrypted) $ yield >> awaitPublish + timeout 10000000 awaitPublish `shouldReturn` Just () + -- This fails on the old order: its worker is blocked in notify while + -- the encrypted chunk directory still exists. + doesDirectoryExist tmpPath `shouldReturn` False + rfGet rcp `shouldReturn` ("barrier", rfId, RFPROG 0 (kb 128)) + rfGet rcp `shouldReturn` ("", rfId, RFDONE outputPath) + -- Model the consumer immediately deleting its completed receive. + xftpDeleteRcvFile rcp rfId + let awaitIdle = do + idle <- atomically $ isEmptyTMVar doWork + blocked <- (== ThreadBlocked BlockedOnSTM) <$> threadStatus tid + unless (idle && blocked) $ yield >> awaitIdle + timeout 10000000 awaitIdle `shouldReturn` Just () + AEnv.restartCount <$> readTVarIO restarts `shouldReturn` 0 + atomically (tryReadTBQueue $ AC.subQ rcp) `shouldReturn` Nothing + testXFTPAgentPrepareReceive :: (HasCallStack, HasTestEnv) => AFStoreType -> IO () testXFTPAgentPrepareReceive = withXFTPServer $ do filePath <- createRandomFile