Map claimed key packages to clients

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 <opensource@wire.com>
+--
+-- 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;
+     |]

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

@@ -35,6 +35,7 @@ import Brig.App
 import Brig.Data.Activation
 import qualified Brig.Data.Client as Data
 import qualified Brig.Data.Connection as Data
+import qualified Brig.Data.MLS.KeyPackage as Data
 import qualified Brig.Data.User as Data
 import qualified Brig.IO.Intra as Intra
 import Brig.Options hiding (internalEvents, sesQueue)
@@ -72,6 +73,8 @@ import Servant.Swagger.Internal.Orphans ()
 import Servant.Swagger.UI
 import qualified System.Logger.Class as Log
 import Wire.API.ErrorDescription
+import Wire.API.MLS.Credential
+import Wire.API.MLS.KeyPackage
 import qualified Wire.API.Routes.Internal.Brig as BrigIRoutes
 import Wire.API.Routes.Internal.Brig.Connection
 import Wire.API.Routes.Named
@@ -84,7 +87,7 @@ import Wire.API.User.RichInfo
 -- Sitemap (servant)
 
 servantSitemap :: ServerT BrigIRoutes.API (Handler r)
-servantSitemap = ejpdAPI :<|> accountAPI
+servantSitemap = ejpdAPI :<|> accountAPI :<|> mlsAPI
 
 ejpdAPI :: ServerT BrigIRoutes.EJPD_API (Handler r)
 ejpdAPI =
@@ -95,6 +98,9 @@ ejpdAPI =
     :<|> getConnectionsStatusUnqualified
     :<|> getConnectionsStatus
 
+mlsAPI :: ServerT BrigIRoutes.MLSAPI (Handler r)
+mlsAPI = getClientByKeyPackageRef
+
 accountAPI :: ServerT BrigIRoutes.AccountAPI (Handler r)
 accountAPI = Named @"createUserNoVerify" createUserNoVerify
 
@@ -115,6 +121,9 @@ deleteAccountFeatureConfig uid =
 swaggerDocsAPI :: Servant.Server BrigIRoutes.SwaggerDocsAPI
 swaggerDocsAPI = swaggerSchemaUIServer BrigIRoutes.swaggerDoc
 
+getClientByKeyPackageRef :: KeyPackageRef -> Handler r (Maybe ClientIdentity)
+getClientByKeyPackageRef ref = runMaybeT $ Data.derefKeyPackage ref
+
 ---------------------------------------------------------------------------
 -- Sitemap (wai-route)
 

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

@@ -65,8 +65,8 @@ claimLocalKeyPackages lusr target = do
     mkEntry :: ClientId -> AppIO r (Maybe KeyPackageBundleEntry)
     mkEntry c =
       runMaybeT $
-        KeyPackageBundleEntry (qUntagged target) c
-          <$> Data.claimKeyPackage (tUnqualified target) c
+        uncurry (KeyPackageBundleEntry (qUntagged target) c)
+          <$> Data.claimKeyPackage target c
 
 countKeyPackages :: Local UserId -> ClientId -> Handler r KeyPackageCount
 countKeyPackages lusr c =

services/brig/src/Brig/App.hs

@@ -136,7 +136,7 @@ import Util.Options
 import Wire.API.User.Identity (Email)
 
 schemaVersion :: Int32
-schemaVersion = 68
+schemaVersion = 69
 
 -------------------------------------------------------------------------------
 -- Environment

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

@@ -67,6 +67,7 @@ import Control.Monad.Random (randomRIO)
 import Control.Retry
 import qualified Data.ByteString.Base64 as B64
 import Data.ByteString.Conversion (toByteString, toByteString')
+import qualified Data.ByteString.Lazy as LBS
 import qualified Data.HashMap.Strict as HashMap
 import Data.Id
 import Data.Json.Util (UTCTimeMillis, toUTCTimeMillis)
@@ -272,18 +273,18 @@ lookupMLSPublicKey u c ss =
 addMLSPublicKeys ::
   UserId ->
   ClientId ->
-  [(SignatureSchemeTag, LByteString)] ->
+  [(SignatureSchemeTag, ByteString)] ->
   ExceptT ClientDataError (AppIO r) ()
 addMLSPublicKeys u c = traverse_ (uncurry (addMLSPublicKey u c))
 
 addMLSPublicKey ::
   UserId ->
   ClientId ->
   SignatureSchemeTag ->
-  LByteString ->
+  ByteString ->
   ExceptT ClientDataError (AppIO r) ()
 addMLSPublicKey u c ss pk = do
-  rows <- trans insertMLSPublicKeys (params LocalQuorum (u, c, ss, Blob pk))
+  rows <- trans insertMLSPublicKeys (params LocalQuorum (u, c, ss, Blob (LBS.fromStrict pk)))
   case rows of
     [row]
       | C.fromRow 0 row /= Right (Just True) ->
@@ -380,7 +381,7 @@ toClient keys (cid, cty, tme, lbl, cls, cok, lat, lon, mdl, cps) =
       clientLocation = location <$> lat <*> lon,
       clientModel = mdl,
       clientCapabilities = ClientCapabilityList $ maybe Set.empty (Set.fromList . C.fromSet) cps,
-      clientMLSPublicKeys = fmap fromBlob (Map.fromList keys)
+      clientMLSPublicKeys = fmap (LBS.toStrict . fromBlob) (Map.fromList keys)
     }
 
 toPubClient :: (ClientId, Maybe ClientClass) -> PubClient