From fc77ea3e22d179644bc8251ad06ed1ea9d03e19d Mon Sep 17 00:00:00 2001
From: David Mosbach <david.mosbach@live.de>
Date: Sat, 23 Dec 2023 04:53:56 +0100
Subject: [PATCH] redirecting auth code back to client

---
 oauth2-mock-server.cabal |  3 +++
 package.yaml             |  1 +
 src/AuthCode.hs          |  2 +-
 src/Server.hs            | 21 ++++++++++++++++++---
 4 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/oauth2-mock-server.cabal b/oauth2-mock-server.cabal
index 83ad5af..25276d8 100644
--- a/oauth2-mock-server.cabal
+++ b/oauth2-mock-server.cabal
@@ -32,6 +32,7 @@ library
   build-depends:
       aeson
     , base >=4.7 && <5
+    , bytestring
     , containers
     , http-client
     , servant
@@ -56,6 +57,7 @@ executable oauth2-mock-server-exe
   build-depends:
       aeson
     , base >=4.7 && <5
+    , bytestring
     , containers
     , http-client
     , oauth2-mock-server
@@ -82,6 +84,7 @@ test-suite oauth2-mock-server-test
   build-depends:
       aeson
     , base >=4.7 && <5
+    , bytestring
     , containers
     , http-client
     , oauth2-mock-server
diff --git a/package.yaml b/package.yaml
index 47736b4..4b1c903 100644
--- a/package.yaml
+++ b/package.yaml
@@ -32,6 +32,7 @@ dependencies:
 - stm
 - time
 - transformers
+- bytestring
 
 ghc-options:
 - -Wall
diff --git a/src/AuthCode.hs b/src/AuthCode.hs
index cf373a6..4324e36 100644
--- a/src/AuthCode.hs
+++ b/src/AuthCode.hs
@@ -33,7 +33,7 @@ genUnencryptedCode client url expiration state = do
   now <- getCurrentTime
   let
     expiresAt = expiration `addUTCTime` now
-    simpleCode = client <> url <> show now <> show expiresAt
+    simpleCode = filter (/= ' ') $ client <> url <> show now <> show expiresAt
   success <- atomically . stateTVar state $ \s ->
     let mEntry = M.lookup simpleCode s.activeCodes
     in
diff --git a/src/Server.hs b/src/Server.hs
index 0a477a9..f3990ee 100644
--- a/src/Server.hs
+++ b/src/Server.hs
@@ -17,8 +17,11 @@ import Control.Monad.IO.Class
 import Control.Monad.Trans.Reader
 
 import Data.Aeson
-import Data.List (find)
-import Data.Text hiding (elem, find, head, map, words)
+import Data.ByteString (ByteString (..))
+import Data.List (find, elemIndex)
+import Data.Maybe (fromMaybe)
+import Data.String (IsString (..))
+import Data.Text hiding (elem, find, head, length, map, null, splitAt, tail, words)
 import Data.Text.Encoding (decodeUtf8)
 
 import qualified Data.Map.Strict as Map
@@ -97,7 +100,19 @@ authServer = handleAuth
       mAuthCode <- asks (genUnencryptedCode client url 600) >>= liftIO
       liftIO $ print mAuthCode
       liftIO . putStrLn $ "user: " ++ show u ++ " | scopes: " ++ show (map (showScope @user @userData) scopes')
-      return uData
+      -- return uData
+      redirect $ url `withCode` mAuthCode
+    redirect :: Maybe ByteString -> AuthHandler userData
+    redirect (Just url) = throwError err302 { errHeaders = [("Location", url)]}
+    redirect Nothing = throwError err500 { errBody = "Could not generate authorisation code."}
+    withCode :: String -> Maybe String -> Maybe ByteString
+    withCode url Nothing = Nothing
+    withCode url (Just code) =
+      let qPos = fromMaybe (length url) $ elemIndex '?' url
+          (pre, post) = splitAt qPos url
+          post' = if not (null post) then '&' : tail post else post
+      in Just . fromString $ pre ++ "?authorization_code=" ++ code ++ post'
+
 
 exampleAuthServer :: AuthServer (Auth User (Map.Map Text Text))
 exampleAuthServer = authServer