refactor: improve error interceptor

internal/middleware/access-middleware.go

@@ -16,7 +16,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 import (
-	"encoding/json"
 	"fmt"
 	"github.com/jkaninda/goma-gateway/pkg/logger"
 	"github.com/jkaninda/goma-gateway/util"
@@ -31,16 +30,7 @@ func (blockList AccessListMiddleware) AccessMiddleware(next http.Handler) http.H
 		for _, block := range blockList.List {
 			if isPathBlocked(r.URL.Path, util.ParseURLPath(blockList.Path+block)) {
 				logger.Error("%s: %s access forbidden", getRealIP(r), r.URL.Path)
-				w.Header().Set("Content-Type", "application/json")
-				w.WriteHeader(http.StatusForbidden)
-				err := json.NewEncoder(w).Encode(ProxyResponseError{
-					Success: false,
-					Code:    http.StatusForbidden,
-					Message: fmt.Sprintf("You do not have permission to access this resource"),
-				})
-				if err != nil {
-					return
-				}
+				RespondWithError(w, http.StatusForbidden, fmt.Sprintf("%d you do not have permission to access this resource", http.StatusForbidden), blockList.ErrorInterceptor)
 				return
 			}
 		}

internal/middleware/block-common-exploits.go

@@ -18,15 +18,19 @@
 package middleware
 
 import (
-	"encoding/json"
 	"fmt"
+	errorinterceptor "github.com/jkaninda/goma-gateway/pkg/error-interceptor"
 	"github.com/jkaninda/goma-gateway/pkg/logger"
 	"net/http"
 	"regexp"
 )
 
+type BlockCommon struct {
+	ErrorInterceptor errorinterceptor.ErrorInterceptor
+}
+
 // BlockExploitsMiddleware Middleware to block common exploits
-func BlockExploitsMiddleware(next http.Handler) http.Handler {
+func (blockCommon BlockCommon) BlockExploitsMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		// Patterns to detect SQL injection attempts
 		sqlInjectionPattern := regexp.MustCompile(sqlPatterns)
@@ -42,36 +46,18 @@ func BlockExploitsMiddleware(next http.Handler) http.Handler {
 			pathTraversalPattern.MatchString(r.URL.Path) ||
 			xssPattern.MatchString(r.URL.RawQuery) {
 			logger.Error("%s: %s Forbidden - Potential exploit detected", getRealIP(r), r.URL.Path)
-			w.Header().Set("Content-Type", "application/json")
-			w.WriteHeader(http.StatusForbidden)
-			err := json.NewEncoder(w).Encode(ProxyResponseError{
-				Success: false,
-				Code:    http.StatusForbidden,
-				Message: fmt.Sprintf("Forbidden - Potential exploit detected"),
-			})
-			if err != nil {
-				return
-			}
+			RespondWithError(w, http.StatusForbidden, fmt.Sprintf("%d Forbidden - Potential exploit detected", http.StatusForbidden), blockCommon.ErrorInterceptor)
 			return
-		}
 
+		}
 		// Check form data (for POST requests)
 		if r.Method == http.MethodPost {
 			if err := r.ParseForm(); err == nil {
 				for _, values := range r.Form {
 					for _, value := range values {
 						if sqlInjectionPattern.MatchString(value) || xssPattern.MatchString(value) {
-							logger.Error("%s: %s Forbidden - Potential exploit detected", getRealIP(r), r.URL.Path)
-							w.Header().Set("Content-Type", "application/json")
-							w.WriteHeader(http.StatusForbidden)
-							err := json.NewEncoder(w).Encode(ProxyResponseError{
-								Success: false,
-								Code:    http.StatusForbidden,
-								Message: fmt.Sprintf("Forbidden - Potential exploit detected"),
-							})
-							if err != nil {
-								return
-							}
+							logger.Error("%s: %s %s Forbidden - Potential exploit detected", getRealIP(r), r.Method, r.URL.Path)
+							RespondWithError(w, http.StatusForbidden, fmt.Sprintf("%d Forbidden - Potential exploit detected", http.StatusForbidden), blockCommon.ErrorInterceptor)
 							return
 						}
 					}

internal/middleware/error-interceptor.go

@@ -22,6 +22,7 @@ import (
 	"github.com/jkaninda/goma-gateway/pkg/logger"
 	"io"
 	"net/http"
+	"slices"
 )
 
 func newResponseRecorder(w http.ResponseWriter) *responseRecorder {
@@ -62,6 +63,7 @@ func (intercept InterceptErrors) ErrorInterceptor(next http.Handler) http.Handle
 			if err != nil {
 				return
 			}
+			return
 		} else {
 			// No error: write buffered response to client
 			w.WriteHeader(rec.statusCode)
@@ -69,18 +71,12 @@ func (intercept InterceptErrors) ErrorInterceptor(next http.Handler) http.Handle
 			if err != nil {
 				return
 			}
+			return
 
 		}
 
 	})
 }
 func canIntercept(code int, errors []int) bool {
-	for _, er := range errors {
-		if er == code {
-			return true
-		}
-		continue
-
-	}
-	return false
+	return slices.Contains(errors, code)
 }

internal/middleware/helpers.go

@@ -17,7 +17,12 @@
 
 package middleware
 
-import "net/http"
+import (
+	"encoding/json"
+	"fmt"
+	errorinterceptor "github.com/jkaninda/goma-gateway/pkg/error-interceptor"
+	"net/http"
+)
 
 func getRealIP(r *http.Request) string {
 	if ip := r.Header.Get("X-Real-IP"); ip != "" {
@@ -38,3 +43,53 @@ func allowedOrigin(origins []string, origin string) bool {
 	return false
 
 }
+func canInterceptError(code int, errors []errorinterceptor.Error) bool {
+	for _, er := range errors {
+		if er.Code == code {
+			return true
+		}
+		continue
+
+	}
+	return false
+}
+func errMessage(code int, errors []errorinterceptor.Error) (string, error) {
+	for _, er := range errors {
+		if er.Code == code {
+			if len(er.Message) != 0 {
+				return er.Message, nil
+			}
+			continue
+		}
+	}
+	return "", fmt.Errorf("%d errors occurred", code)
+}
+
+// RespondWithError is a helper function to handle error responses with flexible content type
+func RespondWithError(w http.ResponseWriter, statusCode int, logMessage string, errorIntercept errorinterceptor.ErrorInterceptor) {
+	message, err := errMessage(statusCode, errorIntercept.Errors)
+	if err != nil {
+		message = logMessage
+	}
+	if errorIntercept.ContentType == errorinterceptor.ApplicationJson {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(statusCode)
+		err := json.NewEncoder(w).Encode(ProxyResponseError{
+			Success: false,
+			Code:    statusCode,
+			Message: message,
+		})
+		if err != nil {
+			return
+		}
+		return
+	} else {
+		w.Header().Set("Content-Type", "plain/text;charset=utf-8")
+		w.WriteHeader(statusCode)
+		_, err2 := w.Write([]byte(message))
+		if err2 != nil {
+			return
+		}
+		return
+	}
+}

internal/middleware/middleware.go

@@ -17,7 +17,6 @@ limitations under the License.
 */
 import (
 	"encoding/base64"
-	"encoding/json"
 	"github.com/jkaninda/goma-gateway/pkg/logger"
 	"io"
 	"net/http"
@@ -38,48 +37,23 @@ func (jwtAuth JwtAuth) AuthMiddleware(next http.Handler) http.Handler {
 				if allowedOrigin(jwtAuth.Origins, r.Header.Get("Origin")) {
 					w.Header().Set("Access-Control-Allow-Origin", r.Header.Get("Origin"))
 				}
-				w.WriteHeader(http.StatusUnauthorized)
-				err := json.NewEncoder(w).Encode(ProxyResponseError{
-					Message: http.StatusText(http.StatusUnauthorized),
-					Code:    http.StatusUnauthorized,
-					Success: false,
-				})
-				if err != nil {
-					return
-				}
+				RespondWithError(w, http.StatusUnauthorized, http.StatusText(http.StatusUnauthorized), jwtAuth.ErrorInterceptor)
 				return
+
 			}
 		}
 		//token := r.Header.Get("Authorization")
 		authURL, err := url.Parse(jwtAuth.AuthURL)
 		if err != nil {
 			logger.Error("Error parsing auth URL: %v", err)
-			w.Header().Set("Content-Type", "application/json")
-			w.WriteHeader(http.StatusInternalServerError)
-			err = json.NewEncoder(w).Encode(ProxyResponseError{
-				Message: "Internal Server Error",
-				Code:    http.StatusInternalServerError,
-				Success: false,
-			})
-			if err != nil {
-				return
-			}
+			RespondWithError(w, http.StatusInternalServerError, http.StatusText(http.StatusInternalServerError), jwtAuth.ErrorInterceptor)
 			return
 		}
 		// Create a new request for /authentication
 		authReq, err := http.NewRequest("GET", authURL.String(), nil)
 		if err != nil {
 			logger.Error("Proxy error creating authentication request: %v", err)
-			w.Header().Set("Content-Type", "application/json")
-			w.WriteHeader(http.StatusInternalServerError)
-			err = json.NewEncoder(w).Encode(ProxyResponseError{
-				Message: "Internal Server Error",
-				Code:    http.StatusInternalServerError,
-				Success: false,
-			})
-			if err != nil {
-				return
-			}
+			RespondWithError(w, http.StatusInternalServerError, http.StatusText(http.StatusInternalServerError), jwtAuth.ErrorInterceptor)
 			return
 		}
 		logger.Trace("JWT Auth response headers: %v", authReq.Header)
@@ -99,16 +73,7 @@ func (jwtAuth JwtAuth) AuthMiddleware(next http.Handler) http.Handler {
 		if err != nil || authResp.StatusCode != http.StatusOK {
 			logger.Debug("%s %s %s %s", r.Method, getRealIP(r), r.URL, r.UserAgent())
 			logger.Debug("Proxy authentication error")
-			w.Header().Set("Content-Type", "application/json")
-			w.WriteHeader(http.StatusUnauthorized)
-			err = json.NewEncoder(w).Encode(ProxyResponseError{
-				Message: "Unauthorized",
-				Code:    http.StatusUnauthorized,
-				Success: false,
-			})
-			if err != nil {
-				return
-			}
+			RespondWithError(w, http.StatusUnauthorized, http.StatusText(http.StatusUnauthorized), jwtAuth.ErrorInterceptor)
 			return
 		}
 		defer func(Body io.ReadCloser) {
@@ -146,31 +111,14 @@ func (basicAuth AuthBasic) AuthMiddleware(next http.Handler) http.Handler {
 		if authHeader == "" {
 			logger.Debug("Proxy error, missing Authorization header")
 			w.Header().Set("WWW-Authenticate", `Basic realm="Restricted"`)
-			w.Header().Set("Content-Type", "application/json")
-			w.WriteHeader(http.StatusUnauthorized)
-			err := json.NewEncoder(w).Encode(ProxyResponseError{
-				Success: false,
-				Code:    http.StatusUnauthorized,
-				Message: http.StatusText(http.StatusUnauthorized),
-			})
-			if err != nil {
-				return
-			}
+			RespondWithError(w, http.StatusUnauthorized, http.StatusText(http.StatusUnauthorized), basicAuth.ErrorInterceptor)
 			return
 		}
 		// Check if the Authorization header contains "Basic" scheme
 		if !strings.HasPrefix(authHeader, "Basic ") {
 			logger.Error("Proxy error, missing Basic Authorization header")
-			w.Header().Set("Content-Type", "application/json")
-			w.WriteHeader(http.StatusUnauthorized)
-			err := json.NewEncoder(w).Encode(ProxyResponseError{
-				Success: false,
-				Code:    http.StatusUnauthorized,
-				Message: http.StatusText(http.StatusUnauthorized),
-			})
-			if err != nil {
-				return
-			}
+			RespondWithError(w, http.StatusUnauthorized, http.StatusText(http.StatusUnauthorized), basicAuth.ErrorInterceptor)
+
 			return
 		}
 
@@ -178,16 +126,7 @@ func (basicAuth AuthBasic) AuthMiddleware(next http.Handler) http.Handler {
 		payload, err := base64.StdEncoding.DecodeString(authHeader[len("Basic "):])
 		if err != nil {
 			logger.Debug("Proxy error, missing Basic Authorization header")
-			w.Header().Set("Content-Type", "application/json")
-			w.WriteHeader(http.StatusUnauthorized)
-			err := json.NewEncoder(w).Encode(ProxyResponseError{
-				Success: false,
-				Code:    http.StatusUnauthorized,
-				Message: http.StatusText(http.StatusUnauthorized),
-			})
-			if err != nil {
-				return
-			}
+			RespondWithError(w, http.StatusUnauthorized, http.StatusText(http.StatusUnauthorized), basicAuth.ErrorInterceptor)
 			return
 		}
 
@@ -195,16 +134,7 @@ func (basicAuth AuthBasic) AuthMiddleware(next http.Handler) http.Handler {
 		pair := strings.SplitN(string(payload), ":", 2)
 		if len(pair) != 2 || pair[0] != basicAuth.Username || pair[1] != basicAuth.Password {
 			w.Header().Set("WWW-Authenticate", `Basic realm="Restricted"`)
-			w.Header().Set("Content-Type", "application/json")
-			w.WriteHeader(http.StatusUnauthorized)
-			err := json.NewEncoder(w).Encode(ProxyResponseError{
-				Success: false,
-				Code:    http.StatusUnauthorized,
-				Message: http.StatusText(http.StatusUnauthorized),
-			})
-			if err != nil {
-				return
-			}
+			RespondWithError(w, http.StatusUnauthorized, http.StatusText(http.StatusUnauthorized), basicAuth.ErrorInterceptor)
 			return
 		}
 

internal/middleware/rate-limit.go

@@ -16,7 +16,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 import (
-	"encoding/json"
+	"fmt"
 	"github.com/gorilla/mux"
 	"github.com/jkaninda/goma-gateway/pkg/logger"
 	"net/http"
@@ -28,20 +28,17 @@ func (rl *TokenRateLimiter) RateLimitMiddleware() mux.MiddlewareFunc {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			if !rl.Allow() {
+				logger.Error("Too many requests from IP: %s %s %s", getRealIP(r), r.URL, r.UserAgent())
+				//RespondWithError(w, http.StatusUnauthorized, http.StatusText(http.StatusUnauthorized), basicAuth.ErrorInterceptor)
+
 				// Rate limit exceeded, return a 429 Too Many Requests response
-				w.Header().Set("Content-Type", "application/json")
 				w.WriteHeader(http.StatusTooManyRequests)
-				err := json.NewEncoder(w).Encode(ProxyResponseError{
-					Success: false,
-					Code:    http.StatusTooManyRequests,
-					Message: "Too many requests, API rate limit exceeded. Please try again later.",
-				})
+				_, err := w.Write([]byte(fmt.Sprintf("%d Too many requests, API rate limit exceeded. Please try again later", http.StatusTooManyRequests)))
 				if err != nil {
 					return
 				}
 				return
 			}
-
 			// Proceed to the next handler if rate limit is not exceeded
 			next.ServeHTTP(w, r)
 		})
@@ -66,21 +63,12 @@ func (rl *RateLimiter) RateLimitMiddleware() mux.MiddlewareFunc {
 			rl.mu.Unlock()
 
 			if client.RequestCount > rl.Requests {
-				logger.Debug("Too many requests from IP: %s %s %s", clientID, r.URL, r.UserAgent())
-				w.Header().Set("Content-Type", "application/json")
-				w.WriteHeader(http.StatusTooManyRequests)
+				logger.Error("Too many requests from IP: %s %s %s", clientID, r.URL, r.UserAgent())
 				//Update Origin Cors Headers
 				if allowedOrigin(rl.Origins, r.Header.Get("Origin")) {
 					w.Header().Set("Access-Control-Allow-Origin", r.Header.Get("Origin"))
 				}
-				err := json.NewEncoder(w).Encode(ProxyResponseError{
-					Success: false,
-					Code:    http.StatusTooManyRequests,
-					Message: "Too many requests, API rate limit exceeded. Please try again later.",
-				})
-				if err != nil {
-					return
-				}
+				RespondWithError(w, http.StatusTooManyRequests, fmt.Sprintf("%d Too many requests, API rate limit exceeded. Please try again later", http.StatusTooManyRequests), rl.ErrorInterceptor)
 				return
 			}
 			// Proceed to the next handler if rate limit is not exceeded

internal/middleware/route_error_interceptor.go

@@ -0,0 +1,58 @@
+package middleware
+
+/*
+ * Copyright 2024 Jonas Kaninda
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+import (
+	errorinterceptor "github.com/jkaninda/goma-gateway/pkg/error-interceptor"
+	"github.com/jkaninda/goma-gateway/pkg/logger"
+	"io"
+	"net/http"
+)
+
+// RouteErrorInterceptor contains backend status code errors to intercept
+type RouteErrorInterceptor struct {
+	Origins          []string
+	ErrorInterceptor errorinterceptor.ErrorInterceptor
+}
+
+// RouteErrorInterceptor Middleware intercepts backend route errors
+func (intercept RouteErrorInterceptor) RouteErrorInterceptor(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		rec := newResponseRecorder(w)
+		next.ServeHTTP(rec, r)
+		if canInterceptError(rec.statusCode, intercept.ErrorInterceptor.Errors) {
+			logger.Debug("Backend error")
+			logger.Error("An error occurred from the backend with the status code: %d", rec.statusCode)
+			//Update Origin Cors Headers
+			if allowedOrigin(intercept.Origins, r.Header.Get("Origin")) {
+				w.Header().Set("Access-Control-Allow-Origin", r.Header.Get("Origin"))
+			}
+			RespondWithError(w, rec.statusCode, http.StatusText(rec.statusCode), intercept.ErrorInterceptor)
+			return
+		} else {
+			// No error: write buffered response to client
+			w.WriteHeader(rec.statusCode)
+			_, err := io.Copy(w, rec.body)
+			if err != nil {
+				return
+			}
+			return
+
+		}
+
+	})
+}

internal/middleware/types.go

@@ -19,6 +19,7 @@ package middleware
 
 import (
 	"bytes"
+	errorinterceptor "github.com/jkaninda/goma-gateway/pkg/error-interceptor"
 	"net/http"
 	"sync"
 	"time"
@@ -26,11 +27,12 @@ import (
 
 // RateLimiter defines rate limit properties.
 type RateLimiter struct {
-	Requests  int
-	Window    time.Duration
-	ClientMap map[string]*Client
-	mu        sync.Mutex
-	Origins   []string
+	Requests         int
+	Window           time.Duration
+	ClientMap        map[string]*Client
+	mu               sync.Mutex
+	Origins          []string
+	ErrorInterceptor errorinterceptor.ErrorInterceptor
 }
 
 // Client stores request count and window expiration for each client.
@@ -67,11 +69,12 @@ type ProxyResponseError struct {
 
 // JwtAuth  stores JWT configuration
 type JwtAuth struct {
-	AuthURL         string
-	RequiredHeaders []string
-	Headers         map[string]string
-	Params          map[string]string
-	Origins         []string
+	AuthURL          string
+	RequiredHeaders  []string
+	Headers          map[string]string
+	Params           map[string]string
+	Origins          []string
+	ErrorInterceptor errorinterceptor.ErrorInterceptor
 }
 
 // AuthenticationMiddleware Define struct
@@ -82,17 +85,19 @@ type AuthenticationMiddleware struct {
 	Params          map[string]string
 }
 type AccessListMiddleware struct {
-	Path        string
-	Destination string
-	List        []string
+	Path             string
+	Destination      string
+	List             []string
+	ErrorInterceptor errorinterceptor.ErrorInterceptor
 }
 
 // AuthBasic contains Basic auth configuration
 type AuthBasic struct {
-	Username string
-	Password string
-	Headers  map[string]string
-	Params   map[string]string
+	Username         string
+	Password         string
+	Headers          map[string]string
+	Params           map[string]string
+	ErrorInterceptor errorinterceptor.ErrorInterceptor
 }
 
 // InterceptErrors contains backend status code errors to intercept
@@ -120,10 +125,11 @@ type Oauth struct {
 	// Scope specifies optional requested permissions.
 	Scopes []string
 	// contains filtered or unexported fields
-	State     string
-	Origins   []string
-	JWTSecret string
-	Provider  string
+	State            string
+	Origins          []string
+	JWTSecret        string
+	Provider         string
+	ErrorInterceptor errorinterceptor.ErrorInterceptor
 }
 type OauthEndpoint struct {
 	AuthURL     string