freedakgmail
105 lines
3.0 KiB
Go
105 lines
3.0 KiB
Go
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"runtime/debug"
|
|
"time"
|
|
|
|
"github.com/edgeai/gateway/internal/observability"
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
type contextKey string
|
|
|
|
const (
|
|
RequestIDKey contextKey = "request_id"
|
|
)
|
|
|
|
// RequestID middleware generates a unique request ID and sets it in context and response header.
|
|
func RequestID(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
requestID := r.Header.Get("X-Request-ID")
|
|
if requestID == "" {
|
|
requestID = uuid.New().String()
|
|
}
|
|
w.Header().Set("X-Request-ID", requestID)
|
|
ctx := context.WithValue(r.Context(), RequestIDKey, requestID)
|
|
next.ServeHTTP(w, r.WithContext(ctx))
|
|
})
|
|
}
|
|
|
|
// BodyLimit middleware rejects requests with bodies exceeding the given size.
|
|
func BodyLimit(maxMB int) func(http.Handler) http.Handler {
|
|
maxBytes := int64(maxMB) * 1024 * 1024
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.ContentLength > maxBytes {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusBadRequest)
|
|
fmt.Fprintf(w, `{"error":{"code":"INVALID_REQUEST","message":"request body exceeds %dMB limit"}}`, maxMB)
|
|
return
|
|
}
|
|
r.Body = http.MaxBytesReader(w, r.Body, maxBytes)
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
}
|
|
|
|
// Recovery middleware catches panics and returns 500.
|
|
func Recovery(logger *observability.Logger) func(http.Handler) http.Handler {
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
defer func() {
|
|
if rec := recover(); rec != nil {
|
|
logger.Error("panic recovered",
|
|
observability.F().Event("panic").
|
|
RequestID(r.Header.Get("X-Request-ID")).
|
|
Reason(fmt.Sprintf("%v\n%s", rec, debug.Stack())))
|
|
http.Error(w, `{"error":{"code":"INTERNAL_ERROR","message":"internal server error"}}`,
|
|
http.StatusInternalServerError)
|
|
}
|
|
}()
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
}
|
|
|
|
// Logging middleware logs request method, path, status, and duration.
|
|
func Logging(logger *observability.Logger) func(http.Handler) http.Handler {
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
start := time.Now()
|
|
rw := &responseWriter{ResponseWriter: w, status: 200}
|
|
next.ServeHTTP(rw, r)
|
|
logger.Info("http request",
|
|
observability.F().
|
|
Event("http_request").
|
|
RequestID(r.Header.Get("X-Request-ID")).
|
|
Set("method", r.Method).
|
|
Set("path", r.URL.Path).
|
|
Set("status", rw.status).
|
|
Set("duration_ms", time.Since(start).Milliseconds()))
|
|
})
|
|
}
|
|
}
|
|
|
|
type responseWriter struct {
|
|
http.ResponseWriter
|
|
status int
|
|
}
|
|
|
|
func (rw *responseWriter) WriteHeader(code int) {
|
|
rw.status = code
|
|
rw.ResponseWriter.WriteHeader(code)
|
|
}
|
|
|
|
// GetRequestID extracts the request ID from context.
|
|
func GetRequestID(ctx context.Context) string {
|
|
if v, ok := ctx.Value(RequestIDKey).(string); ok {
|
|
return v
|
|
}
|
|
return ""
|
|
}
|