Merge pull request #539 from IntersectMBO/jdral/corrupt-snapshot

jorisdral · web-flow · commit ac32ec919a38 · 2025-01-24T16:12:40.000Z
Test snapshot corruption
diff --git a/test/Test/Database/LSMTree/Internal/Snapshot/FS.hs b/test/Test/Database/LSMTree/Internal/Snapshot/FS.hs
@@ -1,17 +1,36 @@
+{-# LANGUAGE OverloadedStrings #-}
+
 -- | Tests for snapshots and their interaction with the file system
 module Test.Database.LSMTree.Internal.Snapshot.FS (tests) where
 
 import           Codec.CBOR.Read (DeserialiseFailure)
-import           Control.Exception
+import           Control.Monad.Class.MonadThrow
+import           Control.Monad.IOSim (runSimOrThrow)
+import           Control.Tracer
+import           Data.Bifunctor (Bifunctor (..))
+import           Data.Maybe (fromJust)
+import qualified Data.Set as Set
+import qualified Data.Vector as V
+import           Data.Word
+import           Database.LSMTree.Extras (showPowersOf10)
+import           Database.LSMTree.Extras.Generators ()
+import           Database.LSMTree.Internal
+import           Database.LSMTree.Internal.Config
 import           Database.LSMTree.Internal.CRC32C
+import           Database.LSMTree.Internal.Entry
+import           Database.LSMTree.Internal.Paths
+import           Database.LSMTree.Internal.Serialise
 import           Database.LSMTree.Internal.Snapshot
 import           Database.LSMTree.Internal.Snapshot.Codec
+import qualified System.FS.API as FS
 import           System.FS.API
 import           System.FS.Sim.Error hiding (genErrors)
 import qualified System.FS.Sim.MockFS as MockFS
 import           Test.Database.LSMTree.Internal.Snapshot.Codec ()
+import           Test.QuickCheck
+import           Test.QuickCheck.Gen (genDouble)
 import           Test.Tasty
-import           Test.Tasty.QuickCheck as QC
+import           Test.Tasty.QuickCheck
 import           Test.Util.FS
 
 tests :: TestTree
@@ -20,6 +39,7 @@ tests = testGroup "Test.Database.LSMTree.Internal.Snapshot.FS" [
         prop_fsRoundtripSnapshotMetaData
     , testProperty "prop_fault_fsRoundtripSnapshotMetaData"
         prop_fault_fsRoundtripSnapshotMetaData
+    , testProperty "prop_flipSnapshotBit" prop_flipSnapshotBit
     ]
 
 -- | @readFileSnapshotMetaData . writeFileSnapshotMetaData = id@
@@ -56,12 +76,12 @@ prop_fault_fsRoundtripSnapshotMetaData testErrs metadata =
     ioProperty $
     withSimErrorHasFS propNoOpenHandles MockFS.empty emptyErrors $ \hfs _fsVar errsVar -> do
       writeResult <-
-        try @FsError $
+        try @_ @FsError $
           withErrors errsVar (writeErrors testErrs) $
             writeFileSnapshotMetaData hfs metadataPath checksumPath metadata
 
       readResult <-
-        try @SomeException $
+        try @_ @SomeException $
           withErrors errsVar (readErrors testErrs) $
             readFileSnapshotMetaData hfs metadataPath checksumPath
 
@@ -137,3 +157,106 @@ instance Arbitrary TestErrors where
       [ TestErrors writeErrors' readErrors'
       | (writeErrors', readErrors') <- shrink (writeErrors, readErrors)
       ]
+
+{-------------------------------------------------------------------------------
+  Snapshot corruption
+-------------------------------------------------------------------------------}
+
+-- | A 'Double' in the @[0, 1)@ range.
+newtype Double_0_1 = Double_0_1 Double
+  deriving stock (Show, Eq)
+
+instance Arbitrary Double_0_1 where
+  arbitrary = Double_0_1 <$> genDouble
+  shrink (Double_0_1 x) = [Double_0_1 x' | x' <- shrink x, 0 <= x', x' < 1]
+
+-- TODO: an alternative to generating doubles a priori is to run the monadic
+-- code in @PropertyM (IOSim s)@, and then we can do quantification inside the
+-- monadic property using @pick@. This complicates matters, however, because
+-- functions like @withSimHasBlockIO@ and @withTable@ would have to run in
+-- @PropertyM (IOSim s)@ as well. It's not clear whether the refactoring is
+-- worth it.
+prop_flipSnapshotBit ::
+     Positive (Small Int)
+  -> V.Vector (Word64, Entry Word64 Word64)
+  -> Double_0_1 -- ^ Used to pick which file to corrupt
+  -> Double_0_1 -- ^ Used to pick which bit to flip in the file we picked
+  -> Property
+prop_flipSnapshotBit
+  (Positive (Small bufferSize))
+  es
+  (Double_0_1 pickFile)
+  (Double_0_1 pickBit) =
+    runSimOrThrow $
+    withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _fsVar ->
+    withSession nullTracer hfs hbio root $ \s ->
+    withTable s conf $ \t -> do
+      -- Create a table, populate it, and create a snapshot
+      updates resolve es' t
+      createSnap t
+
+      -- Pick a random file from the named snapshot directory
+      files <- listDirectoryRecursiveFiles hfs (getNamedSnapshotDir namedSnapDir)
+      let i = round (fromIntegral (Set.size files - 1) * pickFile)
+      let file = Set.elemAt i files
+      let path = getNamedSnapshotDir namedSnapDir </> file
+      -- Pick a random bit from the file that we want to corrupt
+      n <- withFile hfs path ReadMode $ hGetSize hfs
+      let j = round (fromIntegral (n * 8 - 1) * pickBit)
+
+      -- Some info for the test output
+      let
+        tabCorruptedFile = tabulate "Corrupted file" [show path]
+        counterCorruptedFile = counterexample ("Corrupted file: " ++ show path)
+        tabFlippedBit = tabulate "Flipped bit" [showPowersOf10 j]
+        counterFlippedBit = counterexample ("Flipped bit: " ++ show j)
+
+      let isUncheckedFile =
+               path == getNamedSnapshotDir namedSnapDir </> FS.mkFsPath ["0.keyops"]
+            || path == getNamedSnapshotDir namedSnapDir </> FS.mkFsPath ["0.blobs"]
+            || path == getNamedSnapshotDir namedSnapDir </> FS.mkFsPath ["0.checksums"]
+
+      -- TODO: remove once write buffer files have checksum verification
+      if isUncheckedFile then
+        pure discard
+      else if n <= 0 then -- file is empty
+        pure $ tabulate "Result" ["No corruption applied"] True
+      else do -- file is non-empty
+
+        -- Flip a bit and try to open the snapshot
+        flipFileBit hfs path j
+        t' <- try @_ @SomeException $ bracket (openSnap s) close $ \_ -> pure ()
+
+        pure $
+          tabCorruptedFile $ counterCorruptedFile $ tabFlippedBit $ counterFlippedBit $
+          case t' of
+            -- If we find an error, we detected corruption. Success!
+            Left e ->
+              tabulate
+                "Result"
+                ["Corruption detected: " <> getConstructorName e]
+                True
+            -- The corruption was not detected. Failure!
+            Right _ -> property False
+  where
+    root = FS.mkFsPath []
+    namedSnapDir = namedSnapshotDir (SessionRoot root) snapName
+
+    conf = defaultTableConfig {
+        confWriteBufferAlloc = AllocNumEntries (NumEntries bufferSize)
+      }
+    es' = fmap (bimap serialiseKey (bimap serialiseValue serialiseBlob)) es
+
+    resolve (SerialisedValue x) (SerialisedValue y) =
+        SerialisedValue (x <> y)
+
+    snapName = fromJust $ mkSnapshotName "snap"
+    snapLabel = SnapshotLabel "label"
+
+    createSnap t =
+        createSnapshot snapName snapLabel SnapFullTable t
+
+    openSnap s =
+        openSnapshot s snapLabel SnapFullTable configNoOverride snapName resolve
+
+    getConstructorName e = takeWhile (/= ' ') (show e)
diff --git a/test/Test/FS.hs b/test/Test/FS.hs
@@ -6,12 +6,19 @@ module Test.FS (tests) where
 import           Control.Concurrent.Class.MonadSTM (MonadSTM (atomically))
 import           Control.Concurrent.Class.MonadSTM.Strict.TMVar
 import           Control.Monad
+import           Control.Monad.Class.MonadThrow
 import           Control.Monad.IOSim (runSimOrThrow)
-import           Data.Char (isAsciiLower, isAsciiUpper)
-import qualified Data.List as List
-import qualified Data.Text as Text
+import           Control.Monad.ST (runST)
+import           Data.Bit (cloneFromByteString, cloneToByteString, flipBit)
+import           Data.ByteString (ByteString)
+import qualified Data.ByteString as BS
+import           Data.Set (Set)
+import qualified Data.Set as Set
+import           Data.Vector.Unboxed (thaw, unsafeFreeze)
 import           GHC.Generics (Generic)
 import           System.FS.API
+import           System.FS.API.Lazy
+import           System.FS.API.Strict
 import           System.FS.Sim.Error
 import qualified System.FS.Sim.MockFS as MockFS
 import qualified System.FS.Sim.Stream as S
@@ -29,6 +36,8 @@ tests = testGroup "Test.FS" [
       -- * Simulated file system properties
       testProperty "prop_numOpenHandles" prop_numOpenHandles
     , testProperty "prop_numDirEntries" prop_numDirEntries
+      -- * Corruption
+    , testProperty "prop_flipFileBit" prop_flipFileBit
       -- * Equality
     , testClassLaws "Stream" $
         eqLaws (Proxy @(Stream Int))
@@ -40,43 +49,16 @@ tests = testGroup "Test.FS" [
   Simulated file system properties
 -------------------------------------------------------------------------------}
 
-newtype Path = Path FsPath
-  deriving stock (Show, Eq)
-
-newtype UniqueList a = UniqueList [a]
-  deriving stock Show
-
-instance (Arbitrary a, Eq a) => Arbitrary (UniqueList a) where
-  arbitrary = do
-      xs <- arbitrary
-      pure (UniqueList (List.nub xs))
-  shrink (UniqueList []) = []
-  shrink (UniqueList xs) = UniqueList . List.nub <$> shrink xs
-
-instance Arbitrary Path where
-  arbitrary = Path . mkFsPath . (:[]) <$> ((:) <$> genChar <*> listOf genChar)
-    where
-      genChar = elements (['A'..'Z'] ++ ['a'..'z'])
-  shrink (Path p) = case fsPathToList p of
-      [] -> []
-      t:_ -> [
-          Path p'
-        | t' <- shrink t
-        , let t'' = Text.filter (\c -> isAsciiUpper c || isAsciiLower c) t'
-        , not (Text.null t'')
-        , let p' = fsPathFromList [t']
-        ]
-
 -- | Sanity check for 'propNoOpenHandles' and 'propNumOpenHandles'
-prop_numOpenHandles :: UniqueList Path -> Property
-prop_numOpenHandles (UniqueList paths) = runSimOrThrow $
+prop_numOpenHandles :: Set FsPathComponent -> Property
+prop_numOpenHandles (Set.toList -> paths) = runSimOrThrow $
     withSimHasFS propTrivial MockFS.empty $ \hfs fsVar -> do
       -- No open handles initially
       fs <- atomically $ readTMVar fsVar
       let prop = propNoOpenHandles fs
 
       -- Open n handles
-      hs <- forM paths $ \(Path p) -> hOpen hfs p (WriteMode MustBeNew)
+      hs <- forM paths $ \(fsPathComponentFsPath -> p) -> hOpen hfs p (WriteMode MustBeNew)
 
       -- Now there should be precisely n open handles
       fs' <- atomically $ readTMVar fsVar
@@ -94,8 +76,12 @@ prop_numOpenHandles (UniqueList paths) = runSimOrThrow $
     n = length paths
 
 -- | Sanity check for 'propNoDirEntries' and 'propNumDirEntries'
-prop_numDirEntries :: Path -> InfiniteList Bool -> UniqueList Path -> Property
-prop_numDirEntries (Path dir) isFiles (UniqueList paths) = runSimOrThrow $
+prop_numDirEntries ::
+     FsPathComponent
+  -> InfiniteList Bool
+  -> Set FsPathComponent
+  -> Property
+prop_numDirEntries (fsPathComponentFsPath -> dir) isFiles (Set.toList -> paths) = runSimOrThrow $
     withSimHasFS propTrivial MockFS.empty $ \hfs fsVar -> do
       createDirectoryIfMissing hfs False dir
 
@@ -104,17 +90,17 @@ prop_numDirEntries (Path dir) isFiles (UniqueList paths) = runSimOrThrow $
       let prop = propNoDirEntries dir fs
 
       -- Create n entries
-      forM_ xs $ \(isFile, Path p) ->
+      forM_ xs $ \(isFile, fsPathComponentFsPath -> p) ->
         if isFile
-          then withFile hfs (dir </> p) (WriteMode MustBeNew) $ \_ -> pure ()
+          then createFile hfs (dir </> p)
           else createDirectory hfs (dir </> p)
 
       -- Now there should be precisely n entries
       fs' <- atomically $ readTMVar fsVar
       let prop' = propNumDirEntries dir n fs'
 
       -- Remove n entries
-      forM_ xs $ \(isFile, Path p) ->
+      forM_ xs $ \(isFile, fsPathComponentFsPath -> p) ->
         if isFile
           then removeFile hfs (dir </> p)
           else removeDirectoryRecursive hfs (dir </> p)
@@ -128,6 +114,50 @@ prop_numDirEntries (Path dir) isFiles (UniqueList paths) = runSimOrThrow $
     n = length paths
     xs = zip (getInfiniteList isFiles) paths
 
+createFile :: MonadThrow m => HasFS m h -> FsPath -> m ()
+createFile hfs p = withFile hfs p (WriteMode MustBeNew) $ \_ -> pure ()
+
+{-------------------------------------------------------------------------------
+  Corruption
+-------------------------------------------------------------------------------}
+
+data WithBitOffset a = WithBitOffset Int a
+  deriving stock Show
+
+instance Arbitrary (WithBitOffset ByteString) where
+  arbitrary = do
+      bs <- arbitrary `suchThat` (\bs -> BS.length bs > 0)
+      bitOffset <- chooseInt (0, BS.length bs - 1)
+      pure $ WithBitOffset bitOffset bs
+  shrink (WithBitOffset bitOffset bs) =
+      [ WithBitOffset bitOffset' bs'
+      | bs' <- shrink bs
+      , BS.length bs' > 0
+      , let bitOffset' = max 0 $ min (BS.length bs' - 1) bitOffset
+      ] ++ [
+        WithBitOffset bitOffset' bs
+      | bitOffset' <- max 0 <$> shrink bitOffset
+      , bitOffset' >= 0
+      ]
+
+prop_flipFileBit :: WithBitOffset ByteString -> Property
+prop_flipFileBit (WithBitOffset bitOffset bs) =
+    ioProperty $
+    withSimHasFS propTrivial MockFS.empty $ \hfs _fsVar -> do
+      void $ withFile hfs path (WriteMode MustBeNew) $ \h -> hPutAllStrict hfs h bs
+      flipFileBit hfs path bitOffset
+      bs' <- withFile hfs path ReadMode $ \h -> BS.toStrict <$> hGetAll hfs h
+      pure (spec_flipFileBit bs bitOffset === bs')
+  where
+    path = mkFsPath ["file"]
+
+spec_flipFileBit :: ByteString -> Int -> ByteString
+spec_flipFileBit bs bitOffset = runST $ do
+    mv <- thaw $ cloneFromByteString bs
+    flipBit mv bitOffset
+    v <- unsafeFreeze mv
+    pure $ cloneToByteString v
+
 {-------------------------------------------------------------------------------
   Equality
 -------------------------------------------------------------------------------}
diff --git a/test/Test/Util/FS.hs b/test/Test/Util/FS.hs
