511 lines
19 KiB
Haskell
511 lines
19 KiB
Haskell
module Main (main) where
|
||
|
||
import Control.Concurrent.Async (concurrently, mapConcurrently)
|
||
import Control.Exception (SomeException, bracket_, try)
|
||
import Control.Monad (unless, void)
|
||
import Crypto.Hash (Digest, SHA256, hashlazy)
|
||
import Data.Bits ((.&.))
|
||
import qualified Data.ByteArray as BA
|
||
import qualified Data.ByteString as BS
|
||
import qualified Data.ByteString.Base16 as B16
|
||
import qualified Data.ByteString.Lazy as BL
|
||
import Data.Int (Int64)
|
||
import Data.List (sort)
|
||
import qualified Data.Pool as Pool
|
||
import Data.Text (Text)
|
||
import qualified Data.Text.Encoding as TE
|
||
import Data.Word (Word64)
|
||
import Database.PostgreSQL.Simple (Connection, Only (..))
|
||
import qualified Database.PostgreSQL.Simple as PG
|
||
import Nuthatch.Config
|
||
( Config (..)
|
||
, ProductConfig (..)
|
||
, S3Config (..)
|
||
, loadConfig
|
||
, revealSecret
|
||
, secret
|
||
)
|
||
import Nuthatch.Db
|
||
( Db
|
||
, configurePoolLimit
|
||
, ensurePersonalPool
|
||
, migrate
|
||
, migrationAdvisoryLockKey
|
||
, withDb
|
||
, withTransaction
|
||
)
|
||
import qualified Nuthatch.Db as Db
|
||
import Nuthatch.Protocol
|
||
import Nuthatch.S3
|
||
import Nuthatch.Types
|
||
|
||
main :: IO ()
|
||
main = do
|
||
config <- loadConfig
|
||
withDb (cfgPostgres config) $ \database -> do
|
||
migrate database "sql" (cfgProduct config)
|
||
migrate database "sql" (cfgProduct config)
|
||
verifyMigrationLock database
|
||
verifySchema database
|
||
verifyAdmission database config
|
||
verifyRowDecoder database
|
||
store <- newS3 (cfgS3 config)
|
||
verifyStorage config store
|
||
putStrLn "Storage foundation integration gate passed"
|
||
|
||
verifyMigrationLock :: Db -> IO ()
|
||
verifyMigrationLock database = do
|
||
assertEqual
|
||
"migration lock key follows SHA-256 namespace derivation"
|
||
migrationLockKeyFromHash
|
||
migrationAdvisoryLockKey
|
||
Pool.withResource database $ \holder ->
|
||
Pool.withResource database $ \contender -> do
|
||
bracket_
|
||
(acquireAdvisoryLock holder)
|
||
(void (releaseAdvisoryLock holder))
|
||
(do
|
||
admitted <- tryAdvisoryLock contender
|
||
assert "migration lock excludes another session" (not admitted))
|
||
admitted <- tryAdvisoryLock contender
|
||
assert "migration lock admits another session after release" admitted
|
||
released <- releaseAdvisoryLock contender
|
||
assert "contending session releases migration lock" released
|
||
|
||
migrationLockKeyFromHash :: Int64
|
||
migrationLockKeyFromHash =
|
||
fromIntegral
|
||
(BS.foldl' appendByte 0 (BS.take 8 digestBytes) .&. maxSignedKey)
|
||
where
|
||
digestBytes = BA.convert (hashlazy "nuthatch" :: Digest SHA256)
|
||
appendByte accumulator byte =
|
||
accumulator * 256 + fromIntegral byte
|
||
maxSignedKey = 0x7fffffffffffffff :: Word64
|
||
|
||
acquireAdvisoryLock :: Connection -> IO ()
|
||
acquireAdvisoryLock connection =
|
||
void
|
||
(PG.query connection
|
||
"SELECT pg_advisory_lock(?) IS NULL"
|
||
(Only migrationAdvisoryLockKey)
|
||
:: IO [Only Bool])
|
||
|
||
tryAdvisoryLock :: Connection -> IO Bool
|
||
tryAdvisoryLock connection = do
|
||
[Only admitted] <- PG.query connection
|
||
"SELECT pg_try_advisory_lock(?)"
|
||
(Only migrationAdvisoryLockKey)
|
||
pure admitted
|
||
|
||
releaseAdvisoryLock :: Connection -> IO Bool
|
||
releaseAdvisoryLock connection = do
|
||
[Only released] <- PG.query connection
|
||
"SELECT pg_advisory_unlock(?)"
|
||
(Only migrationAdvisoryLockKey)
|
||
pure released
|
||
|
||
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_initial_schema.sql"]
|
||
(map fromOnly migrations)
|
||
|
||
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
|
||
productCfg = cfgProduct config
|
||
ensurePersonalPool database
|
||
productCfg { cfgPersonalQuotaBytes = 30 }
|
||
personalUser
|
||
ensurePersonalPool database
|
||
productCfg { cfgPersonalQuotaBytes = 40 }
|
||
personalUser
|
||
_ <- Db.execute database
|
||
"UPDATE pools SET reserved_bytes = 35 WHERE pool_id = ?"
|
||
(Only (unPoolId personal))
|
||
personalLower <- try
|
||
(ensurePersonalPool database
|
||
productCfg { 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 (DispatchTiming 600 90)
|
||
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
|
||
|
||
verifyIdempotencyExpiry database
|
||
|
||
verifyIdempotencyExpiry :: Db -> IO ()
|
||
verifyIdempotencyExpiry database = do
|
||
resetAuthority database 100
|
||
first <- admitWithWindow database "expiring-operation" "expiring-key" 10 1
|
||
let rec = preparedRecord [first]
|
||
finalized <- withTransaction database $ \connection -> do
|
||
applying <- markApplying connection rec
|
||
pending <- markExternalPending connection rec (DispatchTiming 600 90)
|
||
settled <- markExternalSettled connection rec
|
||
succeeded <- finalizeSuccess connection rec 10 (Just fixedIdentity)
|
||
pure (applying && pending && settled && succeeded)
|
||
assert "terminal expiry fixture finalized" finalized
|
||
_ <- Db.execute database
|
||
"UPDATE mutations SET updated_at = now() - interval '5 seconds' \
|
||
\WHERE user_id = ? AND kind = ? AND operation_id = ?"
|
||
( "integration-user" :: Text
|
||
, "put" :: Text
|
||
, "expiring-operation" :: Text
|
||
)
|
||
outcomes <- mapConcurrently
|
||
(\size ->
|
||
admitWithWindow database "expiring-operation" "expiring-key" size 1)
|
||
[11, 12]
|
||
assertEqual "one first request after terminal expiry" 1 (countPrepared outcomes)
|
||
assertEqual "atomic conflicting reuse after terminal expiry" 1
|
||
(countConflict outcomes)
|
||
assertAuthority database 1 (mrGrowth (preparedRecord outcomes))
|
||
|
||
resetAuthority database 100
|
||
active <- admitWithWindow database "active-operation" "active-key" 10 1
|
||
_ <- Db.execute database
|
||
"UPDATE mutations SET updated_at = now() - interval '5 seconds' \
|
||
\WHERE user_id = ? AND kind = ? AND operation_id = ?"
|
||
( "integration-user" :: Text
|
||
, "put" :: Text
|
||
, "active-operation" :: Text
|
||
)
|
||
replay <- admitWithWindow database "active-operation" "active-key" 10 1
|
||
blocked <- admitWithWindow database "other-operation" "active-key" 10 1
|
||
assertEqual "aged active identity remains replayable" 1
|
||
(countReplay [active, replay])
|
||
assertEqual "aged active key remains locked" 1 (countRejected [blocked])
|
||
assertAuthority database 1 10
|
||
|
||
verifyRowDecoder :: Db -> IO ()
|
||
verifyRowDecoder database = do
|
||
resetAuthority database 100
|
||
outcome <- admit database "invalid-phase-operation" "invalid-phase-key" 10
|
||
let record = preparedRecord [outcome]
|
||
roundTripped <- withTransaction database $ \connection ->
|
||
lookupMutation
|
||
connection
|
||
(mrUserId record)
|
||
(mrKind record)
|
||
(mrOperationId record)
|
||
assertEqual "typed mutation row round-trip" (Just record) roundTripped
|
||
_ <- Db.execute database
|
||
"UPDATE mutations SET external_state = 'pending', \
|
||
\dispatch_started_at = now(), request_deadline_at = now(), \
|
||
\worker_lease_until = now() \
|
||
\WHERE user_id = ? AND kind = ? AND operation_id = ?"
|
||
( unUserId (mrUserId record)
|
||
, kindLabel (mrKind record)
|
||
, unOperationId (mrOperationId record)
|
||
)
|
||
invalidDecode <-
|
||
( try $ withTransaction database $ \connection ->
|
||
lookupMutation
|
||
connection
|
||
(mrUserId record)
|
||
(mrKind record)
|
||
(mrOperationId record)
|
||
) :: IO (Either SomeException (Maybe MutationRecord))
|
||
case invalidDecode of
|
||
Left _ -> pure ()
|
||
Right value ->
|
||
fail ("invalid persisted mutation phase decoded as " <> show value)
|
||
|
||
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 <- expectRight "observe object without identity metadata"
|
||
=<< observeObjectFully 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
|
||
let s3Cfg = cfgS3 config
|
||
deniedStore <- newS3 s3Cfg
|
||
{ cfgS3SecretKey =
|
||
secret (revealSecret (cfgS3SecretKey s3Cfg) <> "-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 =
|
||
admitWithWindow database operation key size 3600
|
||
|
||
admitWithWindow :: Db -> Text -> Text -> Integer -> Integer -> IO AdmitOutcome
|
||
admitWithWindow database operation key size window =
|
||
withTransaction database $ \connection ->
|
||
admitMutation
|
||
connection
|
||
AdmissionRequest
|
||
{ admissionIdentity =
|
||
AdmissionIdentity (UserId "integration-user") Put
|
||
(OperationId operation)
|
||
, admissionPoolId = sharedPool
|
||
, admissionRelativePath = RelativePath key
|
||
, admissionFingerprint =
|
||
MutationFingerprint
|
||
Put
|
||
(ObjectKey ("shared/" <> key)) size (Just fixedIdentity)
|
||
}
|
||
AdmissionLimits
|
||
{ admissionMaxObjectBytes = 1000
|
||
, admissionIdempotencyWindowSec = window
|
||
}
|
||
|
||
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
|