ServantAuthGoogle.hs

module ServantAuthGoogle where

import           ClassyPrelude                      hiding (Handler)
import           Control.Error.Util                 (hush)
import           Control.Lens                       ((.~), (^.), (^?))
import           Control.Monad.Except
import qualified Crypto.JOSE                        as Jose
import qualified Crypto.JWT                         as Jose
import           Crypto.Util                        (constTimeEq)
import           Data.Aeson
import qualified Data.ByteString                    as BS
import qualified Data.ByteString.Lazy               as LBS
import           Data.Function                      ((&))
import           Network.HTTP.Simple                hiding (Proxy)
import           Network.URI                        hiding (path)
import           Network.Wai                        (requestHeaders)
import           Servant.Auth.Server                hiding (audienceMatches,
                                                     key)
import           Servant.Auth.Server.Internal.Class (IsAuth (..))

import qualified Harbinger.Type.Email               as Email

--------------------------------------------------------------------------------
-- | Payload decoded from an @idToken@ JWT issued by Google's Identity service.
data GoogleAuth
  = GoogleAuth
  { gaEmail         :: !Email.EmailAddress
  , gaEmailVerified :: !(Maybe Bool)
  , gaName          :: !(Maybe Text)
  , gaPicture       :: !(Maybe Text)
  , gaGivenName     :: !(Maybe Text)
  , gaFamilyName    :: !(Maybe Text)
  , gaLocale        :: !(Maybe Text)
  , gaHostDomain    :: !(Maybe Text)
  } deriving (Generic, Show)

instance FromJSON GoogleAuth where
  parseJSON = withObject "GoogleAuth" $ \o -> do
    gaEmail      <- o .:  "email"
    gaName       <- o .:? "name"
    gaPicture    <- o .:? "picture"
    gaGivenName  <- o .:? "givenName"
    gaFamilyName <- o .:? "familyName"
    gaLocale     <- o .:? "locale"
    gaHostDomain <- o .:? "hd"

    e <- o .:? "emailVerified"
    let gaEmailVerified = toBool =<< e

    pure GoogleAuth{..}
    where
      toBool :: Text -> Maybe Bool
      toBool "true"  = Just True
      toBool "false" = Just False
      toBool _       = Nothing

instance FromJWT GoogleAuth where
  decodeJWT m =
    let claims = m ^. Jose.unregisteredClaims
    in case fromJSON (Object claims) of
      Error   e -> Left $ pack e
      Success a -> Right a

instance ToJSON GoogleAuth
instance ToJWT GoogleAuth

--------------------------------------------------------------------------------
-- | Newtype that wraps some function checking if a Google JWT Audience matches
-- | some string.
newtype GoogleJWTAudienceMatches
  = GoogleJWTAudienceMatches (Jose.StringOrURI -> IsMatch)

-- | Newtype that wraps some function checking if a Google JWT Issuer matches
-- | some string.
newtype GoogleJWTIssuerMatches
  = GoogleJWTIssuerMatches (Jose.StringOrURI -> IsMatch)

-- | Authentication settings for @idToken@ JWTs issued by Google's Identity
-- | service.
data GoogleAuthSettings
  = GoogleAuthSettings
  { gasKeySet          :: !Jose.JWKSet
  , gasAudienceMatches :: !GoogleJWTAudienceMatches
  , gasIssuerMatches   :: !GoogleJWTIssuerMatches
  , gasHostDomain      :: !(Maybe Text)
  }

-- | Check that a given @String@ matches a @Jose.StringOrURI@.
stringOrUriMatches :: String -> Jose.StringOrURI -> IsMatch
stringOrUriMatches inputString stringOrUri =
  fromMaybe DoesNotMatch $ tryUri <|> tryString
  where
    tryUri :: Maybe IsMatch
    tryUri = do
      uri <- stringOrUri ^? Jose.uri
      inputUri <- parseURI inputString
      guard $ inputUri == uri
      pure Matches

    tryString :: Maybe IsMatch
    tryString = do
      string <- stringOrUri ^? Jose.string
      guard $ inputString == string
      pure Matches

-- | Smart constructor for @GoogleAudienceMatches@ given some @String@.
googleAudienceMatches :: String -> GoogleJWTAudienceMatches
googleAudienceMatches = GoogleJWTAudienceMatches . stringOrUriMatches

-- | Smart constructor for @GoogleIssuerMatches@ given some @String@.
googleIssuerMatches :: String -> GoogleJWTIssuerMatches
googleIssuerMatches = GoogleJWTIssuerMatches . stringOrUriMatches

