fix: ignore pending XFTP files in storage accounting (#1814)

* fix: ignore pending XFTP files in storage accounting

* style

---------

Co-authored-by: Paul Bottinelli <paul.bottinelli@trailofbits.com>

src/Simplex/FileTransfer/Server/Store.hs

@@ -22,11 +22,11 @@ where
 import Data.Kind (Type)
 
 import Control.Concurrent.STM
-import Control.Monad (forM, void)
+import Control.Monad
 import qualified Data.Attoparsec.ByteString.Char8 as A
 import Data.Int (Int64)
 import qualified Data.Map.Strict as M
-import Data.Maybe (catMaybes)
+import Data.Maybe (catMaybes, isJust)
 import Data.Set (Set)
 import qualified Data.Set as S
 import Data.Word (Word32)
@@ -175,8 +175,10 @@ instance FileStoreClass STMFileStore where
           pure $ Just (sId, path, size)
         else pure Nothing
 
-  getUsedStorage STMFileStore {files} =
-    M.foldl' (\acc FileRec {fileInfo = FileInfo {size}} -> acc + fromIntegral size) 0 <$> readTVarIO files
+  getUsedStorage STMFileStore {files} = foldM addSize 0 =<< readTVarIO files
+    where
+      addSize acc FileRec {fileInfo = FileInfo {size}, filePath} =
+        ifM (isJust <$> readTVarIO filePath) (pure $! acc + fromIntegral size) (pure acc)
 
   getFileCount STMFileStore {files} = M.size <$> readTVarIO files
 

src/Simplex/FileTransfer/Server/Store/Postgres.hs

@@ -164,7 +164,7 @@ instance FileStoreClass PostgresFileStore where
 
   getUsedStorage st =
     withTransaction (dbStore st) $ \db -> do
-      [Only total] <- DB.query_ db "SELECT COALESCE(SUM(file_size::BIGINT), 0)::BIGINT FROM files"
+      [Only total] <- DB.query_ db "SELECT COALESCE(SUM(file_size::BIGINT), 0)::BIGINT FROM files WHERE file_path IS NOT NULL"
       pure total
 
   getFileCount st =

tests/CoreTests/XFTPStoreTests.hs

@@ -23,18 +23,21 @@ import Util
 import XFTPClient (testXFTPPostgresCfg)
 
 xftpStoreTests :: Spec
-xftpStoreTests = describe "PostgresFileStore operations" $ do
-  it "should add and get file by sender" testAddGetFileSender
-  it "should add and get file by recipient" testAddGetFileRecipient
-  it "should reject duplicate file" testDuplicateFile
-  it "should return AUTH for nonexistent file" testGetNonexistent
-  it "should set file path with IS NULL guard" testSetFilePath
-  it "should reject duplicate recipient" testDuplicateRecipient
-  it "should delete file and cascade recipients" testDeleteFileCascade
-  it "should block file and update status" testBlockFile
-  it "should ack file reception" testAckFile
-  it "should return expired files with limit" testExpiredFiles
-  it "should compute used storage and file count" testStorageAndCount
+xftpStoreTests = do
+  describe "STMFileStore operations" $
+    it "should compute committed used storage and file count" testSTMStorageAndCount
+  describe "PostgresFileStore operations" $ do
+    it "should add and get file by sender" testAddGetFileSender
+    it "should add and get file by recipient" testAddGetFileRecipient
+    it "should reject duplicate file" testDuplicateFile
+    it "should return AUTH for nonexistent file" testGetNonexistent
+    it "should set file path with IS NULL guard" testSetFilePath
+    it "should reject duplicate recipient" testDuplicateRecipient
+    it "should delete file and cascade recipients" testDeleteFileCascade
+    it "should block file and update status" testBlockFile
+    it "should ack file reception" testAckFile
+    it "should return expired files with limit" testExpiredFiles
+    it "should compute committed used storage and file count" testStorageAndCount
 
 xftpMigrationTests :: Spec
 xftpMigrationTests = describe "XFTP migration round-trip" $ do
@@ -201,16 +204,32 @@ testExpiredFiles = withPgStore $ \st -> do
 
 testStorageAndCount :: Expectation
 testStorageAndCount = withPgStore $ \st -> do
+  testStorageAndCountForStore st
+
+testSTMStorageAndCount :: Expectation
+testSTMStorageAndCount = do
+  st <- newFileStore () :: IO STMFileStore
+  testStorageAndCountForStore st
+  closeFileStore st
+
+testStorageAndCountForStore :: FileStoreClass s => s -> Expectation
+testStorageAndCountForStore st = do
   g <- C.newRandom
   (sndKey, _) <- atomically $ C.generateAuthKeyPair C.SEd25519 g
   getUsedStorage st `shouldReturn` 0
   getFileCount st `shouldReturn` 0
-  let fileInfo = testFileInfo sndKey
-  addFile st (EntityId "file_a__________") fileInfo testCreatedAt EntityActive `shouldReturn` Right ()
-  addFile st (EntityId "file_b__________") fileInfo testCreatedAt EntityActive `shouldReturn` Right ()
+  let fileInfoA = testFileInfo sndKey
+      fileInfoB = fileInfoA {size = 64000}
+      fileA = EntityId "file_a__________"
+      fileB = EntityId "file_b__________"
+  addFile st fileA fileInfoA testCreatedAt EntityActive `shouldReturn` Right ()
+  addFile st fileB fileInfoB testCreatedAt EntityActive `shouldReturn` Right ()
   getFileCount st `shouldReturn` 2
-  used <- getUsedStorage st
-  used `shouldBe` 256000 -- 128000 * 2
+  getUsedStorage st `shouldReturn` 0
+  setFilePath st fileA "/tmp/file_a" `shouldReturn` Right ()
+  getUsedStorage st `shouldReturn` 128000
+  setFilePath st fileB "/tmp/file_b" `shouldReturn` Right ()
+  getUsedStorage st `shouldReturn` 192000
 
 -- Migration round-trip test