[FS-937] Validate Remotely Claimed Key Packages (#2692)

* Provide a custom Show instance for ClientIdentity
* Align function types in KeyPackages and Validation
* Validate remotely claimed key packages
* Restrict validation of MLS public keys to locals
* Update Brig integration test utils
* Update the remote key package claim test

changelog.d/1-api-changes/validate-remotely-claimed-key-packages

@@ -0,0 +1 @@
+Validate remotely claimed key packages

libs/wire-api/src/Wire/API/Error/Brig.hs

@@ -64,6 +64,8 @@ data BrigError
   | PasswordAuthenticationFailed
   | TooManyTeamInvitations
   | InsufficientTeamPermissions
+  | KeyPackageDecodingError
+  | InvalidKeyPackageRef
 
 instance KnownError (MapError e) => IsSwaggerError (e :: BrigError) where
   addToSwagger = addStaticErrorToSwagger @(MapError e)
@@ -172,3 +174,7 @@ type instance MapError 'PasswordAuthenticationFailed = 'StaticError 403 "passwor
 type instance MapError 'TooManyTeamInvitations = 'StaticError 403 "too-many-team-invitations" "Too many team invitations for this team"
 
 type instance MapError 'InsufficientTeamPermissions = 'StaticError 403 "insufficient-permissions" "Insufficient team permissions"
+
+type instance MapError 'KeyPackageDecodingError = 'StaticError 409 "decoding-error" "Key package could not be TLS-decoded"
+
+type instance MapError 'InvalidKeyPackageRef = 'StaticError 409 "invalid-reference" "Key package's reference does not match its data"

libs/wire-api/src/Wire/API/MLS/Credential.hs

@@ -139,9 +139,17 @@ data ClientIdentity = ClientIdentity
     ciUser :: UserId,
     ciClient :: ClientId
   }
-  deriving stock (Eq, Ord, Show, Generic)
+  deriving stock (Eq, Ord, Generic)
   deriving (FromJSON, ToJSON, S.ToSchema) via Schema ClientIdentity
 
+instance Show ClientIdentity where
+  show (ClientIdentity dom u c) =
+    show u
+      <> ":"
+      <> T.unpack (client c)
+      <> "@"
+      <> T.unpack (domainText dom)
+
 cidQualifiedClient :: ClientIdentity -> Qualified (UserId, ClientId)
 cidQualifiedClient cid = Qualified (ciUser cid, ciClient cid) (ciDomain cid)
 

services/brig/src/Brig/API/Error.hs

