6 files changed, 60 insertions, 17 deletions

diff --git a/overlays/spm/frontend/src/app/spm/spm.component.ts b/overlays/spm/frontend/src/app/spm/spm.component.ts
index f6892d7c..361a1ce4 100644
--- a/overlays/spm/frontend/src/app/spm/spm.component.ts
+++ b/overlays/spm/frontend/src/app/spm/spm.component.ts
@@ -36,7 +36,7 @@ export interface ClaimedSpmMail {
 }
 
 export interface SpmJwtUnregisteredPayload {
-  local: string;
+  'li.yggdrasil.local': string;
 }
 
 type SpmJwtPayload = JwtPayload & SpmJwtUnregisteredPayload;
@@ -136,7 +136,7 @@ export class SpmComponent implements OnInit, OnDestroy {
       const curr: Map<string, SpmMail> = this.spmMails$.getValue();
       curr.set(k, {
         state: 'loaded',
-        local: payload.local,
+        local: payload['li.yggdrasil.local'],
         domain: payload.aud,
         jwt: encoded,
         expiration: payload.exp ? new Date(1000 * payload.exp) : undefined,
diff --git a/overlays/spm/lib/Spm/Api.hs b/overlays/spm/lib/Spm/Api.hs
index 04dff2c9..8285cc55 100644
--- a/overlays/spm/lib/Spm/Api.hs
+++ b/overlays/spm/lib/Spm/Api.hs
@@ -5,6 +5,7 @@ module Spm.Api
   , SpmMailbox, SpmDomain
   , SpmLocal(..), SpmExtension(..)
   , SpmMapping(..), SpmMappingState(..), SpmMappingListingItem(..), SpmMappingListing(..)
+  , SpmJWTClaims(..), _spmjwtLocal, _SpmJWTLocal
   , _SpmMappingText, _SpmMappingStateReject
   , spmMappingAncestors
   , SpmApi, spmApi
@@ -24,18 +25,21 @@ import Type.Reflection (Typeable)
 
 import Control.Lens
 
+import Control.Monad (guard)
+
 import Data.CaseInsensitive (CI)
 import qualified Data.CaseInsensitive as CI
 import Data.CaseInsensitive.Instances ()
 
 import Crypto.JOSE.JWK (JWKSet)
-import Crypto.JWT (SignedJWT)
+import Crypto.JWT (SignedJWT, ClaimsSet, HasClaimsSet(..), emptyClaimsSet)
 import Crypto.JWT.Instances ()
 
 import Data.UUID (UUID)
 import Data.UUID.Instances ()
 
 import qualified Data.Aeson as JSON
+import qualified Data.Aeson.Lens as JSON
 import Data.Aeson.TH (deriveJSON)
 import Data.Aeson.Casing
 
@@ -163,6 +167,31 @@ deriveJSON (aesonPrefix trainCase) ''SpmMappingListingItem
 instance ToJSON SpmMappingListing where
   toJSON SpmMappingListing{..} = JSON.object [ "mappings" JSON..= unSpmMappingListing ]
 
+data SpmJWTClaims = SpmJWTClaims
+  { spmjwtStdClaims :: ClaimsSet
+  , spmjwtLocal :: SpmLocal
+  } deriving stock (Eq, Show, Generic, Typeable)
+
+makeLensesFor [("spmjwtStdClaims", "_stdClaims"), ("spmjwtLocal", "_spmjwtLocal")] ''SpmJWTClaims
+
+instance HasClaimsSet SpmJWTClaims where
+  claimsSet = _stdClaims
+
+_SpmJWTLocal :: (Wrapped l, Unwrapped l ~ Unwrapped SpmLocal) => Prism' SpmJWTClaims l
+_SpmJWTLocal = prism' toClaims fromClaims
+  where toClaims (view $ _Wrapped' . _Unwrapped' -> spmjwtLocal) = SpmJWTClaims{..}
+          where spmjwtStdClaims = emptyClaimsSet
+        fromClaims SpmJWTClaims{..} = view (_Wrapped' . _Unwrapped') spmjwtLocal <$ guard (spmjwtStdClaims == emptyClaimsSet)
+
+instance JSON.ToJSON SpmJWTClaims where
+  toJSON SpmJWTClaims{..} = JSON.toJSON spmjwtStdClaims
+    & JSON._Object . at "li.yggdrasil.local" .~ Just (JSON.toJSON spmjwtLocal)
+
+instance JSON.FromJSON SpmJWTClaims where
+  parseJSON = JSON.withObject "SpmJWTClaims" $ \o -> SpmJWTClaims
+    <$> JSON.parseJSON (JSON._Object # o)
+    <*> o JSON..: "li.yggdrasil.local"
+
 
 type SpmApi = "whoami" :> Get '[PlainText, JSON] SpmMailbox
          :<|> "domain" :> Get '[PlainText, JSON] SpmDomain
diff --git a/overlays/spm/package.yaml b/overlays/spm/package.yaml
index e188dc77..6d4abe1c 100644
--- a/overlays/spm/package.yaml
+++ b/overlays/spm/package.yaml
@@ -1,5 +1,5 @@
 name: spm
-version: 0.1.0
+version: 0.1.1
 
 default-extensions:
   - NoImplicitPrelude
@@ -36,6 +36,7 @@ library:
     - lens
     - case-insensitive
     - aeson
+    - lens-aeson
     - jose
     - uuid
     - containers
diff --git a/overlays/spm/server/Crypto/JOSE/Error/Instances.hs b/overlays/spm/server/Crypto/JOSE/Error/Instances.hs
new file mode 100644
index 00000000..9f3eea51
--- /dev/null
+++ b/overlays/spm/server/Crypto/JOSE/Error/Instances.hs
@@ -0,0 +1,11 @@
+{-# OPTIONS_GHC -Wno-orphans #-}
+
+module Crypto.JOSE.Error.Instances () where
+
+import Data.Monoid (Monoid)
+import Control.Monad (Monad)
+import Control.Applicative (Alternative)
+
+import Crypto.JOSE.Error
+
+deriving newtype instance (Monad m, Monoid e) => Alternative (JOSE e m)
diff --git a/overlays/spm/server/Spm/Server.hs b/overlays/spm/server/Spm/Server.hs
index df12e3d8..8e7f8786 100644
--- a/overlays/spm/server/Spm/Server.hs
+++ b/overlays/spm/server/Spm/Server.hs
@@ -44,7 +44,7 @@ import Database.Persist
 import Database.Persist.Postgresql
 import UnliftIO.Pool
 
-import Control.Monad.Trans.Reader (ReaderT, runReaderT)
+import Control.Monad.Trans.Reader (ReaderT, runReaderT, mapReaderT)
 
 import Control.Monad.Logger
 
@@ -58,12 +58,12 @@ import qualified Data.UUID as UUID
 import qualified Data.UUID.V4 as UUID
 
 import qualified Data.Aeson as JSON
-import Data.Aeson.Lens (_JSON)
 
 import System.FilePath ((</>), isRelative)
 
 import Crypto.JOSE.JWK hiding (Context)
 import Crypto.JOSE.JWK.Instances ()
+import Crypto.JOSE.Error.Instances ()
 
 import Crypto.Random.Instances ()
 import qualified Crypto.Random as Crypto
@@ -197,10 +197,13 @@ mkSpmApp = do
 spmSql :: ReaderT SqlBackend Handler' a -> Handler' a
 spmSql act = do
   sqlPool <- view sctxSqlPool
-  withResource sqlPool $ runReaderT act
+  mapReaderT (mapLoggingT $ either throwError pure <=< liftIO) . withResource sqlPool $ mapReaderT (mapLoggingT runHandler) . runReaderT act
 
-spmJWT :: forall error a. Show error => ServerError -> ExceptT error IO a -> Handler' a
-spmJWT errTemplate = either (\err -> throwError errTemplate{ errBody = LBS.fromStrict . Text.encodeUtf8 . Text.pack $ show err }) return <=< liftIO . runExceptT
+spmJWT :: forall error a. Show error => ServerError -> JOSE error IO a -> Handler' a
+spmJWT errTemplate = either (\err -> throwError errTemplate{ errBody = LBS.fromStrict . Text.encodeUtf8 . Text.pack $ show err }) return <=< liftIO . runJOSE
+
+withJOSE :: forall m e e' a. Functor m => (e -> e') -> JOSE e m a -> JOSE e' m a
+withJOSE f = JOSE . withExceptT f . unwrapJOSE
 
 generateLocal :: MonadIO m => SpmStyle -> m MailLocal
 generateLocal SpmWords = fmap (review _Wrapped . CI.mk) . liftIO $ do
@@ -254,24 +257,23 @@ spmServer dom mbox = whoami
       instanceId' <- view sctxInstanceId
       jwks <- view $ sctxJwkSet . _Wrapped
       tokenId <- liftIO UUID.nextRandom
-      let claimsSet = emptyClaimsSet
+      let jwtClaims = (_SpmJWTLocal # local)
             & claimIss ?~ (JWT.string # UUID.toText instanceId')
             & claimAud ?~ Audience (pure $ dom ^. _Wrapped . to CI.original . re JWT.string)
             & claimNbf ?~ NumericDate t
             & claimIat ?~ NumericDate t
             & claimExp ?~ NumericDate (600 `addUTCTime` t)
             & claimJti ?~ UUID.toText tokenId
-            & unregisteredClaims . at "local" ?~ view (_Wrapped . re _JSON) local
       spmJWT @JWT.Error err500 $ do
-        (jwsAlg, selectedJwk) <- withExceptT (fromMaybe JWT.NoUsableKeys . getFirst) . asum $ map (\jwk' -> (, jwk') <$> withExceptT (First . Just) (bestJWSAlg jwk')) jwks
-        signClaims selectedJwk (newJWSHeader ((), jwsAlg)) claimsSet
+        (jwsAlg, selectedJwk) <- withJOSE (fromMaybe JWT.NoUsableKeys . getFirst) . asum $ map (\jwk' -> (, jwk') <$> withJOSE (First . Just) (bestJWSAlg jwk')) jwks
+        signJWT selectedJwk (newJWSHeader ((), jwsAlg)) jwtClaims
 
     claim jwt = do
       jwks <- view sctxJwkSet
       let validationSettings' = defaultJWTValidationSettings ((== Just dom) . fmap (review _Wrapped . CI.mk) . preview JWT.string)
             & jwtValidationSettingsAllowedSkew .~ 5
-      claims <- spmJWT @JWT.JWTError err403 $ verifyClaims validationSettings' jwks jwt
-      mailboxMappingLocal <- maybe (throwError err400{ errBody = "Claim ‘local’ missing" }) (return . Just) $ claims ^? unregisteredClaims . ix "local" . _JSON . to CI.mk . re _Wrapped
+      jwtClaims <- spmJWT @JWT.JWTError err403 $ verifyJWT validationSettings' jwks jwt
+      let mailboxMappingLocal = Just $ jwtClaims ^. _spmjwtLocal . _Wrapped . _Unwrapped
 
       spmSql $ do
         Entity mailboxMappingMailbox _ <- maybe (throwError err404) return <=< getBy $ UniqueMailbox mbox
diff --git a/overlays/spm/spm.nix b/overlays/spm/spm.nix
index 7c9eb933..2d1f7eed 100644
--- a/overlays/spm/spm.nix
+++ b/overlays/spm/spm.nix
@@ -10,13 +10,13 @@
 }:
 mkDerivation {
   pname = "spm";
-  version = "0.1.0";
+  version = "0.1.1";
   src = ./.;
   isLibrary = true;
   isExecutable = true;
   libraryHaskellDepends = [
     aeson aeson-casing base case-insensitive containers jose lens
-    servant text uuid
+    lens-aeson servant text uuid
   ];
   libraryToolDepends = [ hpack ];
   executableHaskellDepends = [