-- | An @AuthCheck@ for JWTs requiring @audienceMatches@ and @issuerMatches@
-- | predicates, in our case specifically for @idToken@ JWTs that come from
-- | Google's Identity service.
googleAuthCheck :: GoogleAuthSettings -> AuthCheck GoogleAuth
googleAuthCheck config = do
  req <- ask
  token <- maybe mempty return $ do
    authHdr <- lookup "Authorization" $ requestHeaders req
    let bearer = "Bearer "
        (mbearer, rest) = BS.splitAt (BS.length bearer) authHdr
    guard (mbearer `constTimeEq` bearer)
    pure rest

  verifiedJWT <- liftIO $ runExceptT $ do
    unverifiedJWT <- Jose.decodeCompact $ LBS.fromStrict token
    Jose.verifyClaims (googleAuthSettingsToJwtValidationSettings config)
                      (gasKeySet config)
                      unverifiedJWT

  let mVerifiedJWT = do
        verified <- hush (verifiedJWT :: Either Jose.JWTError Jose.ClaimsSet)
        decoded  <- hush $ decodeJWT verified :: Maybe GoogleAuth
        if (isNothing $ gasHostDomain config)
        then Just decoded
        else do
          hd <- gasHostDomain config
          jd <- gaHostDomain decoded
          if (hd == jd) then (Just decoded) else Nothing
  maybe mzero pure mVerifiedJWT

-- | Slightly modified version of @jwtSettingsToJwtValidationSettings@ that
-- | verifies both @audienceMatches@ and @issuerMatches@ predicates.
googleAuthSettingsToJwtValidationSettings
  :: GoogleAuthSettings
  -> Jose.JWTValidationSettings
googleAuthSettingsToJwtValidationSettings gas =
  let (GoogleJWTAudienceMatches gam) = gasAudienceMatches gas
      (GoogleJWTIssuerMatches gim) = gasIssuerMatches gas
      validationSettings = Jose.defaultJWTValidationSettings (toBool <$> gam)
  in  validationSettings & Jose.issuerPredicate .~ (toBool <$> gim)
  where
    toBool Matches      = True
    toBool DoesNotMatch = False

instance IsAuth GoogleAuth GoogleAuth where
  type AuthArgs GoogleAuth = '[GoogleAuthSettings]
  runAuth _ _ = googleAuthCheck

--------------------------------------------------------------------------------
getJWKSet :: ByteString -> ByteString -> IO (Maybe Jose.JWKSet)
getJWKSet host path = do
  let req =
        setRequestMethod "GET"
        . setRequestSecure True
        . setRequestHost host
        . setRequestPath path
        . setRequestPort 443
        $ defaultRequest
  hush . getResponseBody <$> httpJSONEither req

ServantAuthMain.hs

module ServantAuthMain.hs where

import           ClassyPrelude                        hiding (keys, killThread)
import           Control.Concurrent                   (killThread)
import           Control.Monad.Metrics                (initializeWith)
import           Database.Persist.Postgresql          (ConnectionPool,
                                                       ConnectionString,
                                                       createPostgresqlPool)
import           Harbinger.Config
import           Harbinger.Logger
import qualified Harbinger.Server                     as Server
import           Harbinger.Type.App
import           Harbinger.Type.AppEnv
import           Harbinger.Type.Config
import           Harbinger.Type.ServerSettings
import           Network.Wai.Handler.Warp             (run)
import qualified Network.Wai.Metrics                  as NWM
import           Network.Wai.Middleware.RequestLogger
import           Servant                              (hoistServerWithContext,
                                                       serveWithContext)

import           System.IO                            (BufferMode(..), hSetBuffering)
import           System.Remote.Monitoring.Wai         (forkServer,
                                                       serverMetricStore,
                                                       serverThreadId)

--------------------------------------------------------------------------------
-- | Initialize the application and serve the API.
main :: IO ()
main = do
  hSetBuffering stdout NoBuffering
  config@Config {..} <- getConfig
  print config
  logState@LogState{..} <- mkLogState cServerEnv
  connPool <- mkConnectionPool config logState
  ServerSettings {..} <- mkServerSettings config

  bracket
    (forkServer "localhost" 8081)
    (killThread . serverThreadId) $
    \metricServer -> do
      -- Set up EKG metrics store and WAI metrics logging
      let metricStore = serverMetricStore metricServer
      metrics <- initializeWith metricStore
      waiMetrics <- NWM.registerWaiMetrics metricStore

      -- Set up WAI request logging middlware
      let waiLogger = case cServerEnv of
            Development -> logStdoutDev
            Testing     -> id
            _           -> logStdout

          middleware = (NWM.metrics waiMetrics) . waiLogger

          appEnv = AppEnv metrics connPool lsContexts lsNamespace lsLogEnv

          handler =
            hoistServerWithContext
              Server.routes
              ssAuthContext
              (appToHandler appEnv)
              Server.handler

          server =
            serveWithContext
            Server.routes
            ssAuthConfig
            handler

      run cServerPort . middleware $ server

