From b73aef6e3a8f5d2cba5dae9c1467e287c0494f76 Mon Sep 17 00:00:00 2001
From: Dmitrii Kozlov <dmitrii.kozlov@madappgang.com>
Date: Tue, 10 Sep 2024 14:04:09 +0300
Subject: [PATCH] more context for JWT middleware

---
 jwt/middleware/middleware.go | 90 +++++++++++++++++++++++++++++++-----
 1 file changed, 78 insertions(+), 12 deletions(-)

diff --git a/jwt/middleware/middleware.go b/jwt/middleware/middleware.go
index 547a170d..269b1f10 100644
--- a/jwt/middleware/middleware.go
+++ b/jwt/middleware/middleware.go
@@ -27,41 +27,107 @@ const (
 // next is a reference to the next handler in the handler chain
 type Handler func(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc)
 
-// ErrorHandler interface for handling error from middleware
-// rw - readwriter to write error to (JSON, HTML or other error to the client)
-// errorType - error returned from middleware, you can get description by calling errorType.Description()
-// status - http status code
-type ErrorHandler interface {
-	Error(rw http.ResponseWriter, errorType Error, status int, description string)
+// JWTErrorHandlerContext is a context for handling errors from JWTv2 middleware.
+type JWTErrorHandlerContext struct {
+	// ResponseWriter is an http.ResponseWriter to write error to (JSON, HTML or other error to the client).
+	ResponseWriter http.ResponseWriter
+	// Request is an http.Request to get more information about the request.
+	Request *http.Request
+	// ErrorType is an error returned from middleware, you can get description by calling errorType.Description().
+	ErrorType Error
+	// Status is an http status code.
+	Status int
+	// Description is a description of the error.
+	Description string
+	// Token is a parsed but not valid token.
+	Token model.Token
 }
 
-// JWT returns middleware function you can use to handle JWT token auth
-func JWT(eh ErrorHandler, c validator.Config) (Handler, error) {
+// JWTErrorHandler is an interface for handling errors from JWTv2 middleware.
+type JWTErrorHandler interface {
+	Error(errContext *JWTErrorHandlerContext)
+}
+
+// JWTv2 returns middleware function you can use to handle JWT token auth.
+// It extracts token from Authorization header, validates it and stores the parsed token in the context,
+// use TokenFromContext to get token from context.
+// It uses JWTErrorHandler to handle errors.
+func JWTv2(eh JWTErrorHandler, c validator.Config) (Handler, error) {
 	v, err := validator.NewValidatorWithConfig(c)
 	if err != nil {
 		return nil, err
 	}
+
+	reportError := func(
+		rw http.ResponseWriter,
+		r *http.Request,
+		errType Error,
+		status int,
+		description string,
+		token model.Token,
+	) {
+		ec := &JWTErrorHandlerContext{
+			ResponseWriter: rw,
+			Request:        r,
+			ErrorType:      errType,
+			Status:         status,
+			Description:    description,
+			Token:          token,
+		}
+
+		eh.Error(ec)
+	}
+
 	// Middleware middleware functions extracts token and validates it and store the parsed token in the context
 	return func(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
 		tokenBytes := jwt.ExtractTokenFromBearerHeader(r.Header.Get(AuthorizationHeaderKey))
 		if tokenBytes == nil {
-			eh.Error(rw, ErrorTokenIsEmpty, http.StatusBadRequest, "")
+			reportError(rw, r, ErrorTokenIsEmpty, http.StatusBadRequest, "", nil)
 			return
 		}
 		tokenString := string(tokenBytes)
 
 		token, err := v.ValidateString(tokenString)
 		if err != nil {
-			eh.Error(rw, ErrorTokenIsInvalid, http.StatusBadRequest, err.Error())
+			invalidToken, _ := jwt.ParseTokenString(tokenString)
+			reportError(rw, r, ErrorTokenIsInvalid, http.StatusBadRequest, err.Error(), invalidToken)
 			return
 		}
 
-		ctx := context.WithValue(r.Context(), model.TokenContextKey, token)
-		r = r.WithContext(ctx)
+		r = appendRequestContextValue(r, model.TokenContextKey, token)
+
 		next.ServeHTTP(rw, r)
 	}, nil
 }
 
+func appendRequestContextValue(r *http.Request, key, value interface{}) *http.Request {
+	ctx := context.WithValue(r.Context(), key, value)
+	return r.WithContext(ctx)
+}
+
+// ErrorHandler interface for handling error from middleware
+// rw - http.ResponseWriter to write error to (JSON, HTML or other error to the client)
+// errorType - error returned from middleware, you can get description by calling errorType.Description()
+// status - http status code
+type ErrorHandler interface {
+	Error(rw http.ResponseWriter, errorType Error, status int, description string)
+}
+
+// jwtErrorHandler is a wrapper for ErrorHandler to make it compatible with JWTErrorHandler
+type jwtErrorHandler struct {
+	eh ErrorHandler
+}
+
+func (j jwtErrorHandler) Error(errContext *JWTErrorHandlerContext) {
+	j.eh.Error(errContext.ResponseWriter, errContext.ErrorType, errContext.Status, errContext.Description)
+}
+
+// JWT returns middleware function you can use to handle JWT token auth
+// Deprecated: use JWT instead
+func JWT(eh ErrorHandler, c validator.Config) (Handler, error) {
+	return JWTv2(jwtErrorHandler{eh}, c)
+}
+
 // TokenFromContext returns token from request context.
 // Or nil if there is no token in context.
 func TokenFromContext(ctx context.Context) model.Token {