From e951e42b4dc8348ba4a82fc04ca39674a6892178 Mon Sep 17 00:00:00 2001
From: miloschwartz <mschwartz10612@gmail.com>
Date: Sat, 5 Apr 2025 22:31:23 -0400
Subject: [PATCH] remove explicit access token check and pass query params and
 headers in verify session

---
 .traefik.yml |  1 -
 README.md    |  1 -
 main.go      | 51 ++++++++++++++++++++++++++++++++++++---------------
 3 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/.traefik.yml b/.traefik.yml
index 8c27f6c..cbcbc78 100644
--- a/.traefik.yml
+++ b/.traefik.yml
@@ -8,5 +8,4 @@ summary: Middleware auth bouncer for Pangolin
 testData:
     apiBaseUrl: "http://localhost:3001/api/v1"
     userSessionCookieName: "p_session_token"
-    accessTokenQueryParam: "p_token"
     resourceSessionRequestParam: "p_session_request"
diff --git a/README.md b/README.md
index 10181d2..b176014 100644
--- a/README.md
+++ b/README.md
@@ -17,7 +17,6 @@ Badger requires the following configuration parameters to be specified in your [
 ```yaml
 apiBaseUrl: "http://localhost:3001/api/v1"
 userSessionCookieName: "p_session_token"
-accessTokenQueryParam: "p_token"
 resourceSessionRequestParam: "p_session_request"
 ```
 
diff --git a/main.go b/main.go
index 5c01735..2d37c34 100644
--- a/main.go
+++ b/main.go
@@ -12,7 +12,6 @@ import (
 type Config struct {
 	APIBaseUrl                  string `json:"apiBaseUrl"`
 	UserSessionCookieName       string `json:"userSessionCookieName"`
-	AccessTokenQueryParam       string `json:"accessTokenQueryParam"`
 	ResourceSessionRequestParam string `json:"resourceSessionRequestParam"`
 }
 
@@ -21,7 +20,6 @@ type Badger struct {
 	name                        string
 	apiBaseUrl                  string
 	userSessionCookieName       string
-	accessTokenQueryParam       string
 	resourceSessionRequestParam string
 }
 
@@ -32,15 +30,17 @@ type VerifyBody struct {
 	RequestHost        *string           `json:"host"`
 	RequestPath        *string           `json:"path"`
 	RequestMethod      *string           `json:"method"`
-	AccessToken        *string           `json:"accessToken,omitempty"`
 	TLS                bool              `json:"tls"`
 	RequestIP          *string           `json:"requestIp,omitempty"`
+	Headers            map[string]string `json:"headers,omitempty"`
+	Query              map[string]string `json:"query,omitempty"`
 }
 
 type VerifyResponse struct {
 	Data struct {
-		Valid       bool    `json:"valid"`
-		RedirectURL *string `json:"redirectUrl"`
+		Valid           bool              `json:"valid"`
+		RedirectURL     *string           `json:"redirectUrl"`
+		ResponseHeaders map[string]string `json:"responseHeaders,omitempty"`
 	} `json:"data"`
 }
 
@@ -52,8 +52,9 @@ type ExchangeSessionBody struct {
 
 type ExchangeSessionResponse struct {
 	Data struct {
-		Valid  bool    `json:"valid"`
-		Cookie *string `json:"cookie"`
+		Valid           bool              `json:"valid"`
+		Cookie          *string           `json:"cookie"`
+		ResponseHeaders map[string]string `json:"responseHeaders,omitempty"`
 	} `json:"data"`
 }
 
@@ -67,7 +68,6 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
 		name:                        name,
 		apiBaseUrl:                  config.APIBaseUrl,
 		userSessionCookieName:       config.UserSessionCookieName,
-		accessTokenQueryParam:       config.AccessTokenQueryParam,
 		resourceSessionRequestParam: config.ResourceSessionRequestParam,
 	}, nil
 }
@@ -75,7 +75,6 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
 func (p *Badger) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	cookies := p.extractCookies(req)
 
-	var accessToken *string
 	queryValues := req.URL.Query()
 
 	if sessionRequestValue := queryValues.Get(p.resourceSessionRequestParam); sessionRequestValue != "" {
@@ -116,17 +115,18 @@ func (p *Badger) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 				originalRequestURL = fmt.Sprintf("%s?%s", originalRequestURL, cleanedQuery)
 			}
 
+			if result.Data.ResponseHeaders != nil {
+				for key, value := range result.Data.ResponseHeaders {
+					rw.Header().Add(key, value)
+				}
+			}
+
 			fmt.Println("Got exchange token, redirecting to", originalRequestURL)
 			http.Redirect(rw, req, originalRequestURL, http.StatusFound)
 			return
 		}
 	}
 
-	if token := queryValues.Get(p.accessTokenQueryParam); token != "" {
-		accessToken = &token
-		queryValues.Del(p.accessTokenQueryParam)
-	}
-
 	cleanedQuery := queryValues.Encode()
 	originalRequestURL := fmt.Sprintf("%s://%s%s", p.getScheme(req), req.Host, req.URL.Path)
 	if cleanedQuery != "" {
@@ -135,6 +135,20 @@ func (p *Badger) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 
 	verifyURL := fmt.Sprintf("%s/badger/verify-session", p.apiBaseUrl)
 
+	headers := make(map[string]string)
+	for name, values := range req.Header {
+		if len(values) > 0 {
+			headers[name] = values[0] // Send only the first value for simplicity
+		}
+	}
+
+	queryParams := make(map[string]string)
+	for key, values := range queryValues {
+		if len(values) > 0 {
+			queryParams[key] = values[0]
+		}
+	}
+
 	cookieData := VerifyBody{
 		Sessions:           cookies,
 		OriginalRequestURL: originalRequestURL,
@@ -142,9 +156,10 @@ func (p *Badger) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 		RequestHost:        &req.Host,
 		RequestPath:        &req.URL.Path,
 		RequestMethod:      &req.Method,
-		AccessToken:        accessToken,
 		TLS:                req.TLS != nil,
 		RequestIP:          &req.RemoteAddr,
+		Headers:            headers,
+		Query:              queryParams,
 	}
 
 	jsonData, err := json.Marshal(cookieData)
@@ -176,6 +191,12 @@ func (p *Badger) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 		return
 	}
 
+	if result.Data.ResponseHeaders != nil {
+		for key, value := range result.Data.ResponseHeaders {
+			rw.Header().Add(key, value)
+		}
+	}
+
 	if result.Data.RedirectURL != nil && *result.Data.RedirectURL != "" {
 		fmt.Println("Badger: Redirecting to", *result.Data.RedirectURL)
 		http.Redirect(rw, req, *result.Data.RedirectURL, http.StatusFound)