@@ -182,6 +182,8 @@ clientDataError (ClientReAuthError e) = reauthError e
 clientDataError ClientMissingAuth = StdError (errorToWai @'E.MissingAuth)
 clientDataError MalformedPrekeys = StdError (errorToWai @'E.MalformedPrekeys)
 clientDataError MLSPublicKeyDuplicate = StdError (errorToWai @'E.MLSDuplicatePublicKey)
+clientDataError KeyPackageDecodingError = StdError (errorToWai @'E.KeyPackageDecodingError)
+clientDataError InvalidKeyPackageRef = StdError (errorToWai @'E.InvalidKeyPackageRef)
 
 deleteUserError :: DeleteUserError -> Error
 deleteUserError DeleteUserInvalid = StdError (errorToWai @'E.InvalidUser)

services/brig/src/Brig/API/MLS/KeyPackages.hs

@@ -42,6 +42,7 @@ import Wire.API.Federation.API
 import Wire.API.Federation.API.Brig
 import Wire.API.MLS.Credential
 import Wire.API.MLS.KeyPackage
+import Wire.API.MLS.Serialisation
 import Wire.API.Team.LegalHold
 import Wire.API.User.Client
 
@@ -57,12 +58,11 @@ claimKeyPackages ::
   Maybe ClientId ->
   Handler r KeyPackageBundle
 claimKeyPackages lusr target skipOwn =
-  withExceptT clientError $
-    foldQualified
-      lusr
-      (claimLocalKeyPackages (qUntagged lusr) skipOwn)
-      (claimRemoteKeyPackages lusr)
-      target
+  foldQualified
+    lusr
+    (withExceptT clientError . claimLocalKeyPackages (qUntagged lusr) skipOwn)
+    (claimRemoteKeyPackages lusr)
+    target
 
 claimLocalKeyPackages ::
   Qualified UserId ->
@@ -96,22 +96,35 @@ claimLocalKeyPackages qusr skipOwn target = do
 claimRemoteKeyPackages ::
   Local UserId ->
   Remote UserId ->
-  ExceptT ClientError (AppT r) KeyPackageBundle
+  Handler r KeyPackageBundle
 claimRemoteKeyPackages lusr target = do
   bundle <-
-    (handleFailure =<<) $
-      withExceptT ClientFederationError $
+    withExceptT clientError
+      . (handleFailure =<<)
+      $ withExceptT ClientFederationError $
         runBrigFederatorClient (tDomain target) $
           fedClient @'Brig @"claim-key-packages" $
             ClaimKeyPackageRequest
               { ckprClaimant = tUnqualified lusr,
                 ckprTarget = tUnqualified target
               }
 
-  -- set up mappings for all claimed key packages
-  wrapClientE $
-    for_ (kpbEntries bundle) $ \e ->
-      Data.mapKeyPackageRef (kpbeRef e) (kpbeUser e) (kpbeClient e)
+  -- validate and set up mappings for all claimed key packages
+  for_ (kpbEntries bundle) $ \e -> do
+    let cid = mkClientIdentity (kpbeUser e) (kpbeClient e)
+    kpRaw <-
+      withExceptT (const . clientDataError $ KeyPackageDecodingError)
+        . except
+        . decodeMLS'
+        . kpData
+        . kpbeKeyPackage
+        $ e
+    (refVal, _) <- validateKeyPackage cid kpRaw
+    unless (refVal == kpbeRef e)
+      . throwE
+      . clientDataError
+      $ InvalidKeyPackageRef
+    wrapClientE $ Data.mapKeyPackageRef (kpbeRef e) (kpbeUser e) (kpbeClient e)
 
   pure bundle
   where

services/brig/src/Brig/API/MLS/KeyPackages/Validation.hs

@@ -35,6 +35,7 @@ import Brig.Options
 import Control.Applicative
 import Control.Lens (view)
 import qualified Data.ByteString.Lazy as LBS
+import Data.Qualified
 import Data.Time.Clock
 import Data.Time.Clock.POSIX
 import Imports
@@ -46,8 +47,12 @@ import Wire.API.MLS.Extension
 import Wire.API.MLS.KeyPackage
 import Wire.API.MLS.Serialisation
 
-validateKeyPackage :: ClientIdentity -> RawMLS KeyPackage -> Handler r (KeyPackageRef, KeyPackageData)
+validateKeyPackage ::
+  ClientIdentity ->
+  RawMLS KeyPackage ->
+  Handler r (KeyPackageRef, KeyPackageData)
 validateKeyPackage identity (RawMLS (KeyPackageData -> kpd) kp) = do
+  loc <- qualifyLocal ()
   -- get ciphersuite
   cs <-
     maybe
@@ -60,19 +65,32 @@ validateKeyPackage identity (RawMLS (KeyPackageData -> kpd) kp) = do
   when (signatureScheme ss /= bcSignatureScheme (kpCredential kp)) $
     mlsProtocolError "Signature scheme incompatible with ciphersuite"
 
-  -- authenticate signature key
-  key <-
-    fmap LBS.toStrict $
-      maybe
-        (mlsProtocolError "No key associated to the given identity and signature scheme")
-        pure
-        =<< lift (wrapClient (Data.lookupMLSPublicKey (ciUser identity) (ciClient identity) ss))
-  when (key /= bcSignatureKey (kpCredential kp)) $
-    mlsProtocolError "Unrecognised signature key"
+  -- Authenticate signature key. This is performed only upon uploading a key
+  -- package for a local client.
+  foldQualified
+    loc
+    ( \_ -> do
+        key <-
+          fmap LBS.toStrict $
+            maybe
+              (mlsProtocolError "No key associated to the given identity and signature scheme")
+              pure
+              =<< lift (wrapClient (Data.lookupMLSPublicKey (ciUser identity) (ciClient identity) ss))
+        when (key /= bcSignatureKey (kpCredential kp)) $
+          mlsProtocolError "Unrecognised signature key"
+    )
+    (pure . const ())
+    (cidQualifiedClient identity)
 
   -- validate signature
-  unless (csVerifySignature cs key (rmRaw (kpTBS kp)) (kpSignature kp)) $
-    mlsProtocolError "Invalid signature"
+  unless
+    ( csVerifySignature
+        cs
+        (bcSignatureKey (kpCredential kp))
+        (rmRaw (kpTBS kp))
+        (kpSignature kp)
+    )
+    $ mlsProtocolError "Invalid signature"
   -- validate protocol version
   maybe
     (mlsProtocolError "Unsupported protocol version")

services/brig/src/Brig/Data/Client.hs

@@ -97,6 +97,8 @@ data ClientDataError
   | ClientMissingAuth
   | MalformedPrekeys
   | MLSPublicKeyDuplicate
+  | KeyPackageDecodingError
+  | InvalidKeyPackageRef
 
 -- | Re-authentication policy.
 --

services/brig/test/integration/API/MLS.hs

@@ -23,7 +23,6 @@ import Bilge.Assert
 import Brig.Options
 import Control.Timeout
 import qualified Data.Aeson as Aeson
-import qualified Data.ByteString as BS
 import Data.ByteString.Conversion
 import Data.Default
 import Data.Id
@@ -32,14 +31,14 @@ import qualified Data.Set as Set
 import Data.Timeout
 import Federation.Util
 import Imports
-import Test.QuickCheck hiding ((===))
 import Test.Tasty
 import Test.Tasty.HUnit
 import UnliftIO.Temporary
 import Util
 import Web.HttpApiData
 import Wire.API.MLS.Credential
 import Wire.API.MLS.KeyPackage
+import Wire.API.MLS.Serialisation
 import Wire.API.User
 import Wire.API.User.Client
 
@@ -186,13 +185,18 @@ testKeyPackageRemoteClaim opts brig = do
 
   u' <- userQualifiedId <$> randomUser brig
 
-  entries <-
-    liftIO . replicateM 2 . generate $
-      -- claimed key packages are not validated by the backend, so it is fine to
-      -- make up some random data here
-      KeyPackageBundleEntry u <$> arbitrary
-        <*> (KeyPackageRef . BS.pack <$> vector 32)
-        <*> (KeyPackageData . BS.pack <$> vector 64)
+  qcid <- mkClientIdentity u <$> randomClient
+  entries <- withSystemTempDirectory "mls" $ \tmp -> do
+    initStore tmp qcid
+    replicateM 2 $ do
+      (r, kp) <- generateKeyPackage tmp qcid Nothing
+      pure $
+        KeyPackageBundleEntry
+          { kpbeUser = u,
+            kpbeClient = ciClient qcid,
+            kpbeRef = kp,
+            kpbeKeyPackage = KeyPackageData . rmRaw $ r
+          }
   let mockBundle = KeyPackageBundle (Set.fromList entries)
   (bundle :: KeyPackageBundle, _reqs) <-
     liftIO . withTempMockFederator opts (Aeson.encode mockBundle) $

services/brig/test/integration/API/MLS/Util.hs

@@ -23,19 +23,20 @@ import Bilge.Assert
 import Data.Aeson (object, toJSON, (.=))
 import Data.ByteString.Conversion
 import Data.Default
-import Data.Domain
 import Data.Id
 import Data.Json.Util
 import qualified Data.Map as Map
 import Data.Qualified
-import qualified Data.Text as T
+import qualified Data.Text as Text
 import Data.Timeout
 import Imports
 import System.FilePath
 import System.Process
+import Test.Tasty.HUnit
 import Util
 import Wire.API.MLS.Credential
 import Wire.API.MLS.KeyPackage
+import Wire.API.MLS.Serialisation
 import Wire.API.User.Client
 
 data SetKey = SetKey | DontSetKey
@@ -49,6 +50,39 @@ data KeyingInfo = KeyingInfo
 instance Default KeyingInfo where
   def = KeyingInfo SetKey Nothing
 
+cliCmd :: FilePath -> ClientIdentity -> [String]
+cliCmd tmp qcid =
+  ["mls-test-cli", "--store", tmp </> (show qcid <> ".db")]
+
+initStore ::
+  HasCallStack =>
+  MonadIO m =>
+  FilePath ->
+  ClientIdentity ->
+  m ()
+initStore tmp qcid = do
+  let cmd0 = cliCmd tmp qcid
+  void . liftIO . flip spawn Nothing . shell . unwords $
+    cmd0 <> ["init", show qcid]
+
+generateKeyPackage ::
+  HasCallStack =>
+  MonadIO m =>
+  FilePath ->
+  ClientIdentity ->
+  Maybe Timeout ->
+  m (RawMLS KeyPackage, KeyPackageRef)
+generateKeyPackage tmp qcid lifetime = do
+  let cmd0 = cliCmd tmp qcid
+  kp <-
+    liftIO $
+      decodeMLSError <=< (flip spawn Nothing . shell . unwords) $
+        cmd0
+          <> ["key-package", "create"]
+          <> (("--lifetime " <>) . show . (#> Second) <$> maybeToList lifetime)
+  let ref = fromJust (kpRef' kp)
+  pure (kp, ref)
+
 uploadKeyPackages ::
   HasCallStack =>
   Brig ->
@@ -59,20 +93,10 @@ uploadKeyPackages ::
   Int ->
   Http ()
 uploadKeyPackages brig tmp KeyingInfo {..} u c n = do
-  let cmd0 = ["mls-test-cli", "--store", tmp </> (clientId <> ".db")]
-      clientId =
-        show (qUnqualified u)
-          <> ":"
-          <> T.unpack (client c)
-          <> "@"
-          <> T.unpack (domainText (qDomain u))
-  void . liftIO . flip spawn Nothing . shell . unwords $
-    cmd0 <> ["init", clientId]
-  kps <-
-    replicateM n . liftIO . flip spawn Nothing . shell . unwords $
-      cmd0
-        <> ["key-package", "create"]
-        <> (("--lifetime " <>) . show . (#> Second) <$> maybeToList kiLifetime)
+  let cmd0 = cliCmd tmp cid
+      cid = mkClientIdentity u c
+  initStore tmp cid
+  kps <- replicateM n (fst <$> generateKeyPackage tmp cid kiLifetime)
   when (kiSetKey == SetKey) $
     do
       pk <-
@@ -85,7 +109,7 @@ uploadKeyPackages brig tmp KeyingInfo {..} u c n = do
             . json defUpdateClient {updateClientMLSPublicKeys = Map.fromList [(Ed25519, pk)]}
         )
       !!! const 200 === statusCode
-  let upload = object ["key_packages" .= toJSON (map Base64ByteString kps)]
+  let upload = object ["key_packages" .= toJSON (map (Base64ByteString . rmRaw) kps)]
   post
     ( brig
         . paths ["mls", "key-packages", "self", toByteString' c]
@@ -102,3 +126,8 @@ getKeyPackageCount brig u c =
           . zUser (qUnqualified u)
       )
     <!! const 200 === statusCode
+
+decodeMLSError :: ParseMLS a => ByteString -> IO a
+decodeMLSError s = case decodeMLS' s of
+  Left e -> assertFailure ("Could not parse MLS object: " <> Text.unpack e)
+  Right x -> pure x