Skip to content

Commit

Permalink
feat: Add ExtractTokenFromRequest method in Gate struct
Browse files Browse the repository at this point in the history
  • Loading branch information
TwiN committed Jan 2, 2022
1 parent f18d51d commit 8f6fc57
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 8 deletions.
21 changes: 13 additions & 8 deletions gate.go
Original file line number Diff line number Diff line change
Expand Up @@ -181,12 +181,7 @@ func (gate *Gate) ProtectFuncWithPermissions(handlerFunc http.HandlerFunc, permi
}
}
if gate.authorizationService != nil {
var token string
if gate.customTokenExtractorFunc != nil {
token = gate.customTokenExtractorFunc(request)
} else {
token = extractTokenFromRequest(request)
}
token := gate.ExtractTokenFromRequest(request)
if !gate.authorizationService.IsAuthorized(token, permissions) {
writer.WriteHeader(http.StatusUnauthorized)
_, _ = writer.Write(gate.unauthorizedResponseBody)
Expand All @@ -206,7 +201,17 @@ func (gate *Gate) ProtectFuncWithPermission(handlerFunc http.HandlerFunc, permis
return gate.ProtectFuncWithPermissions(handlerFunc, []string{permission})
}

// extractTokenFromRequest extracts the bearer token from the AuthorizationHeader
func extractTokenFromRequest(request *http.Request) string {
// ExtractTokenFromRequest extracts a token from a request.
//
// By default, it extracts the bearer token from the AuthorizationHeader, but if a customTokenExtractorFunc is defined,
// it will use that instead.
//
// Note that this method is internally used by Protect, ProtectWithPermission, ProtectFunc and
// ProtectFuncWithPermissions, but it is exposed in case you need to use it directly.
func (gate *Gate) ExtractTokenFromRequest(request *http.Request) string {
if gate.customTokenExtractorFunc != nil {
// A custom token extractor function is defined, so we'll use it instead of the default token extraction logic
return gate.customTokenExtractorFunc(request)
}
return strings.TrimPrefix(request.Header.Get(AuthorizationHeader), "Bearer ")
}
27 changes: 27 additions & 0 deletions gate_bench_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -179,3 +179,30 @@ func BenchmarkGate_ProtectWithClientProviderConcurrently(b *testing.B) {
})
b.ReportAllocs()
}

func BenchmarkGate_ProtectWithValidTokenAndCustomTokenExtractorFuncConcurrently(b *testing.B) {
customTokenExtractorFunc := func(request *http.Request) string {
sessionCookie, err := request.Cookie("session")
if err != nil {
return ""
}
return sessionCookie.Value
}
gate := New().WithAuthorizationService(NewAuthorizationService().WithToken("good-token")).WithCustomTokenExtractor(customTokenExtractorFunc)
request, _ := http.NewRequest("GET", "/handle", http.NoBody)
request.AddCookie(&http.Cookie{Name: "session", Value: "good-token"})

router := http.NewServeMux()
router.Handle("/handle", gate.Protect(handler))

b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
responseRecorder := httptest.NewRecorder()
router.ServeHTTP(responseRecorder, request)
if responseRecorder.Code != http.StatusOK {
b.Fatalf("%s %s should have returned %d, but returned %d instead", request.Method, request.URL, http.StatusOK, responseRecorder.Code)
}
}
})
b.ReportAllocs()
}

0 comments on commit 8f6fc57

Please sign in to comment.