From 0c668253a8a8748a9d1c460e3d9578033773dacd Mon Sep 17 00:00:00 2001
From: MrMelon54 <github@mrmelon54.com>
Date: Sat, 10 Feb 2024 02:53:58 +0000
Subject: [PATCH] Fix correct redirection after login flow, and update test
 client

---
 pages/login-memory.go.html |  1 +
 pages/login.go.html        |  1 +
 server/login.go            |  7 ++++++-
 server/server.go           |  3 ++-
 test-client/index.html     | 15 ++++-----------
 test-client/pop2.js        | 11 ++++++-----
 6 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/pages/login-memory.go.html b/pages/login-memory.go.html
index a566467..395eb60 100644
--- a/pages/login-memory.go.html
+++ b/pages/login-memory.go.html
@@ -17,6 +17,7 @@
     </div>
     <div>
         <form method="POST" action="/login">
+            <input type="hidden" name="redirect" value="{{.Redirect}}"/>
             <input type="hidden" name="loginname" value="{{.LoginName}}"/>
             <button type="submit">Continue</button>
         </form>
diff --git a/pages/login.go.html b/pages/login.go.html
index 99380de..134b2cc 100644
--- a/pages/login.go.html
+++ b/pages/login.go.html
@@ -10,6 +10,7 @@
 </header>
 <main>
     <form method="POST" action="/login">
+        <input type="hidden" name="redirect" value="{{.Redirect}}"/>
         <div>
             <label for="field_loginname">Login Name:</label>
             <input type="text" name="loginname" id="field_loginname" required/>
diff --git a/server/login.go b/server/login.go
index 653c85d..5d5cd68 100644
--- a/server/login.go
+++ b/server/login.go
@@ -33,11 +33,13 @@ func (h *HttpServer) loginGet(rw http.ResponseWriter, req *http.Request, _ httpr
 		pages.RenderPageTemplate(rw, "login-memory", map[string]any{
 			"ServiceName": h.conf.ServiceName,
 			"LoginName":   cookie.Value,
+			"Redirect":    req.URL.Query().Get("redirect"),
 		})
 		return
 	}
 	pages.RenderPageTemplate(rw, "login", map[string]any{
 		"ServiceName": h.conf.ServiceName,
+		"Redirect":    req.URL.Query().Get("redirect"),
 	})
 }
 
@@ -85,7 +87,7 @@ func (h *HttpServer) loginPost(rw http.ResponseWriter, req *http.Request, _ http
 
 	// save state for use later
 	state := login.Config.Namespace + ":" + uuid.NewString()
-	h.flowState.Set(state, flowStateData{login}, time.Now().Add(15*time.Minute))
+	h.flowState.Set(state, flowStateData{login, req.PostFormValue("redirect")}, time.Now().Add(15*time.Minute))
 
 	// generate oauth2 config and redirect to authorize URL
 	oa2conf := login.OAuth2Config
@@ -135,6 +137,9 @@ func (h *HttpServer) loginCallback(rw http.ResponseWriter, req *http.Request, _
 		http.Error(rw, "Internal Server Error", http.StatusInternalServerError)
 		return
 	}
+	if flowState.redirect != "" {
+		req.Form.Set("redirect", flowState.redirect)
+	}
 	h.SafeRedirect(rw, req)
 }
 
diff --git a/server/server.go b/server/server.go
index 08f2b8e..75bf0f2 100644
--- a/server/server.go
+++ b/server/server.go
@@ -41,7 +41,8 @@ type HttpServer struct {
 }
 
 type flowStateData struct {
-	sso *issuer.WellKnownOIDC
+	sso      *issuer.WellKnownOIDC
+	redirect string
 }
 
 func NewHttpServer(conf Conf, db *database.DB, signingKey mjwt.Signer) *http.Server {
diff --git a/test-client/index.html b/test-client/index.html
index d57203a..955835b 100644
--- a/test-client/index.html
+++ b/test-client/index.html
@@ -6,24 +6,17 @@
   <script>
       const ssoService = "http://localhost:9090";
 
-      POP2.init(ssoService + "/authorize", "bc36b32c-83cb-404e-8736-cb207f30afe3", "openid profile", 500, 600);
+      POP2.init(ssoService + "/authorize", "f4cdb93d-fe28-427b-b037-f03f44c86a16", "openid profile", 500, 600);
 
       function updateTokenInfo(data) {
           document.getElementById("someTextArea").textContent = JSON.stringify(data, null, 2);
       }
 
-      function parseJwt(token) {
-          const base64Url = token.split('.')[1];
-          const base64 = base64Url.replace(/-/g, '+').replace(/_/g, '/');
-          const jsonPayload = decodeURIComponent(window.atob(base64).split('').map(function (c) {
-              return '%' + ('00' + c.charCodeAt(0).toString(16)).slice(-2);
-          }).join(''));
-          return JSON.parse(jsonPayload);
-      }
-
       function doThisThing() {
           POP2.clientRequest(ssoService + "/userinfo", {}, true).then(function (x) {
-              console.log(x);
+              return x.json();
+          }).then(function (x) {
+              updateTokenInfo(x);
           }).catch(function (x) {
               console.error(x);
           });
diff --git a/test-client/pop2.js b/test-client/pop2.js
index a233715..52ab6a6 100644
--- a/test-client/pop2.js
+++ b/test-client/pop2.js
@@ -139,20 +139,21 @@
                 options.headers['Authorization'] = 'Bearer ' + access_token;
                 return new Promise(function (res, rej) {
                     fetch(resource, options).then(function (x) {
-                        if (x.statusCode >= 200 && x.statusCode < 300) res(x);
+                        if (x.status >= 200 && x.status < 300) res(x);
                         else rej(x);
                     }).catch(function (x) {
-                        rej(x);
+                        rej(["failed to send request", x]);
                     });
                 });
             };
-            const resendRequest = function() {
+            const resendRequest = function () {
                 return new Promise(function (res, rej) {
-                    w.POP2.getToken(function() {
+                    access_token = undefined;
+                    w.POP2.getToken(function () {
                         sendRequest().then(function (x) {
                             res(x);
                         }).catch(function (x) {
-                            rej(x);
+                            rej(["failed to resend request", x]);
                         });
                     });
                 });