Keys.hs

-- | Loading signing-key material, and refreshing it on a live server.
--
-- 'loadKeyMaterial' is __the__ stored→live load path: it reads every publishable key
-- (active + retired) and assembles the three things the runtime needs — the private key
-- that signs new tokens, the public key set the verifier trusts, and the precomputed JWKS
-- document served at @\/.well-known\/jwks.json@. 'bootstrapKeys' wraps it with first-boot
-- key generation; 'reloadKeys' wraps it with the swap-or-keep-last-good policy the
-- periodic refresh and the @SIGHUP@ handler both use.
--
-- Publishing retired keys alongside the active one is what makes rotation zero-downtime:
-- a token signed moments before @shomei-admin keys activate@ keeps verifying until it
-- expires, here and at every downstream that fetches the JWKS.
--
-- Private key material is encrypted at rest (see "Shomei.SigningKey.Protection.Jwt"). Only the
-- __signer__ is decrypted; the verifier key set and the published JWKS are built from the
-- public column, so a missing or wrong key-encryption key can never break verification of
-- outstanding tokens; runtime startup still requires the KEK so minting is always available.
--
-- Runs over a minimal stack — just 'SigningKeyStore' + 'Clock' + 'Database' +
-- @Error AuthError@ + 'IOE' against the pool — because it must run /before/ the full
-- 'Shomei.Server.App.Env' exists (there is no key yet to build it).
module Shomei.Server.Keys
  ( LoadedKeys (..),
    loadKeyMaterial,
    bootstrapKeys,
    reloadKeys,
    publishedKids,
    loadKekFromEnv,
    loadNamedKekFromEnv,
  )
where

import Crypto.JOSE.JWK (JWK, JWKSet)
import Data.Aeson (Value (Object))
import Data.Aeson qualified as Aeson
import Data.Aeson.KeyMap qualified as KM
import Data.Bifunctor (first)
import Data.Foldable (toList)
import Data.IORef (IORef, readIORef, writeIORef)
import Data.Text qualified as Text
import Effectful (Eff, IOE, runEff, (:>))
import Effectful.Error.Static (runErrorNoCallStack)
import Shomei.Error (AuthError)
import Shomei.Persistence.Database.Postgres (runDatabasePool)
import Shomei.Persistence.Pool.Postgres (Pool)
import Shomei.Prelude
import Shomei.SigningKey.Domain (SigningAlgorithm, SigningKeyStatus (KeyActive), StoredSigningKey (..))
import Shomei.SigningKey.Jwks.Jwt (KeySet (..), jwksDocument, keySetPublicJwks)
import Shomei.SigningKey.Key.Jwt (generateSigningKeyFor, keyKid, toStoredSigningKeyFor)
import Shomei.SigningKey.Postgres (runSigningKeyStorePostgres)
import Shomei.SigningKey.Protection.Jwt
  ( KeyDecryptError (..),
    KeyEncryptionKey,
    decryptStoredSigningKey,
    keyEncryptionKeyFromBase64,
    protectStoredSigningKey,
    publicJwkFromStored,
  )
import Shomei.SigningKey.Store
  ( SigningKeyStore,
    insertSigningKey,
    listActiveSigningKeys,
    listPublishableSigningKeys,
  )
import Shomei.Time.Postgres (runClockIO)
import Shomei.Time.Store (Clock, now)
import System.Environment (lookupEnv)
import System.IO (hPutStrLn, stderr)

-- | Read @SHOMEI_KEY_ENCRYPTION_KEY@ (32 bytes, base64).
loadKekFromEnv :: IO KeyEncryptionKey
loadKekFromEnv = loadNamedKekFromEnv "SHOMEI_KEY_ENCRYPTION_KEY"

-- | Read a required KEK from the named environment variable.
loadNamedKekFromEnv :: String -> IO KeyEncryptionKey
loadNamedKekFromEnv name = do
  raw <- lookupEnv name
  case raw of
    Nothing -> fatal "is not set (32 bytes, base64: head -c 32 /dev/urandom | base64)"
    Just "" -> fatal "is empty (32 bytes, base64: head -c 32 /dev/urandom | base64)"
    Just v -> either (fatal . Text.unpack) pure (keyEncryptionKeyFromBase64 (Text.pack v))
  where
    fatal reason = ioError (userError (name <> " " <> reason))

-- | Everything derived from the signing-key table in one load: the private key that signs
-- new tokens, the public key set the verifier trusts, and the precomputed JWKS document.
-- Swapped atomically on reload — readers see the old or the new record, never a mixture.
data LoadedKeys = LoadedKeys
  { signingKey :: !JWK,
    verifierJwks :: !JWKSet,
    jwksBody :: !Value
  }

-- | Load all publishable keys and assemble 'LoadedKeys'. A corrupt key row (one whose
-- stored JWK JSON does not parse, or whose ciphertext does not authenticate) fails the
-- whole load rather than being skipped silently: publishing a key set that is missing a key
-- downstreams may still hold tokens for is exactly the outage this module exists to prevent.
loadKeyMaterial :: KeyEncryptionKey -> Pool -> IO (Either Text LoadedKeys)
loadKeyMaterial kek pool = do
  result :: Either AuthError [StoredSigningKey] <-
    runEff
      . runErrorNoCallStack
      . runDatabasePool pool
      . runClockIO
      . runSigningKeyStorePostgres
      $ listPublishableSigningKeys
  pure (either (Left . Text.pack . show) (assembleKeys kek) result)

