Parse and validate extensions

Author: pcapriotti

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

@@ -19,7 +19,8 @@
 
 module Wire.API.MLS.CipherSuite where
 
-import Data.Binary
+import Data.Word
+import Wire.API.MLS.Serialisation
 
 newtype CipherSuite = CipherSuite {cipherSuiteNumber :: Word16}
-  deriving newtype (Binary)
+  deriving newtype (ParseMLS)

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

@@ -34,9 +34,13 @@ data Credential = BasicCredential
     bcSignatureScheme :: SignatureScheme,
     bcSignatureKey :: ByteString
   }
-  deriving (Generic)
 
-instance Binary Credential
+instance ParseMLS Credential where
+  parseMLS =
+    BasicCredential
+      <$> parseMLSBytes @Word16
+      <*> parseMLS
+      <*> parseMLSBytes @Word16
 
 data CredentialType = BasicCredentialType
 
@@ -47,7 +51,7 @@ credentialType (BasicCredential _ _ _) = BasicCredentialType
 --
 -- See <https://www.iana.org/assignments/tls-parameters/tls-parameters.xhtml#tls-signaturescheme>.
 newtype SignatureScheme = SignatureScheme {signatureSchemeNumber :: Word16}
-  deriving newtype (Binary)
+  deriving newtype (ParseMLS)
 
 data ClientIdentity = ClientIdentity
   { ciDomain :: Domain,

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

@@ -1,3 +1,6 @@
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE StandaloneKindSignatures #-}
+
 -- This file is part of the Wire Server implementation.
 --
 -- Copyright (C) 2022 Wire Swiss GmbH <[email protected]>
@@ -17,14 +20,19 @@
 
 module Wire.API.MLS.KeyPackage where
 
+import Control.Applicative
+import Control.Error.Util
 import Data.Aeson (FromJSON)
 import Data.Binary
 import Data.Json.Util
 import Data.Schema
+import Data.Singletons
+import Data.Singletons.TH
 import qualified Data.Swagger as S
 import Imports
 import Wire.API.MLS.CipherSuite
 import Wire.API.MLS.Credential
+import Wire.API.MLS.Proposal
 import Wire.API.MLS.Serialisation
 
 data KeyPackageUpload = KeyPackageUpload
@@ -47,16 +55,77 @@ instance ToSchema KeyPackageData where
 --------------------------------------------------------------------------------
 
 data ProtocolVersion = ProtocolReserved | ProtocolMLS
-  deriving stock (Enum)
-  deriving (Binary) via EnumBinary Word8 ProtocolVersion
+  deriving stock (Bounded, Enum)
+  deriving (ParseMLS) via EnumMLS Word8 ProtocolVersion
 
 data Extension = Extension
   { extType :: Word16,
     extData :: ByteString
   }
-  deriving (Generic)
 
-instance Binary Extension
+instance ParseMLS Extension where
+  parseMLS = Extension <$> parseMLS <*> parseMLSBytes @Word32
+
+data ExtensionTag
+  = ReservedExtensionTag
+  | CapabilitiesExtensionTag
+  | LifetimeExtensionTag
+  deriving (Bounded, Enum)
+
+$(genSingletons [''ExtensionTag])
+
+type family ExtensionType (t :: ExtensionTag) :: * where
+  ExtensionType 'ReservedExtensionTag = ()
+  ExtensionType 'CapabilitiesExtensionTag = Capabilities
+  ExtensionType 'LifetimeExtensionTag = Lifetime
+
+parseExtension :: Sing t -> Get (ExtensionType t)
+parseExtension SReservedExtensionTag = pure ()
+parseExtension SCapabilitiesExtensionTag = parseMLS
+parseExtension SLifetimeExtensionTag = parseMLS
+
+data SomeExtension where
+  SomeExtension :: Sing t -> ExtensionType t -> SomeExtension
+
+decodeExtension :: Extension -> Maybe SomeExtension
+decodeExtension e = do
+  t <- safeToEnum (fromIntegral (extType e))
+  hush $
+    withSomeSing t $ \st ->
+      decodeMLSWith' (SomeExtension st <$> parseExtension st) (extData e)
+
+-- t <- parse
+-- parseMLS = do
+--   t <- parseMLS
+--   case toSing t of
+--     SomeSing st -> SomeExtension st <$> parseExtension st
+
+data Capabilities = Capabilities
+  { capVersions :: [ProtocolVersion],
+    capCiphersuites :: [CipherSuite],
+    capExtensions :: [Word16],
+    capProposals :: [ProposalType]
+  }
+
+instance ParseMLS Capabilities where
+  parseMLS =
+    Capabilities
+      <$> parseMLSVector @Word8 parseMLS
+      <*> parseMLSVector @Word8 parseMLS
+      <*> parseMLSVector @Word8 parseMLS
+      <*> parseMLSVector @Word8 parseMLS
+
+-- | Seconds since the UNIX epoch.
+newtype Timestamp = Timestamp {timestampSeconds :: Word64}
+  deriving newtype (ParseMLS)
+
+data Lifetime = Lifetime
+  { ltNotBefore :: Timestamp,
+    ltNotAfter :: Timestamp
+  }
+
+instance ParseMLS Lifetime where
+  parseMLS = Lifetime <$> parseMLS <*> parseMLS
 
 data KeyPackageTBS = KeyPackageTBS
   { kpProtocolVersion :: ProtocolVersion,
@@ -65,15 +134,20 @@ data KeyPackageTBS = KeyPackageTBS
     kpCredential :: Credential,
     kpExtensions :: [Extension]
   }
-  deriving (Generic)
 
-instance Binary KeyPackageTBS
+instance ParseMLS KeyPackageTBS where
+  parseMLS =
+    KeyPackageTBS
+      <$> parseMLS
+      <*> parseMLS
+      <*> parseMLSBytes @Word16
+      <*> parseMLS
+      <*> parseMLSVector @Word32 parseMLS
 
 data KeyPackage = KeyPackage
   { kpTBS :: KeyPackageTBS,
     kpSignature :: ByteString
   }
-  deriving (Generic)
-  deriving (ParseMLS) via BinaryMLS KeyPackage
 
-instance Binary KeyPackage
+instance ParseMLS KeyPackage where
+  parseMLS = KeyPackage <$> parseMLS <*> parseMLSBytes @Word16

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

@@ -0,0 +1,35 @@
+-- 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 Wire.API.MLS.Proposal where
+
+import Data.Binary
+import Imports
+import Wire.API.MLS.Serialisation
+
+data ProposalType
+  = AddProposal
+  | UpdateProposal
+  | RemoveProposal
+  | PreSharedKeyProposal
+  | ReInitProposal
+  | ExternalInitProposal
+  | AppAckProposal
+  | GroupContextExtensionsProposal
+  | ExternalProposal
+  deriving stock (Bounded, Enum)
+  deriving (ParseMLS) via (EnumMLS Word16 ProposalType)

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

@@ -17,13 +17,19 @@
 
 module Wire.API.MLS.Serialisation
   ( ParseMLS (..),
+    parseMLSVector,
+    parseMLSBytes,
     BinaryMLS (..),
-    EnumBinary (..),
+    EnumMLS (..),
+    safeToEnum,
     decodeMLS,
     decodeMLS',
+    decodeMLSWith,
+    decodeMLSWith',
   )
 where
 
+import Control.Applicative
 import Data.Binary
 import Data.Binary.Get
 import qualified Data.ByteString.Lazy as LBS
@@ -34,8 +40,30 @@ import Imports
 class ParseMLS a where
   parseMLS :: Get a
 
-instance ParseMLS ByteString where
-  parseMLS = get
+parseMLSVector :: forall w a. (Binary w, Integral w) => Get a -> Get [a]
+parseMLSVector getItem = do
+  len <- get @w
+  pos <- bytesRead
+  isolate (fromIntegral len) $ go (pos + fromIntegral len)
+  where
+    go :: Int64 -> Get [a]
+    go endPos = do
+      x <- getItem
+      pos <- bytesRead
+      (:) <$> pure x <*> if pos < endPos then go endPos else pure []
+
+parseMLSBytes :: forall w. (Binary w, Integral w) => Get ByteString
+parseMLSBytes = do
+  len <- fromIntegral <$> get @w
+  getByteString len
+
+instance ParseMLS Word8 where parseMLS = get
+
+instance ParseMLS Word16 where parseMLS = get
+
+instance ParseMLS Word32 where parseMLS = get
+
+instance ParseMLS Word64 where parseMLS = get
 
 -- | A wrapper to generate a 'ParseMLS' instance given a 'Binary' instance.
 newtype BinaryMLS a = BinaryMLS a
@@ -44,11 +72,15 @@ instance Binary a => ParseMLS (BinaryMLS a) where
   parseMLS = BinaryMLS <$> get
 
 -- | A wrapper to generate a 'Binary' instance for an enumerated type.
-newtype EnumBinary w a = EnumBinary {unEnumBinary :: a}
+newtype EnumMLS w a = EnumMLS {unEnumMLS :: a}
+
+safeToEnum :: forall a f. (Bounded a, Enum a, Alternative f) => Int -> f a
+safeToEnum n = guard (n >= fromEnum @a minBound && n <= fromEnum @a maxBound) $> toEnum n
 
-instance (Binary w, Integral w, Enum a) => Binary (EnumBinary w a) where
-  get = EnumBinary . toEnum . fromIntegral <$> get @w
-  put = put @w . fromIntegral . fromEnum . unEnumBinary
+instance (Binary w, Integral w, Bounded a, Enum a) => ParseMLS (EnumMLS w a) where
+  parseMLS = do
+    n <- fromIntegral <$> get @w
+    EnumMLS <$> safeToEnum n
 
 -- | Decode an MLS value from a lazy bytestring. Return an error message in case of failure.
 decodeMLS :: ParseMLS a => LByteString -> Either Text a
@@ -65,3 +97,6 @@ decodeMLSWith p b = case runGetOrFail p b of
   Right (remainder, pos, x)
     | LBS.null remainder -> Right x
     | otherwise -> Left $ "Trailing data at position " <> T.pack (show pos)
+
+decodeMLSWith' :: Get a -> ByteString -> Either Text a
+decodeMLSWith' p = decodeMLSWith p . LBS.fromStrict

libs/wire-api/wire-api.cabal

@@ -39,6 +39,7 @@ library
       Wire.API.MLS.CipherSuite
       Wire.API.MLS.Credential
       Wire.API.MLS.KeyPackage
+      Wire.API.MLS.Proposal
       Wire.API.MLS.Serialisation
       Wire.API.Notification
       Wire.API.Properties

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

@@ -15,11 +15,15 @@
 -- 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 Brig.API.MLS.KeyPackages where
+module Brig.API.MLS.KeyPackages
+  ( uploadKeyPackages,
+  )
+where
 
 import Brig.API.Error
 import Brig.API.Handler
 import qualified Brig.Data.MLS.KeyPackage as Data
+import Control.Applicative
 import Control.Monad.Trans.Except
 import Data.Id
 import Data.Json.Util
@@ -45,6 +49,7 @@ parseKeyPackage (fromBase64ByteString . kpData -> kpd) =
 validateKeyPackage :: ClientIdentity -> KeyPackage -> Handler ()
 validateKeyPackage identity (kpTBS -> kp) = do
   validateCredential identity (kpCredential kp)
+  validateExtensions (kpExtensions kp)
 
 validateCredential :: ClientIdentity -> Credential -> Handler ()
 validateCredential identity cred = do
@@ -54,3 +59,38 @@ validateCredential identity cred = do
       decodeMLS' (bcIdentity cred)
   when (identity /= identity') $
     throwE mlsIdentityMismatch
+
+data RequiredExtensions f = RequiredExtensions
+  { reLifetime :: f Lifetime,
+    reCapabilities :: f ()
+  }
+
+instance Alternative f => Semigroup (RequiredExtensions f) where
+  RequiredExtensions lt1 cap1 <> RequiredExtensions lt2 cap2 =
+    RequiredExtensions (lt1 <|> lt2) (cap1 <|> cap2)
+
+instance Alternative f => Monoid (RequiredExtensions f) where
+  mempty = RequiredExtensions empty empty
+
+checkRequiredExtensions :: Applicative f => RequiredExtensions f -> f (RequiredExtensions Identity)
+checkRequiredExtensions re =
+  RequiredExtensions
+    <$> (Identity <$> reLifetime re)
+    <*> (Identity <$> reCapabilities re)
+
+findExtensions :: [Extension] -> Maybe (RequiredExtensions Identity)
+findExtensions = checkRequiredExtensions . foldMap findExtension
+
+findExtension :: Extension -> RequiredExtensions Maybe
+findExtension e = case decodeExtension e of
+  Just (SomeExtension SLifetimeExtensionTag lt) -> RequiredExtensions (pure lt) Nothing
+  Just (SomeExtension SCapabilitiesExtensionTag _) -> RequiredExtensions Nothing (pure ())
+  _ -> mempty
+
+validateExtensions :: [Extension] -> Handler ()
+validateExtensions exts = do
+  _re <-
+    maybe (throwE (mlsProtocolError "Missing required extensions")) pure $
+      findExtensions exts
+  -- TODO: validate lifetime
+  pure ()