Map claimed key packages to clients (#2192)

* Add key package ref mapping table
* Add key packages to mapping when claimed
* Internal key package deref API
* Test mapping of key package refs

Co-authored-by: Marko Dimjašević <[email protected]>

Author: pcapriotti

changelog.d/1-api-changes/mls-client-api

@@ -8,6 +8,6 @@ The client JSON object now has an additional field `mls_public_keys`, containing
 ```
 At the moment, `ed25519` is the only supported signature scheme, corresponding to MLS ciphersuite 1.
 
-When creating a new client with `POST /clients`, the field `mls_public_keys` can be set, and the corresponding public keys are bound to the device identity on the backend, and will be used to veriy uploaded key packages with a matching signature scheme.
+When creating a new client with `POST /clients`, the field `mls_public_keys` can be set, and the corresponding public keys are bound to the device identity on the backend, and will be used to verify uploaded key packages with a matching signature scheme.
 
 When updating a client with `PUT /clients/:client`, the field `mls_public_keys` can also be set, with a similar effect. If a given signature scheme already has a public key set for that device, the request will fail.

changelog.d/2-features/mls

@@ -0,0 +1,3 @@
+MLS implementation progress:
+
+ - key package refs are now mapped after being claimed

docs/reference/cassandra-schema.cql

@@ -765,6 +765,28 @@ CREATE TABLE brig_test.mls_key_packages (
     AND read_repair_chance = 0.0
     AND speculative_retry = '99PERCENTILE';
 
+CREATE TABLE brig_test.mls_key_package_refs (
+    ref blob PRIMARY KEY,
+    client text,
+    conv uuid,
+    conv_domain text,
+    domain text,
+    user uuid
+) WITH bloom_filter_fp_chance = 0.1
+    AND caching = {'keys': 'ALL', 'rows_per_partition': 'NONE'}
+    AND comment = ''
+    AND compaction = {'class': 'org.apache.cassandra.db.compaction.LeveledCompactionStrategy'}
+    AND compression = {'chunk_length_in_kb': '64', 'class': 'org.apache.cassandra.io.compress.LZ4Compressor'}
+    AND crc_check_chance = 1.0
+    AND dclocal_read_repair_chance = 0.1
+    AND default_time_to_live = 0
+    AND gc_grace_seconds = 864000
+    AND max_index_interval = 2048
+    AND memtable_flush_period_in_ms = 0
+    AND min_index_interval = 128
+    AND read_repair_chance = 0.0
+    AND speculative_retry = '99PERCENTILE';
+
 CREATE TABLE brig_test.excluded_phones (
     prefix text PRIMARY KEY,
     comment text

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

@@ -19,15 +19,18 @@
 
 module Wire.API.MLS.Credential where
 
-import Data.Aeson
-import Data.Aeson.Types
+import Data.Aeson (FromJSON (..), FromJSONKey (..), ToJSON (..), ToJSONKey (..))
+import qualified Data.Aeson as Aeson
+import qualified Data.Aeson.Types as Aeson
 import Data.Binary
 import Data.Binary.Get
 import Data.Binary.Parser
 import Data.Binary.Parser.Char8
 import Data.Domain
 import Data.Id
 import Data.Qualified
+import Data.Schema
+import qualified Data.Swagger as S
 import qualified Data.Text as T
 import Data.UUID
 import Imports
@@ -101,23 +104,32 @@ parseSignatureScheme name =
     (signatureSchemeFromName name)
 
 instance FromJSON SignatureSchemeTag where
-  parseJSON = withText "SignatureScheme" parseSignatureScheme
+  parseJSON = Aeson.withText "SignatureScheme" parseSignatureScheme
 
 instance FromJSONKey SignatureSchemeTag where
-  fromJSONKey = FromJSONKeyTextParser parseSignatureScheme
+  fromJSONKey = Aeson.FromJSONKeyTextParser parseSignatureScheme
 
 instance ToJSON SignatureSchemeTag where
-  toJSON = String . signatureSchemeName
+  toJSON = Aeson.String . signatureSchemeName
 
 instance ToJSONKey SignatureSchemeTag where
-  toJSONKey = toJSONKeyText signatureSchemeName
+  toJSONKey = Aeson.toJSONKeyText signatureSchemeName
 
 data ClientIdentity = ClientIdentity
   { ciDomain :: Domain,
     ciUser :: UserId,
     ciClient :: ClientId
   }
   deriving stock (Eq, Show, Generic)
+  deriving (FromJSON, ToJSON, S.ToSchema) via Schema ClientIdentity
+
+instance ToSchema ClientIdentity where
+  schema =
+    object "ClientIdentity" $
+      ClientIdentity
+        <$> ciDomain .= field "domain" schema
+        <*> ciUser .= field "user_id" schema
+        <*> ciClient .= field "client_id" schema
 
 instance ParseMLS ClientIdentity where
   parseMLS = do

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

@@ -69,6 +69,7 @@ import Data.Singletons.TH
 import qualified Data.Swagger as S
 import Data.Time.Clock.POSIX
 import Imports
+import Web.HttpApiData
 import Wire.API.Arbitrary
 import Wire.API.MLS.CipherSuite
 import Wire.API.MLS.Credential
@@ -97,6 +98,7 @@ instance ToSchema KeyPackageData where
 data KeyPackageBundleEntry = KeyPackageBundleEntry
   { kpbeUser :: Qualified UserId,
     kpbeClient :: ClientId,
+    kpbeRef :: KeyPackageRef,
     kpbeKeyPackage :: KeyPackageData
   }
   deriving stock (Eq, Ord)
@@ -107,6 +109,7 @@ instance ToSchema KeyPackageBundleEntry where
       KeyPackageBundleEntry
         <$> kpbeUser .= qualifiedObjectSchema "user" schema
         <*> kpbeClient .= field "client" schema
+        <*> kpbeRef .= field "key_package_ref" schema
         <*> kpbeKeyPackage .= field "key_package" schema
 
 newtype KeyPackageBundle = KeyPackageBundle {kpbEntries :: Set KeyPackageBundleEntry}
@@ -127,6 +130,23 @@ instance ToSchema KeyPackageCount where
     object "OwnKeyPackages" $
       KeyPackageCount <$> unKeyPackageCount .= field "count" schema
 
+newtype KeyPackageRef = KeyPackageRef {unKeyPackageRef :: ByteString}
+  deriving stock (Eq, Ord, Show)
+  deriving (FromHttpApiData, ToHttpApiData, S.ToParamSchema) via Base64ByteString
+
+instance ToSchema KeyPackageRef where
+  schema = named "KeyPackageRef" $ unKeyPackageRef .= fmap KeyPackageRef base64Schema
+
+instance ParseMLS KeyPackageRef where
+  parseMLS = KeyPackageRef <$> getByteString 16
+
+kpRef :: CipherSuiteTag -> KeyPackageData -> KeyPackageRef
+kpRef cs =
+  KeyPackageRef
+    . csHash cs "MLS 1.0 KeyPackage Reference"
+    . LBS.toStrict
+    . kpData
+
 --------------------------------------------------------------------------------
 
 newtype ProtocolVersion = ProtocolVersion {pvNumber :: Word8}
@@ -247,19 +267,6 @@ data KeyPackage = KeyPackage
   }
   deriving stock (Eq, Show)
 
-newtype KeyPackageRef = KeyPackageRef {unKeyPackageRef :: ByteString}
-  deriving stock (Eq, Show)
-
-instance ParseMLS KeyPackageRef where
-  parseMLS = KeyPackageRef <$> getByteString 16
-
-kpRef :: CipherSuiteTag -> KeyPackageData -> KeyPackageRef
-kpRef cs =
-  KeyPackageRef
-    . csHash cs "MLS 1.0 KeyPackage Reference"
-    . LBS.toStrict
-    . kpData
-
 instance ParseMLS KeyPackage where
   parseMLS = fst <$> kpSigOffset
 

libs/wire-api/src/Wire/API/Routes/Internal/Brig.hs

@@ -19,6 +19,7 @@ module Wire.API.Routes.Internal.Brig
   ( API,
     EJPD_API,
     AccountAPI,
+    MLSAPI,
     EJPDRequest,
     GetAccountFeatureConfig,
     PutAccountFeatureConfig,
@@ -39,6 +40,8 @@ import Servant.Swagger (HasSwagger (toSwagger))
 import Servant.Swagger.Internal.Orphans ()
 import Servant.Swagger.UI
 import Wire.API.Connection
+import Wire.API.MLS.Credential
+import Wire.API.MLS.KeyPackage
 import Wire.API.Routes.Internal.Brig.Connection
 import Wire.API.Routes.Internal.Brig.EJPD
 import Wire.API.Routes.MultiVerb
@@ -134,9 +137,24 @@ type AccountAPI =
         :> MultiVerb 'POST '[Servant.JSON] RegisterInternalResponses (Either RegisterError SelfProfile)
     )
 
+type MLSAPI = GetClientByKeyPackageRef
+
+type GetClientByKeyPackageRef =
+  Summary "Resolve an MLS key package ref to a qualified client ID"
+    :> "mls"
+    :> "key-packages"
+    :> Capture "ref" KeyPackageRef
+    :> MultiVerb
+         'GET
+         '[Servant.JSON]
+         '[ RespondEmpty 404 "Key package ref not found",
+            Respond 200 "Key package ref found" ClientIdentity
+          ]
+         (Maybe ClientIdentity)
+
 type API =
   "i"
-    :> (EJPD_API :<|> AccountAPI)
+    :> (EJPD_API :<|> AccountAPI :<|> MLSAPI)
 
 type SwaggerDocsAPI = "api" :> "internal" :> SwaggerSchemaUI "swagger-ui" "swagger.json"
 

libs/wire-api/src/Wire/API/User/Client.hs

@@ -450,7 +450,7 @@ data Client = Client
   deriving (Arbitrary) via (GenericUniform Client)
   deriving (FromJSON, ToJSON, Swagger.ToSchema) via Schema Client
 
-type MLSPublicKeys = Map SignatureSchemeTag LByteString
+type MLSPublicKeys = Map SignatureSchemeTag ByteString
 
 instance ToSchema Client where
   schema =
@@ -473,7 +473,7 @@ mlsPublicKeysSchema =
     (fromMaybe mempty)
     ( optField
         "mls_public_keys"
-        (map_ base64SchemaL)
+        (map_ base64Schema)
     )
 
 modelClient :: Doc.Model

services/brig/brig.cabal

@@ -624,6 +624,7 @@ executable brig-schema
       V66_PersonalFeatureConfCallInit
       V67_MLSKeyPackages
       V68_AddMLSPublicKeys
+      V69_MLSKeyPackageRefMapping
       V9
       Paths_brig
   hs-source-dirs:

services/brig/schema/src/Main.hs

@@ -78,6 +78,7 @@ import qualified V65_FederatedConnections
 import qualified V66_PersonalFeatureConfCallInit
 import qualified V67_MLSKeyPackages
 import qualified V68_AddMLSPublicKeys
+import qualified V69_MLSKeyPackageRefMapping
 import qualified V9
 
 main :: IO ()
@@ -145,7 +146,8 @@ main = do
       V65_FederatedConnections.migration,
       V66_PersonalFeatureConfCallInit.migration,
       V67_MLSKeyPackages.migration,
-      V68_AddMLSPublicKeys.migration
+      V68_AddMLSPublicKeys.migration,
+      V69_MLSKeyPackageRefMapping.migration
       -- When adding migrations here, don't forget to update
       -- 'schemaVersion' in Brig.App
 

services/brig/schema/src/V69_MLSKeyPackageRefMapping.hs

@@ -0,0 +1,44 @@
+{-# LANGUAGE QuasiQuotes #-}
+
+-- This file is part of the Wire Server implementation.
+--
+-- Copyright (C) 2022 Wire Swiss GmbH <[email protected]>
+--
+-- This program is free software: you can redistribute it and/or modify it under
+-- the terms of the GNU Affero General Public License as published by the Free
+-- Software Foundation, either version 3 of the License, or (at your option) any
+-- later version.
+--
+-- This program is distributed in the hope that it will be useful, but WITHOUT
+-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
+-- details.
+--
+-- You should have received a copy of the GNU Affero General Public License along
+-- with this program. If not, see <https://www.gnu.org/licenses/>.
+
+module V69_MLSKeyPackageRefMapping
+  ( migration,
+  )
+where
+
+import Cassandra.Schema
+import Imports
+import Text.RawString.QQ
+
+migration :: Migration
+migration =
+  Migration 69 "Add key package ref mapping" $
+    schema'
+      [r|
+        CREATE TABLE mls_key_package_refs
+            ( ref blob
+            , domain text
+            , user uuid
+            , client text
+            , conv_domain text
+            , conv uuid
+            , PRIMARY KEY (ref)
+        ) WITH compaction = {'class': 'org.apache.cassandra.db.compaction.LeveledCompactionStrategy'}
+          AND gc_grace_seconds = 864000;
+     |]