From f0d17fd839f8d39917f509e36ed2504beec868d9 Mon Sep 17 00:00:00 2001
From: Livio Amstutz <livio.a@gmail.com>
Date: Thu, 6 Feb 2020 11:12:00 +0100
Subject: [PATCH] feat: add http interceptor function for auth and token
 endpoints

---
 pkg/op/default_op.go | 10 +++++++++-
 pkg/op/op.go         | 21 +++++++++++++++++----
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/pkg/op/default_op.go b/pkg/op/default_op.go
index 8d140b5..7ce7925 100644
--- a/pkg/op/default_op.go
+++ b/pkg/op/default_op.go
@@ -41,6 +41,7 @@ type DefaultOP struct {
 	http            *http.Server
 	decoder         *schema.Decoder
 	encoder         *schema.Encoder
+	interceptor     HttpInterceptor
 }
 
 type Config struct {
@@ -98,6 +99,13 @@ func WithCustomUserinfoEndpoint(endpoint Endpoint) DefaultOPOpts {
 	}
 }
 
+func WithHttpInterceptor(h HttpInterceptor) DefaultOPOpts {
+	return func(o *DefaultOP) error {
+		o.interceptor = h
+		return nil
+	}
+}
+
 func NewDefaultOP(ctx context.Context, config *Config, storage Storage, opOpts ...DefaultOPOpts) (OpenIDProvider, error) {
 	err := ValidateIssuer(config.Issuer)
 	if err != nil {
@@ -123,7 +131,7 @@ func NewDefaultOP(ctx context.Context, config *Config, storage Storage, opOpts .
 
 	p.discoveryConfig = CreateDiscoveryConfig(p, p.signer)
 
-	router := CreateRouter(p)
+	router := CreateRouter(p, p.interceptor)
 	p.http = &http.Server{
 		Addr:    ":" + config.Port,
 		Handler: router,
diff --git a/pkg/op/op.go b/pkg/op/op.go
index 7db2ff4..4d64e04 100644
--- a/pkg/op/op.go
+++ b/pkg/op/op.go
@@ -21,12 +21,25 @@ type OpenIDProvider interface {
 	HttpHandler() *http.Server
 }
 
-func CreateRouter(o OpenIDProvider) *mux.Router {
+type HttpInterceptor func(http.HandlerFunc) http.HandlerFunc
+
+var (
+	DefaultInterceptor = func(h http.HandlerFunc) http.HandlerFunc {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			h(w, r)
+		})
+	}
+)
+
+func CreateRouter(o OpenIDProvider, h HttpInterceptor) *mux.Router {
+	if h == nil {
+		h = DefaultInterceptor
+	}
 	router := mux.NewRouter()
 	router.HandleFunc(oidc.DiscoveryEndpoint, o.HandleDiscovery)
-	router.HandleFunc(o.AuthorizationEndpoint().Relative(), o.HandleAuthorize)
-	router.HandleFunc(o.AuthorizationEndpoint().Relative()+"/{id}", o.HandleAuthorizeCallback)
-	router.HandleFunc(o.TokenEndpoint().Relative(), o.HandleExchange)
+	router.HandleFunc(o.AuthorizationEndpoint().Relative(), h(o.HandleAuthorize))
+	router.HandleFunc(o.AuthorizationEndpoint().Relative()+"/{id}", h(o.HandleAuthorizeCallback))
+	router.HandleFunc(o.TokenEndpoint().Relative(), h(o.HandleExchange))
 	router.HandleFunc(o.UserinfoEndpoint().Relative(), o.HandleUserinfo)
 	router.HandleFunc(o.KeysEndpoint().Relative(), o.HandleKeys)
 	return router