From f220c75a52f0654c7a77f51282ca0b1056f3d6c5 Mon Sep 17 00:00:00 2001
From: miloschwartz <mschwartz10612@gmail.com>
Date: Tue, 17 Mar 2026 16:59:32 -0700
Subject: [PATCH] use access token params from config

---
 README.md | 14 --------------
 main.go   | 20 +++++++++++++++++---
 2 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index 58fe41c..88e7d55 100644
--- a/README.md
+++ b/README.md
@@ -70,20 +70,6 @@ trustip:
 customIPHeader: "X-Forwarded-For"
 ```
 
-### Configuration Options Reference
-
-| Option                        | Type     | Required\* | Default | Description                                                                         |
-| ----------------------------- | -------- | ---------- | ------- | ----------------------------------------------------------------------------------- |
-| `disableForwardAuth`          | bool     | No         | `false` | Disable forward auth; only IP handling is performed                                 |
-| `apiBaseUrl`                  | string   | Yes\*      | -       | Base URL of the Pangolin API                                                        |
-| `userSessionCookieName`       | string   | Yes\*      | -       | Cookie name for user sessions                                                       |
-| `resourceSessionRequestParam` | string   | Yes\*      | -       | Query parameter name for resource session requests                                  |
-| `trustip`                     | []string | No         | `[]`    | Array of trusted IP ranges in CIDR format                                           |
-| `disableDefaultCFIPs`         | bool     | No         | `false` | Disable default Cloudflare IP ranges                                                |
-| `customIPHeader`              | string   | No         | `""`    | Custom header name to extract IP from (only used if request is from trusted source) |
-
-\* Required only when `disableForwardAuth` is `false` (default)
-
 ## Updating Cloudflare IPs
 
 To update the Cloudflare IP ranges, run:
diff --git a/main.go b/main.go
index 0cf2d92..0258e1c 100644
--- a/main.go
+++ b/main.go
@@ -17,7 +17,9 @@ type Config struct {
 	APIBaseUrl                  string   `json:"apiBaseUrl,omitempty"`
 	UserSessionCookieName       string   `json:"userSessionCookieName,omitempty"`
 	ResourceSessionRequestParam string   `json:"resourceSessionRequestParam,omitempty"`
-	AccessTokenQueryParam       string   `json:"accessTokenQueryParam,omitempty"` // Deprecated: use ResourceSessionRequestParam
+	AccessTokenQueryParam       string   `json:"accessTokenQueryParam,omitempty"`
+	AccessTokenIDHeader         string   `json:"accessTokenIdHeader,omitempty"`
+	AccessTokenHeader           string   `json:"accessTokenHeader,omitempty"`
 	DisableForwardAuth          bool     `json:"disableForwardAuth,omitempty"`
 	TrustIP                     []string `json:"trustip,omitempty"`
 	DisableDefaultCFIPs         bool     `json:"disableDefaultCFIPs,omitempty"`
@@ -39,6 +41,8 @@ type Badger struct {
 	userSessionCookieName       string
 	resourceSessionRequestParam string
 	accessTokenQueryParam       string
+	accessTokenIDHeader         string
+	accessTokenHeader           string
 	disableForwardAuth          bool
 	trustIP                     []*net.IPNet
 	customIPHeader              string
@@ -98,6 +102,8 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
 		userSessionCookieName:       config.UserSessionCookieName,
 		resourceSessionRequestParam: config.ResourceSessionRequestParam,
 		accessTokenQueryParam:       config.AccessTokenQueryParam,
+		accessTokenIDHeader:         config.AccessTokenIDHeader,
+		accessTokenHeader:           config.AccessTokenHeader,
 		disableForwardAuth:          config.DisableForwardAuth,
 		customIPHeader:              config.CustomIPHeader,
 	}
@@ -317,8 +323,7 @@ func (p *Badger) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 
 		p.stripSessionCookies(req)
 		p.stripSessionParam(req)
-		req.Header.Del("P-Access-Token-Id")
-		req.Header.Del("P-Access-Token")
+		p.stripAccessTokenHeaders(req)
 
 		fmt.Println("Badger: Valid session")
 		p.next.ServeHTTP(rw, req)
@@ -439,6 +444,15 @@ func (p *Badger) stripSessionParam(req *http.Request) {
 	}
 }
 
+func (p *Badger) stripAccessTokenHeaders(req *http.Request) {
+	if p.accessTokenIDHeader != "" {
+		req.Header.Del(p.accessTokenIDHeader)
+	}
+	if p.accessTokenHeader != "" {
+		req.Header.Del(p.accessTokenHeader)
+	}
+}
+
 // stripSessionCookies removes session cookies from the request before forwarding to the backend.
 // It processes raw Cookie header pairs so non-target cookies are preserved as-is.
 func (p *Badger) stripSessionCookies(req *http.Request) {