FEAT

Author: jorisdral

src/Database/LSMTree/Internal/BloomFilter.hs

@@ -28,6 +28,7 @@ import qualified Data.Primitive as P
 import qualified Data.Vector as V
 import qualified Data.Vector.Primitive as VP
 import           Data.Word (Word32, Word64, byteSwap32)
+import           Text.Printf (printf)
 
 import           Control.Exception (assert)
 import           Control.Monad (void, when)
@@ -242,23 +243,25 @@ bloomFilterToLBS bf =
 
 {-# SPECIALISE bloomFilterFromFile ::
      HasFS IO h
+  -> Bloom.Salt
   -> Handle h
   -> IO (Bloom a) #-}
 -- | Read a 'Bloom' from a file.
 --
 bloomFilterFromFile ::
      (PrimMonad m, MonadCatch m)
   => HasFS m h
+  -> Bloom.Salt -- ^ Expected salt
   -> Handle h  -- ^ The open file, in read mode
   -> m (Bloom a)
-bloomFilterFromFile hfs h = do
+bloomFilterFromFile hfs expectedSalt h = do
     header <- rethrowEOFError "Doesn't contain a header" $
               hGetByteArrayExactly hfs h 24
 
     let !version = P.indexByteArray header 0 :: Word32
         !nhashes = P.indexByteArray header 1 :: Word32
         !nbits   = P.indexByteArray header 1 :: Word64
-        !salt    = P.indexByteArray header 2 :: Word64
+        !salt    = P.indexByteArray header 2 :: Bloom.Salt
 
     when (version /= bloomFilterVersion) $ throwFormatError $
       if byteSwap32 version == bloomFilterVersion
@@ -268,8 +271,12 @@ bloomFilterFromFile hfs h = do
     when (nbits <= 0) $ throwFormatError "Length is zero"
 
     -- limit to 2^48 bits
-    when (nbits >= 0x1_0000_0000_0000) $ throwFormatError "Too large bloomfilter"
-    --TODO: get max size from bloomfilter lib
+    when (nbits >= fromIntegral Bloom.maxSizeBits) $ throwFormatError "Too large bloomfilter"
+
+    when (expectedSalt /= salt) $ throwFormatError $
+      printf "Expected salt does not match actual salt: %d /= %d"
+        expectedSalt
+        salt
 
     -- read the filter data from the file directly into the bloom filter
     bloom <-

src/Database/LSMTree/Internal/Run.hs

@@ -36,6 +36,7 @@ import           Control.Monad.Class.MonadSTM (MonadSTM (..))
 import           Control.Monad.Class.MonadThrow
 import           Control.Monad.Primitive
 import           Control.RefCount
+import qualified Data.BloomFilter.Hash as Bloom
 import qualified Data.ByteString.Short as SBS
 import           Data.Foldable (for_)
 import           Database.LSMTree.Internal.BlobFile
@@ -268,6 +269,7 @@ fromWriteBuffer fs hbio sessionSalt params fsPaths buffer blobs = do
   -> HasBlockIO IO h
   -> RunDataCaching
   -> IndexType
+  -> Bloom.Salt
   -> RunFsPaths
   -> IO (Ref (Run IO h)) #-}
 -- | Load a previously written run from disk, checking each file's checksum
@@ -293,10 +295,11 @@ openFromDisk ::
   -> HasBlockIO m h
   -> RunDataCaching
   -> IndexType
+  -> Bloom.Salt -- ^ Expected salt
   -> RunFsPaths
   -> m (Ref (Run m h))
 -- TODO: make exception safe
-openFromDisk fs hbio runRunDataCaching indexType runRunFsPaths = do
+openFromDisk fs hbio runRunDataCaching indexType expectedSalt runRunFsPaths = do
     expectedChecksums <-
        CRC.expectValidFile fs (runChecksumsPath runRunFsPaths) CRC.FormatChecksumsFile
            . fromChecksumsFile
@@ -311,7 +314,7 @@ openFromDisk fs hbio runRunDataCaching indexType runRunFsPaths = do
     let filterPath = forRunFilterRaw paths
     checkCRC CacheRunData (forRunFilterRaw expectedChecksums) filterPath
     runFilter <- FS.withFile fs filterPath FS.ReadMode $
-                   bloomFilterFromFile fs
+                   bloomFilterFromFile fs expectedSalt
 
     (runNumEntries, runIndex) <-
       CRC.expectValidFile fs (forRunIndexRaw paths) CRC.FormatIndexFile

src/Database/LSMTree/Internal/Snapshot.hs

@@ -44,6 +44,7 @@ import           Control.Monad.Class.MonadThrow (MonadMask, bracket,
                      bracketOnError)
 import           Control.Monad.Primitive (PrimMonad)
 import           Control.RefCount
+import qualified Data.BloomFilter.Hash as Bloom
 import           Data.Foldable (sequenceA_, traverse_)
 import           Data.String (IsString)
 import           Data.Text (Text)
@@ -585,9 +586,10 @@ snapshotRun hfs hbio snapUc reg (NamedSnapshotDir targetDir) run = do
   -> ActionRegistry IO
   -> NamedSnapshotDir
   -> ActiveDir
+  -> Bloom.Salt
   -> SnapshotRun
   -> IO (Ref (Run IO h)) #-}
--- | @'openRun' _ _ uniqCounter _ sourceDir targetDir snaprun@ takes all run
+-- | @'openRun' _ _ uniqCounter _ sourceDir targetDir _ snaprun@ takes all run
 -- files that are referenced by @snaprun@, and hard links them from @sourceDir@
 -- into @targetDir@ with new, unique names (using @uniqCounter@). Each set of
 -- (hard linked) files that represents a run is opened and verified, returning
@@ -602,10 +604,12 @@ openRun ::
   -> ActionRegistry m
   -> NamedSnapshotDir
   -> ActiveDir
+  -> Bloom.Salt
   -> SnapshotRun
   -> m (Ref (Run m h))
 openRun hfs hbio uc reg
         (NamedSnapshotDir sourceDir) (ActiveDir targetDir)
+        expectedSalt
         SnapshotRun {
           snapRunNumber  = runNum,
           snapRunCaching = caching,
@@ -617,7 +621,7 @@ openRun hfs hbio uc reg
     hardLinkRunFiles hfs hbio reg sourcePaths targetPaths
 
     withRollback reg
-      (Run.openFromDisk hfs hbio caching indexType targetPaths)
+      (Run.openFromDisk hfs hbio caching indexType expectedSalt targetPaths)
       releaseRef
 
 {-------------------------------------------------------------------------------

src/Database/LSMTree/Internal/Unsafe.hs

@@ -1536,11 +1536,11 @@ openTableFromSnapshot policyOveride sesh snap label resolve =
           openWriteBuffer reg resolve hfs hbio uc activeDir snapWriteBufferPaths
 
         -- Hard link runs into the active directory,
-        snapLevels' <- traverse (openRun hfs hbio uc reg snapDir activeDir) snapLevels
+        snapLevels' <- traverse (openRun hfs hbio uc reg snapDir activeDir (Paths.getSessionSalt salt)) snapLevels
         unionLevel <- case mTreeOpt of
               Nothing -> pure NoUnion
               Just mTree -> do
-                snapTree <- traverse (openRun hfs hbio uc reg snapDir activeDir) mTree
+                snapTree <- traverse (openRun hfs hbio uc reg snapDir activeDir (Paths.getSessionSalt salt)) mTree
                 mt <- fromSnapMergingTree hfs hbio salt uc resolve activeDir reg snapTree
                 isStructurallyEmpty mt >>= \case
                   True ->

test/Test/Database/LSMTree/Internal/BloomFilter.hs

@@ -1,14 +1,15 @@
 module Test.Database.LSMTree.Internal.BloomFilter (tests) where
 
 import           Control.DeepSeq (deepseq)
-import           Control.Exception (displayException)
+import           Control.Exception (Exception (..), displayException)
 import           Control.Monad (void)
 import qualified Control.Monad.IOSim as IOSim
 import           Data.Bits ((.&.))
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.Builder as BS.Builder
 import qualified Data.ByteString.Builder.Extra as BS.Builder
 import qualified Data.ByteString.Lazy as LBS
+import qualified Data.List as List
 import qualified Data.Set as Set
 import qualified Data.Vector as V
 import qualified Data.Vector.Primitive as VP
@@ -26,6 +27,8 @@ import           Test.Tasty.QuickCheck hiding ((.&.))
 import           Data.BloomFilter.Blocked (Salt)
 import qualified Data.BloomFilter.Blocked as Bloom
 import           Database.LSMTree.Internal.BloomFilter
+import           Database.LSMTree.Internal.CRC32C (FileCorruptedError (..),
+                     FileFormat (..))
 import           Database.LSMTree.Internal.Paths (SessionSalt (..))
 import           Database.LSMTree.Internal.Serialise (SerialisedKey,
                      serialiseKey)
@@ -69,9 +72,18 @@ limitBits b = b .&. 0xffffff
 prop_total_deserialisation :: BS.ByteString -> Property
 prop_total_deserialisation bs =
     case bloomFilterFromBS bs of
-      Left err -> label (displayException err) $ property True
+      Left err ->
+        label (mkLabel err) $ property True
       Right bf -> label "parsed successfully" $ property $
         bf `deepseq` True
+  where
+    mkLabel err = case err of
+        IOSim.FailureException e
+          | Just (ErrFileFormatInvalid fsep FormatBloomFilterFile msg) <- fromException e
+          , let msg' = "Expected salt does not match actual salt"
+          , msg' `List.isPrefixOf` msg
+          -> displayException $ ErrFileFormatInvalid fsep FormatBloomFilterFile msg'
+        _ -> displayException err
 
 -- | Write the bytestring to a file in the mock file system and then use
 -- 'bloomFilterFromFile'.
@@ -85,7 +97,7 @@ bloomFilterFromBS bs =
         void $ FS.hPutAllStrict hfs h bs
       -- deserialise from file
       FS.withFile hfs file FS.ReadMode $ \h ->
-        bloomFilterFromFile hfs h
+        bloomFilterFromFile hfs testSalt h
 
 -- Length is in bits. A large length would require significant amount of
 -- memory, so we make it 'Small'.

test/Test/Database/LSMTree/Internal/Run.hs

@@ -219,7 +219,8 @@ prop_WriteAndOpen fs hbio wb =
           paths' = paths { runNumber = RunNumber 17}
       hardLinkRunFiles fs hbio reg paths paths'
       loaded <- openFromDisk fs hbio (runParamCaching runParams)
-                             (runParamIndex runParams) (simplePath 17)
+                             (runParamIndex runParams) (getSessionSalt testSessionSalt)
+                             (simplePath 17)
 
       Run.size written @=? Run.size loaded
       withRef written $ \written' ->