From ad468b3d156550096dc73be9aaf5bb3b761a2d3e Mon Sep 17 00:00:00 2001
From: Milo Schwartz <mschwartz10612@gmail.com>
Date: Sat, 11 Jan 2025 18:08:02 -0500
Subject: [PATCH] send access token and pass cookies

---
 .traefik.yml |  1 +
 README.md    |  5 +++--
 main.go      | 24 +++++++++++++++++++++++-
 3 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/.traefik.yml b/.traefik.yml
index ba69c2c..18d81b5 100644
--- a/.traefik.yml
+++ b/.traefik.yml
@@ -9,3 +9,4 @@ testData:
     apiBaseUrl: http://localhost:3001/api/v1
     userSessionCookieName: p_session
     resourceSessionCookieName: p_resource_session
+    accessTokenQueryParam: p_token
diff --git a/README.md b/README.md
index a68c39f..5b1fb04 100644
--- a/README.md
+++ b/README.md
@@ -18,8 +18,9 @@ Badger requires the following configuration parameters to be specified in your [
 
 ```yaml
 apiBaseUrl: "http://localhost:3001/api/v1"
-userSessionCookieName: "session"
-resourceSessionCookieName: "resource_session"
+userSessionCookieName: "p_session"
+resourceSessionCookieName: "p_resource_session"
+accessTokenQueryParam: "p_token"
 ```
 
 ---
diff --git a/main.go b/main.go
index 6f6630c..6d094a7 100644
--- a/main.go
+++ b/main.go
@@ -13,6 +13,7 @@ type Config struct {
 	APIBaseUrl                string `json:"apiBaseUrl"`
 	UserSessionCookieName     string `json:"userSessionCookieName"`
 	ResourceSessionCookieName string `json:"resourceSessionCookieName"`
+	AccessTokenQueryParam     string `json:"accessTokenQueryParam"`
 }
 
 type VerifyBody struct {
@@ -22,6 +23,7 @@ type VerifyBody struct {
 	RequestHost        *string           `json:"host"`
 	RequestPath        *string           `json:"path"`
 	RequestMethod      *string           `json:"method"`
+	AccessToken        *string           `json:"accessToken,omitempty"`
 	TLS                bool              `json:"tls"`
 }
 
@@ -38,6 +40,7 @@ type Badger struct {
 	apiBaseUrl                string
 	userSessionCookieName     string
 	resourceSessionCookieName string
+	accessTokenQueryParam     string
 }
 
 func CreateConfig() *Config {
@@ -51,14 +54,27 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
 		apiBaseUrl:                config.APIBaseUrl,
 		userSessionCookieName:     config.UserSessionCookieName,
 		resourceSessionCookieName: config.ResourceSessionCookieName,
+		accessTokenQueryParam:     config.AccessTokenQueryParam,
 	}, nil
 }
 
 func (p *Badger) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	cookies := p.extractCookies(req)
 
+	var accessToken *string
+	queryValues := req.URL.Query()
+	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 != "" {
+		originalRequestURL = fmt.Sprintf("%s?%s", originalRequestURL, cleanedQuery)
+	}
+
 	verifyURL := fmt.Sprintf("%s/badger/verify-session", p.apiBaseUrl)
-	originalRequestURL := fmt.Sprintf("%s://%s%s", p.getScheme(req), req.Host, req.URL.RequestURI())
 
 	cookieData := VerifyBody{
 		Sessions:           cookies,
@@ -67,6 +83,7 @@ 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,
 	}
 
@@ -83,6 +100,11 @@ func (p *Badger) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	}
 	defer resp.Body.Close()
 
+	// pass through cookies
+	for _, setCookie := range resp.Header["Set-Cookie"] {
+		rw.Header().Add("Set-Cookie", setCookie)
+	}
+
 	if resp.StatusCode != http.StatusOK {
 		http.Error(rw, "Internal Server Error", http.StatusInternalServerError)
 		return