--------------------------------------------------------------------------------
-- | Record of the information @Katip@ needs.
data LogState
  = LogState
  { lsContexts  :: !LogContexts
  , lsNamespace :: !Namespace
  , lsLogEnv    :: !LogEnv
  }

-- | Make a @Katip@ @LogState@ based on the application's environment.
mkLogState :: ServerEnv -> IO LogState
mkLogState env = do
  let katipEnv = Environment $ tshow env
  let logLevel =
        case env of
          Production -> InfoS
          _          -> DebugS
  handleScribe <- mkHandleScribe ColorIfTerminal stdout logLevel V2
  let mkLogEnv = initLogEnv mempty katipEnv
  logEnv <-
    registerScribe "stdout" handleScribe defaultScribeSettings =<< mkLogEnv
  pure $ LogState mempty mempty logEnv

-- | Make a PostgreSQL @ConnectionPool@ based on the @Config@ record; log all
-- | associated database actions out using a @Katip@ logger with the provided
-- | @LogState@.
mkConnectionPool :: Config -> LogState -> IO ConnectionPool
mkConnectionPool config@Config {..} LogState{..} = do
  let connStr = mkConnectionString config
      logLoudness =
        case cServerEnv of
          Testing -> katipNoLogging
          _       -> id
  runKatipContextT lsLogEnv lsContexts lsNamespace
    $ logLoudness
    $ katipAddNamespace "psql"
    $ createPostgresqlPool connStr cPGPoolSize

-- | Make a PostgreSQL @ConnectionString@ based on the @Config@ record.
mkConnectionString :: Config -> ConnectionString
mkConnectionString Config {..} =
  let keys = ["host=", "port=", "user=", "password=", "dbname="]
      params = [cPGHost, tshow cPGPort, cPGUser, cPGPassword, cPGDatabase]
      zipped = zipWith (<>) keys (map encodeUtf8 params)
  in intercalate " " zipped

ServantAuthSettings.hs

module Harbinger.Type.ServerSettings where

import ClassyPrelude hiding (keys)
import Servant
import Servant.Auth.Server

import Harbinger.Type.Config
import Harbinger.Util.GoogleAuth

--------------------------------------------------------------------------------
-- | The @servant-auth@ plugin settings needed to for the route handlers.
type ServerAuthSettings = '[GoogleAuthSettings, JWTSettings, CookieSettings]

-- | The term-level representation of @servant-auth@'s plugin settings.
type ServerAuthConfig = Context ServerAuthSettings

-- | The type-level representation of @servant-auth@'s plugin settings.
type ServerAuthContext = Proxy ServerAuthSettings

-- | Settings needed to run Harbinger's web api server.
data ServerSettings
  = ServerSettings
  { ssAuthConfig  :: ServerAuthConfig
  , ssAuthContext :: ServerAuthContext
  }

-- | An error that can be encountered when trying to retrieve a @JWKSet@ from
-- | Google.
data GoogleJWKSetRetrievalError
  = GoogleJWKSetRetrievalError
  deriving (Show, Typeable)
instance Exception GoogleJWKSetRetrievalError

mkServerSettings :: Config -> IO ServerSettings
mkServerSettings Config{..} = do
  let gasAudienceMatches = googleAudienceMatches $ unpack cGoogleJWTAudience
  let gasIssuerMatches = googleIssuerMatches $ unpack cGoogleJWTIssuer
  let gasHostDomain = cGoogleHostDomain
  mKeySet <- getJWKSet "www.googleapis.com" "/oauth2/v3/certs"
  gasKeySet <-
    case mKeySet of
      Nothing -> throwM GoogleJWKSetRetrievalError
      Just jwkSet -> pure jwkSet

  let googleAuthSettings = GoogleAuthSettings{..}

  -- NOTE: @servant-auth@ requires this, even though it should be unnecessary.
  servantAuthJWTSettings <- fmap defaultJWTSettings generateKey

  let ssAuthContext = Proxy :: ServerAuthContext
  let ssAuthConfig =
        googleAuthSettings
        :. servantAuthJWTSettings
        :. defaultCookieSettings
        :. EmptyContext

  pure ServerSettings{..}