From 1f17a3bd8c13c23b8f7c861a3cc3d904daecf03f Mon Sep 17 00:00:00 2001
From: Stanley Phu <stanley@warrant.dev>
Date: Wed, 5 Apr 2023 23:41:54 -0700
Subject: [PATCH] Allow custom auth middleware on the route level (#62)

* Allow custom auth middleware on the route level

* Rename args to options for auth middleware

* Refactor enableSessionAuth flag check
---
 pkg/service/auth.go   | 15 +++++++++++++--
 pkg/service/router.go |  8 +++++++-
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/pkg/service/auth.go b/pkg/service/auth.go
index ccc8355a..26363323 100644
--- a/pkg/service/auth.go
+++ b/pkg/service/auth.go
@@ -24,14 +24,18 @@ const (
 	authInfoKey key = iota
 )
 
+const (
+	EnableSessionAuthKey = "EnableSessionAuth"
+)
+
 type AuthInfo struct {
 	UserId   string
 	TenantId string
 }
 
-type AuthMiddlewareFunc func(next http.Handler, config *config.Config, enableSessionAuth bool) http.Handler
+type AuthMiddlewareFunc func(next http.Handler, config *config.Config, options map[string]interface{}) http.Handler
 
-func DefaultAuthMiddleware(next http.Handler, config *config.Config, enableSessionAuth bool) http.Handler {
+func DefaultAuthMiddleware(next http.Handler, config *config.Config, options map[string]interface{}) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		logger := hlog.FromRequest(r)
 
@@ -60,6 +64,13 @@ func DefaultAuthMiddleware(next http.Handler, config *config.Config, enableSessi
 			}
 			authInfo = &AuthInfo{}
 		case "Bearer":
+			enableSessionAuth, ok := options["enableSessionAuth"].(bool)
+			if !ok {
+				SendErrorResponse(w, NewUnauthorizedError("Error validating token"))
+				logger.Err(fmt.Errorf("enableSessionAuth must be of type bool"))
+				return
+			}
+
 			if !enableSessionAuth {
 				SendErrorResponse(w, NewUnauthorizedError("Error validating token"))
 				logger.Err(fmt.Errorf("invalid authentication for the endpoint")).Msg("Session authentication not supported for this endpoint")
diff --git a/pkg/service/router.go b/pkg/service/router.go
index 2131ae7c..8eda8085 100644
--- a/pkg/service/router.go
+++ b/pkg/service/router.go
@@ -18,6 +18,7 @@ type Route struct {
 	Pattern           string
 	Method            string
 	Handler           http.Handler
+	AuthMiddleware    AuthMiddlewareFunc
 	DisableAuth       bool
 	EnableSessionAuth bool
 }
@@ -85,11 +86,16 @@ func NewRouter(config *config.Config, pathPrefix string, routes []Route, authMid
 
 	// Setup routes
 	for _, route := range routes {
+		defaultOptions := map[string]interface{}{
+			EnableSessionAuthKey: route.EnableSessionAuth,
+		}
 		routePattern := fmt.Sprintf("%s%s", pathPrefix, route.Pattern)
 		if route.DisableAuth || config.ApiKey == "" {
 			router.Handle(routePattern, route.Handler).Methods(route.Method)
+		} else if route.AuthMiddleware != nil {
+			router.Handle(routePattern, route.AuthMiddleware(route.Handler, config, defaultOptions)).Methods(route.Method)
 		} else {
-			router.Handle(routePattern, authMiddleware(route.Handler, config, route.EnableSessionAuth)).Methods(route.Method)
+			router.Handle(routePattern, authMiddleware(route.Handler, config, defaultOptions)).Methods(route.Method)
 		}
 	}