serve.hs

#!/usr/bin/env nix-shell
#!nix-shell -i runhaskell -p "haskellPackages.ghcWithPackages (pkgs: with pkgs; [servant-server servant-generic servant-auth-server lens resource-pool sqlite-simple string-conv interpolatedstring-perl6 ])"
#!nix-shell -I /home/ielectric/dev


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}


import Network.Wai
import           Network.Wai.Handler.Warp (defaultSettings, runSettings,
                                           setPort)
import Servant
import Servant.Server
import Data.Text
import Data.Monoid
import Servant.Generic
import Network.Wai.Handler.Warp (Port)
import Data.Pool
import Database.SQLite.Simple as Sqlite hiding ((:.))
import Control.Monad.Reader
import Control.Lens
import Data.String.Conv
import Servant.Auth
import Servant.Auth.Server
import Control.Exception.Lifted hiding (Handler)
import Data.ByteString as BS
import System.IO              (stdout)
import Safe
import Data.Aeson

data Env = Env
  { _envDbConn :: Sqlite.Connection
  }
$(makeLensesWith abbreviatedFields ''Env)

type AppM = ReaderT Env IO

data AppError = Redirect302 BS.ByteString
              | OAuthError Text
              | ValidationError Text
              deriving (Eq, Show)
instance Exception AppError

getConn :: AppM Sqlite.Connection
getConn = _envDbConn <$> ask

newtype UserId  = UserId Int deriving (Eq, Show, Generic, ToJSON, FromJSON)
instance FromJWT UserId
instance ToJWT UserId

handleAppError :: AppError -> Handler a
handleAppError e = case e of
  Redirect302 u -> throwError $ err302{errHeaders=("Location", u):(errHeaders err302)}
  ValidationError x -> throwError $ err400{errBody=(toS x)}
  x@(OAuthError _) -> throwError $ err500{errBody=(toS $ show x)}

data Routes route = Routes
  { rTest :: route :- "test" :> Get '[PlainText] Text
  , rLogin :: route :- Auth '[Cookie] UserId :> "login" :> Get '[PlainText] (Headers '[Header "Set-Cookie" SetCookie, Header "Set-Cookie" SetCookie] Text)
  } deriving (Generic)

server :: (CookieSettings, JWTSettings) -> Routes (AsServerT (AppM))
server cookieConfig = Routes
  { rTest = pure "test works"
  , rLogin = login cookieConfig
  }

type RoutesType = ToServant (Routes AsApi)

routesApi :: Proxy RoutesType
routesApi = Proxy

type Config = ()

main :: IO ()
main  = do
  dbPool <- createPool
    (Sqlite.open "chaudhary.db")
    Sqlite.close
    1
    (fromInteger 5)
    5
  cookieKey <- generateKey
  let jwtCfg = defaultJWTSettings cookieKey
      cookieCfg = defaultCookieSettings{cookieIsSecure=Servant.Auth.Server.NotSecure}
      cfg =  cookieCfg :. jwtCfg :. EmptyContext
      appToHandler :: AppM a -> Handler a
      appToHandler action = do
        liftIO $ withResource dbPool $ \conn ->
          runReaderT action Env { _envDbConn=conn
                                }
        `catch` handleAppError
        `catch` (\ (e ::SomeException) -> Servant.throwError $ err500{errBody=(toS $ show e)})

  print "serving"
  runSettings (setPort 8000 defaultSettings) $
    serveWithContext routesApi cfg $
    hoistServerWithContext routesApi (Proxy :: Proxy '[CookieSettings, JWTSettings]) appToHandler (toServant $ server (cookieCfg, jwtCfg))

login :: (CookieSettings, JWTSettings) -> AuthResult UserId -> AppM (Headers '[Header "Set-Cookie" SetCookie, Header "Set-Cookie" SetCookie] Text)
login (cookieCfg, jwtCfg) authResult = case authResult of
  Authenticated userId_ -> do
    -- TODO: Log-out the user
    pure $ noHeader $ noHeader $ "Found auth user " <> toS ( show userId_) <> ", who should now be logged out"
  e -> do
    liftIO $ acceptLogin cookieCfg jwtCfg (UserId 12) >>= \case
      Nothing -> pure $ noHeader $ noHeader "Unable to create the session cookie. Something went wrong."
      Just applyCookies -> pure $ applyCookies "No user faund due to {e}. User ID 12 should be automatically logged in now."