Harden mutation admission and storage observation

test/Integration.hs

@@ -0,0 +1,416 @@
+module Main (main) where
+
+import           Control.Concurrent.Async (concurrently, mapConcurrently)
+import           Control.Exception (SomeException, try)
+import           Control.Monad (unless)
+import           Crypto.Hash (Digest, SHA256, hashlazy)
+import qualified Data.ByteArray as BA
+import qualified Data.ByteString.Base16 as B16
+import qualified Data.ByteString.Lazy as BL
+import           Data.List (sort)
+import           Data.Text (Text)
+import qualified Data.Text.Encoding as TE
+import           Database.PostgreSQL.Simple (Only (..))
+import           Nuthatch.Config (Config (..), loadConfig)
+import           Nuthatch.Db
+  ( Db
+  , configurePoolLimit
+  , ensurePersonalPool
+  , migrate
+  , withDb
+  , withTransaction
+  )
+import qualified Nuthatch.Db as Db
+import           Nuthatch.Protocol
+import           Nuthatch.S3
+import           Nuthatch.Types
+
+main :: IO ()
+main = do
+  config <- loadConfig
+  withDb config $ \database -> do
+    migrate database "sql" config
+    migrate database "sql" config
+    emptyCatalog <- schemaCatalog database
+    resetSchema database
+    installLegacySchema database
+    migrate database "sql" config
+    migrate database "sql" config
+    adoptedCatalog <- schemaCatalog database
+    assertEqual "legacy and empty migration schema equivalence"
+      emptyCatalog adoptedCatalog
+    verifySchema database
+    verifyAdmission database config
+  store <- newS3 config
+  verifyStorage config store
+  putStrLn "WP-04 integration gate passed"
+
+verifySchema :: Db -> IO ()
+verifySchema database = do
+  tables <- Db.query_ database
+    "SELECT table_name FROM information_schema.tables \
+    \WHERE table_schema = 'public' \
+    \AND table_name IN ('pools','mutations','key_locks','objects_meta','schema_migrations') \
+    \ORDER BY table_name"
+    :: IO [Only Text]
+  assertEqual "authoritative schema tables"
+    ["key_locks", "mutations", "objects_meta", "pools", "schema_migrations"]
+    (map fromOnly tables)
+  constraints <- Db.query_ database
+    "SELECT count(*) FROM pg_constraint \
+    \WHERE conname = 'pools_live_bytes_within_limit' \
+    \AND conrelid = 'pools'::regclass"
+    :: IO [Only Int]
+  assertEqual "pool limit constraint" [Only 1] constraints
+  migrations <- Db.query_ database
+    "SELECT filename FROM schema_migrations ORDER BY filename"
+    :: IO [Only Text]
+  assertEqual "migration history"
+    ["001_schema.sql", "002_pool_limit_guard.sql"]
+    (map fromOnly migrations)
+
+schemaCatalog :: Db -> IO [Text]
+schemaCatalog database = do
+  rows <- Db.query_ database
+    "SELECT description FROM ( \
+    \SELECT 'column|' || table_name || '|' || ordinal_position::text || '|' \
+    \  || column_name || '|' || data_type || '|' || is_nullable || '|' \
+    \  || COALESCE(column_default, '') AS description \
+    \FROM information_schema.columns \
+    \WHERE table_schema = 'public' \
+    \UNION ALL \
+    \SELECT 'constraint|' || c.relname || '|' || con.conname || '|' \
+    \  || con.contype::text || '|' || pg_get_constraintdef(con.oid) \
+    \FROM pg_constraint con \
+    \JOIN pg_class c ON c.oid = con.conrelid \
+    \JOIN pg_namespace n ON n.oid = c.relnamespace \
+    \WHERE n.nspname = 'public' \
+    \UNION ALL \
+    \SELECT 'index|' || tablename || '|' || indexname || '|' || indexdef \
+    \FROM pg_indexes WHERE schemaname = 'public' \
+    \) catalog ORDER BY description"
+    :: IO [Only Text]
+  pure (map fromOnly rows)
+
+resetSchema :: Db -> IO ()
+resetSchema database = do
+  _ <- Db.execute_ database
+    "DROP TABLE IF EXISTS schema_migrations, key_locks, objects_meta, \
+    \mutations, pools CASCADE"
+  pure ()
+
+installLegacySchema :: Db -> IO ()
+installLegacySchema database = do
+  mapM_ (Db.execute_ database)
+    [ "CREATE TABLE pools ( \
+      \pool_id TEXT PRIMARY KEY, \
+      \limit_bytes BIGINT NOT NULL CHECK (limit_bytes > 0), \
+      \used_bytes BIGINT NOT NULL DEFAULT 0 CHECK (used_bytes >= 0), \
+      \reserved_bytes BIGINT NOT NULL DEFAULT 0 CHECK (reserved_bytes >= 0))"
+    , "CREATE TABLE mutations ( \
+      \user_id TEXT NOT NULL, kind TEXT NOT NULL, operation_id TEXT NOT NULL, \
+      \pool_id TEXT NOT NULL REFERENCES pools (pool_id), object_key TEXT NOT NULL, \
+      \relative_path TEXT NOT NULL, declared_size BIGINT NOT NULL, \
+      \content_sha256 TEXT, prior_size BIGINT NOT NULL, prior_sha256 TEXT, \
+      \growth_bytes BIGINT NOT NULL, state TEXT NOT NULL, \
+      \external_state TEXT NOT NULL DEFAULT 'idle', \
+      \created_at TIMESTAMPTZ NOT NULL DEFAULT now(), \
+      \updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), \
+      \PRIMARY KEY (user_id, kind, operation_id))"
+    , "CREATE TABLE key_locks ( \
+      \object_key TEXT PRIMARY KEY, user_id TEXT NOT NULL, kind TEXT NOT NULL, \
+      \operation_id TEXT NOT NULL)"
+    , "CREATE TABLE objects_meta ( \
+      \object_key TEXT PRIMARY KEY, \
+      \pool_id TEXT NOT NULL REFERENCES pools (pool_id), \
+      \size_bytes BIGINT NOT NULL, content_sha256 TEXT NOT NULL, \
+      \updated_at TIMESTAMPTZ NOT NULL DEFAULT now())"
+    ]
+  pure ()
+
+verifyAdmission :: Db -> Config -> IO ()
+verifyAdmission database config = do
+  resetAuthority database 100
+  replayOutcomes <- mapConcurrently
+    (const (admit database "same-operation" "same-key" 10))
+    [1 .. 8 :: Int]
+  assertEqual "one first request" 1 (countPrepared replayOutcomes)
+  assertEqual "stable concurrent replay" 7 (countReplay replayOutcomes)
+  assertAuthority database 1 10
+
+  resetAuthority database 100
+  conflictOutcomes <- mapConcurrently
+    (\size -> admit database "conflicting-operation" "conflicting-key" size)
+    [10, 11]
+  assertEqual "one conflicting first request" 1 (countPrepared conflictOutcomes)
+  assertEqual "conflicting reuse" 1 (countConflict conflictOutcomes)
+  assertAuthority database 1 (mrGrowth (preparedRecord conflictOutcomes))
+
+  resetAuthority database 100
+  keyOutcomes <- mapConcurrently
+    (\operation -> admit database operation "locked-key" 10)
+    ["key-operation-a", "key-operation-b"]
+  assertEqual "one same-key authority" 1 (countPrepared keyOutcomes)
+  assertEqual "one same-key rejection" 1 (countRejected keyOutcomes)
+  assertAuthority database 1 10
+
+  resetAuthority database 10
+  quotaOutcomes <- mapConcurrently
+    (\(operation, key) -> admit database operation key 8)
+    [("quota-operation-a", "quota-key-a"), ("quota-operation-b", "quota-key-b")]
+  assertEqual "one quota reservation" 1 (countPrepared quotaOutcomes)
+  assertEqual "one quota rejection" 1 (countRejected quotaOutcomes)
+  assertAuthority database 1 8
+  lower <- try (configurePoolLimit database sharedPool 7)
+    :: IO (Either SomeException ())
+  assert "unsafe pool limit reduction rejected" (isLeft lower)
+  configurePoolLimit database sharedPool 20
+  pool <- Db.query database
+    "SELECT limit_bytes, used_bytes, reserved_bytes FROM pools WHERE pool_id = ?"
+    (Only (unPoolId sharedPool))
+    :: IO [(Integer, Integer, Integer)]
+  assertEqual "safe pool limit reconfiguration" [(20, 0, 8)] pool
+
+  let personalUser = UserId "personal-limit-user"
+      personal = personalPool personalUser
+  ensurePersonalPool database
+    config { cfgPersonalQuotaBytes = 30 }
+    personalUser
+  ensurePersonalPool database
+    config { cfgPersonalQuotaBytes = 40 }
+    personalUser
+  _ <- Db.execute database
+    "UPDATE pools SET reserved_bytes = 35 WHERE pool_id = ?"
+    (Only (unPoolId personal))
+  personalLower <- try
+    (ensurePersonalPool database
+      config { cfgPersonalQuotaBytes = 34 }
+      personalUser)
+    :: IO (Either SomeException ())
+  assert "unsafe personal limit reduction rejected" (isLeft personalLower)
+  personalPoolRow <- Db.query database
+    "SELECT limit_bytes, used_bytes, reserved_bytes FROM pools WHERE pool_id = ?"
+    (Only (unPoolId personal))
+    :: IO [(Integer, Integer, Integer)]
+  assertEqual "personal limit reconfiguration" [(40, 0, 35)] personalPoolRow
+
+  resetAuthority database 100
+  (racingAdmission, racingLimit) <- concurrently
+    (admit database "limit-race-operation" "limit-race-key" 80)
+    (try (configurePoolLimit database sharedPool 50)
+      :: IO (Either SomeException ()))
+  racingPool <- Db.query database
+    "SELECT limit_bytes, used_bytes, reserved_bytes FROM pools WHERE pool_id = ?"
+    (Only (unPoolId sharedPool))
+    :: IO [(Integer, Integer, Integer)]
+  case (racingAdmission, racingLimit, racingPool) of
+    (AdmitPrepared _, Left _, [(100, 0, 80)]) ->
+      assertAuthority database 1 80
+    (AdmitRejected _ _, Right (), [(50, 0, 0)]) ->
+      assertAuthority database 0 0
+    other -> fail ("unsafe admission/limit race result: " <> show other)
+
+  resetAuthority database 100
+  observationOutcome <- admit database
+    "observation-operation"
+    "observation-key"
+    10
+  let observationRecord = preparedRecord [observationOutcome]
+      thirdState =
+        ObjectPresent 10
+          (Just (ContentSha256
+            "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"))
+  settled <- withTransaction database $ \connection -> do
+    applying <- markApplying connection observationRecord
+    pending <- markExternalPending connection observationRecord
+    externalSettled <- markExternalSettled connection observationRecord
+    pure (applying && pending && externalSettled)
+  assert "observation fixture reached applying/settled" settled
+  unavailable <- withTransaction database $ \connection ->
+    finalizeFromObservation connection observationRecord
+      (ObjectUnavailable "timeout")
+  ambiguous <- withTransaction database $ \connection ->
+    finalizeFromObservation connection observationRecord thirdState
+  assertEqual "unavailable observation remains non-terminal" Nothing unavailable
+  assertEqual "third-state observation remains non-terminal" Nothing ambiguous
+  assertAuthority database 1 10
+  durableState <- Db.query database
+    "SELECT state, external_state FROM mutations \
+    \WHERE user_id = ? AND kind = ? AND operation_id = ?"
+    ( "integration-user" :: Text
+    , "put" :: Text
+    , "observation-operation" :: Text
+    )
+    :: IO [(Text, Text)]
+  assertEqual "ambiguous observation retains mutation authority"
+    [("applying", "settled")]
+    durableState
+
+verifyStorage :: Config -> S3 -> IO ()
+verifyStorage config store = do
+  let fallbackKey = ObjectKey "shared/fallback.bin"
+      fallbackBody = BL.replicate (5 * 1024 * 1024) 0
+      fallbackIdentity =
+        ContentSha256
+          "c036cbb7553a909f8b8877d4461924307f27ecb66cff928eeeafd569c3887e29"
+  expectRight "put object without identity metadata"
+    =<< putObject store fallbackKey
+      (toInteger (BL.length fallbackBody))
+      Nothing
+      fallbackBody
+  fallbackHead <- expectJust "observe object without identity metadata"
+    =<< headObject store fallbackKey
+  assertEqual "streaming digest fallback size"
+    (toInteger (BL.length fallbackBody))
+    (oiSize fallbackHead)
+  assertEqual "streaming digest fallback identity"
+    (Just fallbackIdentity)
+    (oiSha256 fallbackHead)
+  expectRight "delete fallback object" =<< deleteObject store fallbackKey
+
+  let unicodeKey = ObjectKey "shared/문서/βeta.txt"
+      firstBody = "same-size-a"
+      secondBody = "same-size-b"
+      firstIdentity = contentIdentity firstBody
+      secondIdentity = contentIdentity secondBody
+  expectRight "put Unicode key"
+    =<< putObject store unicodeKey (toInteger (BL.length firstBody))
+      (Just firstIdentity) firstBody
+  fetched <- expectRight "get Unicode key" =<< getObjectBytes store unicodeKey
+  assertEqual "Unicode key round trip" firstBody fetched
+  firstHead <- expectJust "head Unicode key" =<< headObject store unicodeKey
+  assertEqual "head size" (toInteger (BL.length firstBody)) (oiSize firstHead)
+  assertEqual "head metadata" (Just firstIdentity) (oiSha256 firstHead)
+  expectRight "overwrite same-size content"
+    =<< putObject store unicodeKey (toInteger (BL.length secondBody))
+      (Just secondIdentity) secondBody
+  secondHead <- expectJust "head overwritten key" =<< headObject store unicodeKey
+  assert "same-size content identity differs"
+    (oiSha256 firstHead /= oiSha256 secondHead)
+
+  let paged =
+        [ (ObjectKey "shared/page/α", "a")
+        , (ObjectKey "shared/page/β", "b")
+        , (ObjectKey "shared/page/γ", "c")
+        ]
+  mapM_
+    (\(key, body) ->
+      expectRight "put paged key"
+        =<< putObject store key (toInteger (BL.length body))
+          (Just (contentIdentity body)) body)
+    paged
+  keys <- collectPages store "shared/page/" Nothing []
+  assertEqual "paginated Unicode list"
+    (sort [unObjectKey key | (key, _) <- paged])
+    (sort keys)
+
+  missing <- getObjectBytes store (ObjectKey "shared/missing")
+  assertEqual "typed not-found error" (Left S3NotFound) missing
+  deniedStore <- newS3 config
+    { cfgS3SecretKey = cfgS3SecretKey config <> "-invalid" }
+  denied <- listObjects deniedStore "" Nothing Nothing 1
+  case denied of
+    Left (S3ServiceError code) ->
+      assert "typed provider rejection" (code >= 400 && code < 500)
+    other -> fail ("expected typed provider rejection, got " <> show other)
+  expectRight "delete Unicode key" =<< deleteObject store unicodeKey
+  absent <- headObject store unicodeKey
+  assertEqual "deleted key absent" (Right Nothing) absent
+
+collectPages :: S3 -> Text -> Maybe Text -> [Text] -> IO [Text]
+collectPages store prefix marker accumulated = do
+  page <- expectRight "list page" =<< listObjects store prefix Nothing marker 1
+  let accumulated' = accumulated <> map fst (lpObjects page)
+  if lpIsTruncated page
+    then case lpContinuationToken page of
+      Nothing -> fail "truncated list omitted continuation marker"
+      Just next -> collectPages store prefix (Just next) accumulated'
+    else pure accumulated'
+
+admit :: Db -> Text -> Text -> Integer -> IO AdmitOutcome
+admit database operation key size =
+  withTransaction database $ \connection ->
+    admitMutation
+      connection
+      (UserId "integration-user")
+      Put
+      (OperationId operation)
+      sharedPool
+      (ObjectKey ("shared/" <> key))
+      (RelativePath key)
+      size
+      (Just fixedIdentity)
+      1000
+
+resetAuthority :: Db -> Integer -> IO ()
+resetAuthority database limitBytes = do
+  _ <- Db.execute_ database
+    "TRUNCATE key_locks, mutations, objects_meta, pools CASCADE"
+  configurePoolLimit database sharedPool limitBytes
+
+assertAuthority :: Db -> Int -> Integer -> IO ()
+assertAuthority database expectedLocks expectedReserved = do
+  locks <- Db.query_ database "SELECT count(*) FROM key_locks"
+    :: IO [Only Int]
+  reservations <- Db.query database
+    "SELECT reserved_bytes FROM pools WHERE pool_id = ?"
+    (Only (unPoolId sharedPool))
+    :: IO [Only Integer]
+  assertEqual "active key locks" [Only expectedLocks] locks
+  assertEqual "reserved bytes" [Only expectedReserved] reservations
+
+countPrepared :: [AdmitOutcome] -> Int
+countPrepared outcomes = length [() | AdmitPrepared _ <- outcomes]
+
+countReplay :: [AdmitOutcome] -> Int
+countReplay outcomes = length [() | AdmitReplay MutationPending _ <- outcomes]
+
+countConflict :: [AdmitOutcome] -> Int
+countConflict outcomes = length [() | AdmitConflict <- outcomes]
+
+countRejected :: [AdmitOutcome] -> Int
+countRejected outcomes = length [() | AdmitRejected _ _ <- outcomes]
+
+preparedRecord :: [AdmitOutcome] -> MutationRecord
+preparedRecord outcomes =
+  case [record | AdmitPrepared record <- outcomes] of
+    [record] -> record
+    _ -> error "expected exactly one prepared record"
+
+fixedIdentity :: ContentSha256
+fixedIdentity =
+  ContentSha256 "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+
+contentIdentity :: BL.ByteString -> ContentSha256
+contentIdentity =
+  ContentSha256
+    . TE.decodeUtf8
+    . B16.encode
+    . BA.convert
+    . (hashlazy :: BL.ByteString -> Digest SHA256)
+
+expectRight :: String -> Either S3Error value -> IO value
+expectRight label result =
+  case result of
+    Left err -> fail (label <> ": " <> show err)
+    Right value -> pure value
+
+expectJust :: String -> Either S3Error (Maybe value) -> IO value
+expectJust label result = do
+  value <- expectRight label result
+  case value of
+    Nothing -> fail (label <> ": object absent")
+    Just present -> pure present
+
+assert :: String -> Bool -> IO ()
+assert label condition =
+  unless condition (fail label)
+
+assertEqual :: (Eq value, Show value) => String -> value -> value -> IO ()
+assertEqual label expected actual =
+  unless (expected == actual) $
+    fail (label <> ": expected " <> show expected <> ", got " <> show actual)
+
+isLeft :: Either left right -> Bool
+isLeft (Left _) = True
+isLeft (Right _) = False