fixes

Author: Misfits09

.traefik.yml

@@ -1,4 +1,4 @@
-displayName: Fifteen
+displayName: JWT Field as Header
 type: middleware
 iconPath: .assets/icon.png
 
@@ -7,9 +7,10 @@ import: github.com/birotaio/traefik-plugins
 summary: 'Make custom header from JWT data, can be used for user-based ratelimiting'
 
 testData:
-  jwt-header-name: X-ApiKey
-  jwt-field: customer_id
-  value-header-name: X-UserId-RateLimit
-  fallback-type: ip
+  jwtHeaderName: X-ApiKey
+  jwtField: customer_id
+  valueHeaderName: X-UserId-RateLimit
+  fallbacks:
+    - type: ip
   debug: true
 

Dockerfile

@@ -0,0 +1,2 @@
+FROM traefik:v2.10.4
+COPY . /plugins-local/src/github.com/birotaio/traefik-plugins/

fifteen.go

@@ -5,27 +5,33 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"strings"
 
 	"github.com/golang-jwt/jwt/v5"
 )
 
-type Fallback string
+type FallbackType string
 
 const (
-	FallbackError  Fallback = "error"
-	FallbackPass   Fallback = "pass"
-	FallbackIp     Fallback = "ip"
-	FallbackHeader Fallback = "header"
+	FallbackError  FallbackType = "error"
+	FallbackPass   FallbackType = "pass"
+	FallbackIp     FallbackType = "ip"
+	FallbackHeader FallbackType = "header"
 )
 
+type Fallback struct {
+	Type        FallbackType `yaml:"type,omitempty"`
+	Value       string       `yaml:"value,omitempty"`
+	KeepIfEmpty bool         `yaml:"keepIfEmpty,omitempty"`
+}
+
 // Config the plugin configuration.
 type Config struct {
-	JwtHeaderName      string   `json:"jwt-header-name,omitempty"`
-	JwtField           string   `json:"jwt-field,omitempty"`
-	ValueHeaderName    string   `json:"value-header-name,omitempty"`
-	FallbackType       Fallback `json:"fallback-type,omitempty"`
-	FallbackHeaderName string   `json:"fallback-header-name,omitempty"`
-	Debug              bool     `json:"debug,omitempty"`
+	JwtHeaderName   string     `yaml:"jwtHeaderName,omitempty"`
+	JwtField        string     `yaml:"jwtField,omitempty"`
+	ValueHeaderName string     `yaml:"valueHeaderName,omitempty"`
+	Fallbacks       []Fallback `yaml:"fallbacks,omitempty"`
+	Debug           bool       `yaml:"debug,omitempty"`
 }
 
 // CreateConfig creates the default plugin configuration.
@@ -56,7 +62,11 @@ func (a *Fifteen) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	rawToken := req.Header.Get(a.cfg.JwtHeaderName)
+	rawHeader := req.Header.Get(a.cfg.JwtHeaderName)
+	rawToken := ""
+	if strings.HasPrefix(rawHeader, "Bearer ") {
+		rawToken = rawHeader[len("Bearer "):]
+	}
 	parsedToken, _, err := jwt.NewParser().ParseUnverified(rawToken, jwt.MapClaims{})
 	if err != nil {
 		a.logDebug("Could not parse non-empty jwt token, falling back: %s", err.Error())
@@ -84,26 +94,52 @@ func (a *Fifteen) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 			return
 		}
 	} else {
-		a.logDebug("JWT field value has an unexpected type, falling back")
+		a.logDebug("JWT field value does not hold field %s, falling back", a.cfg.JwtField)
 		a.ServeFallback(rw, req)
 		return
 	}
 
-	a.next.ServeHTTP(rw, req)
+	a.end(rw, req)
 }
 
 func (a *Fifteen) ServeFallback(rw http.ResponseWriter, req *http.Request) {
-	a.logDebug("Fallbacked because JWT was not set, invalid or has unexpected value on field. Using fallback strategy: %s", a.cfg.FallbackType)
-	switch a.cfg.FallbackType {
-	case FallbackError:
-		rw.WriteHeader(http.StatusBadRequest)
-	case FallbackIp:
-		req.Header.Set(a.cfg.ValueHeaderName, req.RemoteAddr)
-	case FallbackHeader:
-		req.Header.Set(a.cfg.ValueHeaderName, req.Header.Get(a.cfg.FallbackHeaderName))
-	default:
-		a.next.ServeHTTP(rw, req)
+	if len(a.cfg.Fallbacks) == 0 {
+		a.logDebug("Fallbacked because JWT was not set, invalid or has unexpected value on field. No fallback strategies, ignoring...")
+	} else {
+		a.logDebug("Fallbacked because JWT was not set, invalid or has unexpected value on field. Finding right fallback strategy")
+		for i, fallback := range a.cfg.Fallbacks {
+			a.logDebug("Strategy %d: %+v", i, fallback)
+			var success bool
+			switch fallback.Type {
+			case FallbackError:
+				rw.Header().Set("Content-Type", "text/plain")
+				rw.WriteHeader(http.StatusBadRequest)
+				rw.Write([]byte("Bad request"))
+				return
+			case FallbackPass:
+				a.logDebug("Passing through")
+				success = true
+			case FallbackIp:
+				req.Header.Set(a.cfg.ValueHeaderName, ipWithNoPort(req.RemoteAddr))
+				success = true
+			case FallbackHeader:
+				headerValue := req.Header.Get(fallback.Value)
+				if headerValue == "" && !fallback.KeepIfEmpty {
+					a.logDebug("Header %s was empty, skipping...", fallback.Value)
+					continue
+				}
+				req.Header.Set(a.cfg.ValueHeaderName, headerValue)
+				success = true
+			default:
+				a.logDebug("Unknown fallback type, skipping...")
+			}
+			if success {
+				a.logDebug("Fallback strategy %d was successful", i)
+				break
+			}
+		}
 	}
+	a.end(rw, req)
 }
 
 func (a *Fifteen) logDebug(format string, args ...any) {
@@ -112,3 +148,15 @@ func (a *Fifteen) logDebug(format string, args ...any) {
 	}
 	os.Stderr.WriteString("[Fifteen middleware]: " + fmt.Sprintf(format, args...) + "\n")
 }
+
+func (a *Fifteen) end(rw http.ResponseWriter, req *http.Request) {
+	a.logDebug("ending with request headers: %+v", req.Header)
+	a.next.ServeHTTP(rw, req)
+}
+
+func ipWithNoPort(addr string) string {
+	if colon := strings.LastIndex(addr, ":"); colon != -1 {
+		return addr[:colon]
+	}
+	return addr
+}

fifteen_test.go

@@ -12,7 +12,9 @@ func TestDemo(t *testing.T) {
 	cfg.JwtHeaderName = "X-ApiKey"
 	cfg.JwtField = "customer_id"
 	cfg.ValueHeaderName = "X-UserId-RateLimit"
-	cfg.FallbackType = FallbackIp
+	cfg.Fallbacks = []Fallback{
+		{Type: FallbackIp},
+	}
 	cfg.Debug = false
 
 	ctx := context.Background()