-- | Assemble the loaded rows. Exactly one active key must exist; multiple active rows
-- fail closed even if an external writer bypassed the database invariant. Every
-- publishable key is published and trusted, the signer included.
--
-- Public material comes from @public_key_jwk@ and needs no key-encryption key; only the
-- signer's @private_key_jwk@ is decrypted.
assembleKeys :: KeyEncryptionKey -> [StoredSigningKey] -> Either Text LoadedKeys
assembleKeys kek stored = do
  publics <- traverse toPublic stored
  signerRow <-
    case [sk | sk <- stored, sk.status == KeyActive] of
      [one] -> Right one
      [] -> Left "no active signing key"
      several ->
        Left
          ( "multiple active signing keys ("
              <> Text.intercalate ", " (map (.keyId) several)
              <> ") violate shomei_signing_keys_one_active; retire all but one with shomei-admin keys retire"
          )
  signerJwk <- first (decryptError signerRow) (decryptStoredSigningKey kek signerRow)
  -- The signer's own row is in 'stored', so its public projection is in 'publics'.
  signerPub <-
    maybe (Left ("signing key " <> signerRow.keyId <> " has no public material")) Right $
      lookup signerRow.keyId [(sk.keyId, jwk) | (sk, jwk) <- publics]
  let others = [jwk | (sk, jwk) <- publics, sk.keyId /= signerRow.keyId]
  pure
    LoadedKeys
      { signingKey = signerJwk,
        verifierJwks = keySetPublicJwks (KeySet signerPub others),
        jwksBody = decodeJwks (jwksDocument (signerPub : others))
      }
  where
    toPublic sk = (\jwk -> (sk, jwk)) <$> first (publicError sk) (publicJwkFromStored sk)
    publicError sk err = "signing key " <> sk.keyId <> " has unparseable public JWK JSON: " <> err
    decryptError sk = \case
      KeyDecryptFailed ->
        "signing key "
          <> sk.keyId
          <> " did not decrypt: SHOMEI_KEY_ENCRYPTION_KEY is wrong, or the row was tampered with"
      MalformedEncryptedKey reason ->
        "signing key " <> sk.keyId <> " has a malformed encrypted private key: " <> reason
      KeyJsonInvalid reason ->
        "signing key " <> sk.keyId <> " has unparseable private JWK JSON: " <> reason
    -- jwksDocument emits a JWKSet object, so the decode cannot fail; the fallback keeps
    -- the total signature rather than partially matching on it.
    decodeJwks = fromMaybe (Object KM.empty) . Aeson.decode

-- | Establish key material at boot. Generates and persists one active key for @alg@ on
-- first boot (guarded on "no active key", so restarts reuse the existing key and
-- previously-issued tokens stay verifiable); then loads via 'loadKeyMaterial'. Unlike a
-- reload, a failure here is fatal: there is no previous material to fall back on.
bootstrapKeys :: KeyEncryptionKey -> SigningAlgorithm -> Pool -> IO LoadedKeys
bootstrapKeys kek alg pool = do
  result :: Either AuthError () <-
    runEff
      . runErrorNoCallStack
      . runDatabasePool pool
      . runClockIO
      . runSigningKeyStorePostgres
      $ ensureActiveKey kek alg
  either (ioError . userError . show) pure result
  loaded <- loadKeyMaterial kek pool
  either (ioError . userError . Text.unpack) pure loaded

-- | Generate and insert an active key for @alg@ when the store has no active key. The new
-- key is encrypted before it is persisted, so a fresh
-- deployment never writes an unencrypted private key.
ensureActiveKey :: (SigningKeyStore :> es, Clock :> es, IOE :> es) => KeyEncryptionKey -> SigningAlgorithm -> Eff es ()
ensureActiveKey kek alg = do
  active <- listActiveSigningKeys
  when (null active) do
    t <- now
    jwk <- liftIO (generateSigningKeyFor alg)
    stored <- liftIO (either (ioError . userError . Text.unpack) pure (toStoredSigningKeyFor alg t jwk))
    protected <- liftIO (protectStoredSigningKey kek stored)
    insertSigningKey protected

-- | Reload key material and swap it in. On failure keep the last good material and log to
-- stderr: an operator mistake mid-rotation (say, retiring the only active key) must not
-- take the server down — it can still verify everything it issued and still sign with the
-- stale key, which downstreams still trust. The @\/health\/ready@ probe fails meanwhile, so
-- orchestration notices.
--
-- Safe from any thread: a plain 'writeIORef' of an immutable record, so a concurrent
-- request sees the old or the new material in full, never a mixture.
reloadKeys :: KeyEncryptionKey -> Pool -> IORef LoadedKeys -> IO ()
reloadKeys kek pool ref = do
  before <- readIORef ref
  loadKeyMaterial kek pool >>= either warn (swap before)
  where
    swap before after = do
      writeIORef ref after
      -- Only announce a reload that changed something, so the steady-state log stays quiet
      -- at one reload per refresh interval.
      when (kidsOf before /= kidsOf after) do
        hPutStrLn stderr $
          "[shomei] signing keys reloaded: active="
            <> Text.unpack (keyKid after.signingKey)
            <> " published="
            <> Text.unpack (Text.intercalate "," (publishedKids after))
    warn reason =
      hPutStrLn stderr ("[shomei] key reload failed: " <> Text.unpack reason <> "; keeping previous key material")
    kidsOf ks = (keyKid ks.signingKey, publishedKids ks)

-- | The @kid@ of every key in the published JWKS document, in document order.
publishedKids :: LoadedKeys -> [Text]
publishedKids ks = case ks.jwksBody of
  Object top | Just (Aeson.Array arr) <- KM.lookup "keys" top -> mapMaybe kidOf (toList arr)
  _ -> []
  where
    kidOf (Object k) | Just (Aeson.String kid) <- KM.lookup "kid" k = Just kid
    kidOf _ = Nothing