From 199194b3526669f6581e0d07dab9947745eefcfc Mon Sep 17 00:00:00 2001
From: Milo Schwartz <mschwartz10612@gmail.com>
Date: Sat, 16 Nov 2024 22:42:09 -0500
Subject: [PATCH] initial implementation of new reource auth method middleware

---
 main.go | 96 +++++++++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 83 insertions(+), 13 deletions(-)

diff --git a/main.go b/main.go
index c81874f..547168e 100644
--- a/main.go
+++ b/main.go
@@ -1,19 +1,42 @@
 package badger
 
 import (
+	"bytes"
 	"context"
+	"encoding/json"
 	"fmt"
 	"net/http"
 	"net/url"
 )
 
-const SessionCookieName = "session"
+const AppSSOSessionCookieName = "session"
+const ResourceSessionCookieName = "resource_session"
 
 type Config struct {
 	AppBaseUrl string `json:"appBaseUrl"`
 	APIBaseUrl string `json:"apiBaseUrl"`
 }
 
+type CookieData struct {
+	Session         *string `json:"session"`
+	ResourceSession *string `json:"resource_session"`
+}
+
+type VerifyBody struct {
+	Cookies            CookieData `json:"cookies"`
+	OriginalRequestURL string     `json:"originalRequestURL"`
+	RequestScheme      *string    `json:"scheme"`
+	RequestHost        *string    `json:"host"`
+	RequestPath        *string    `json:"path"`
+	RequestMethod      *string    `json:"method"`
+	TLS                bool       `json:"tls"`
+}
+
+type VerifyResponse struct {
+	Valid       bool    `json:"valid"`
+	RedirectURL *string `json:"redirectUrl"`
+}
+
 func CreateConfig() *Config {
 	return &Config{}
 }
@@ -35,29 +58,76 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
 }
 
 func (p *Badger) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
-	cookie, err := req.Cookie(SessionCookieName)
+	verifyURL := fmt.Sprintf("%s/badger/verify-session", p.apiBaseUrl)
+	cookies := extractCookies(req)
+
+	originalRequestURL := url.QueryEscape(fmt.Sprintf("%s://%s%s", p.getScheme(req), req.Host, req.URL.RequestURI()))
+
+	cookieData := VerifyBody{
+		Cookies: CookieData{
+			Session:         cookies.Session,
+			ResourceSession: cookies.ResourceSession,
+		},
+		OriginalRequestURL: originalRequestURL,
+		RequestScheme:      &req.URL.Scheme,
+		RequestHost:        &req.Host,
+		RequestPath:        &req.URL.Path,
+		RequestMethod:      &req.Method,
+		TLS:                req.TLS != nil,
+	}
+
+	jsonData, err := json.Marshal(cookieData)
 	if err != nil {
-		originalRequestURL := url.QueryEscape(fmt.Sprintf("%s://%s%s", p.getScheme(req), req.Host, req.URL.RequestURI()))
-		http.Redirect(rw, req, fmt.Sprintf("%s/auth/login?redirect=%s", p.appBaseUrl, originalRequestURL), http.StatusFound)
+		http.Error(rw, "Internal Server Error", http.StatusInternalServerError)
 		return
 	}
 
-	sessionID := cookie.Value
-	verifyURL := fmt.Sprintf("%s/badger/verify-user?sessionId=%s", p.apiBaseUrl, sessionID)
+	resp, err := http.Post(verifyURL, "application/json", bytes.NewBuffer(jsonData))
+	if err != nil {
+		http.Error(rw, "Internal Server Error", http.StatusInternalServerError)
+		return
+	}
+	defer resp.Body.Close()
 
-	resp, err := http.Get(verifyURL)
-	if err != nil || resp.StatusCode != http.StatusOK {
-		if resp != nil && resp.StatusCode == http.StatusUnauthorized {
-			http.Redirect(rw, req, p.appBaseUrl, http.StatusFound)
-		} else {
-			http.Error(rw, "Internal Server Error", http.StatusInternalServerError)
-		}
+	if resp.StatusCode != http.StatusOK {
+		http.Error(rw, "Internal Server Error", http.StatusInternalServerError)
+		return
+	}
+
+	var result VerifyResponse
+	err = json.NewDecoder(resp.Body).Decode(&result)
+	if err != nil {
+		http.Error(rw, "Internal Server Error", http.StatusInternalServerError)
+		return
+	}
+
+	if result.RedirectURL != nil && *result.RedirectURL != "" {
+		http.Redirect(rw, req, *result.RedirectURL, http.StatusFound)
+		return
+	}
+
+	if !result.Valid { // only do this if for some reason the API doesn't return a redirect and it's not valid
+		http.Error(rw, "Unauthorized", http.StatusUnauthorized)
 		return
 	}
 
 	p.next.ServeHTTP(rw, req)
 }
 
+func extractCookies(req *http.Request) CookieData {
+	var cookies CookieData
+
+	if appSSOSessionCookie, err := req.Cookie(AppSSOSessionCookieName); err == nil {
+		cookies.Session = &appSSOSessionCookie.Value
+	}
+
+	if resourceSessionCookie, err := req.Cookie(ResourceSessionCookieName); err == nil {
+		cookies.ResourceSession = &resourceSessionCookie.Value
+	}
+
+	return cookies
+}
+
 func (p *Badger) getScheme(req *http.Request) string {
 	if req.TLS != nil {
 		return "https"