Skip to content

Commit

Permalink
fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
Misfits09 committed Sep 11, 2023
1 parent 9bd47bf commit 4b968dd
Show file tree
Hide file tree
Showing 5 changed files with 101 additions and 295 deletions.
11 changes: 6 additions & 5 deletions .traefik.yml
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
displayName: Fifteen
displayName: JWT Field as Header
type: middleware
iconPath: .assets/icon.png

Expand All @@ -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

2 changes: 2 additions & 0 deletions Dockerfile
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
FROM traefik:v2.10.4
COPY . /plugins-local/src/github.com/birotaio/traefik-plugins/
96 changes: 72 additions & 24 deletions fifteen.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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())
Expand Down Expand Up @@ -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) {
Expand All @@ -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
}
4 changes: 3 additions & 1 deletion fifteen_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading

0 comments on commit 4b968dd

Please sign in to comment.