1 files changed, 13 insertions, 11 deletions

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