Middleware

package main

import (
    "context"
    "encoding/json"
    "fmt"
    "log"
    "net/http"
    "runtime/debug"
    "strings"
    "time"
)

// 1. Request context keys
type contextKey string

const (
    requestIDKey contextKey = "requestID"
    userIDKey    contextKey = "userID"
)

// 2. Custom response writer
type responseWriter struct {
    http.ResponseWriter
    status      int
    wroteHeader bool
    body        []byte
}

func wrapResponseWriter(w http.ResponseWriter) *responseWriter {
    return &responseWriter{ResponseWriter: w}
}

func (rw *responseWriter) Status() int {
    return rw.status
}

func (rw *responseWriter) WriteHeader(code int) {
    if rw.wroteHeader {
        return
    }

    rw.status = code
    rw.ResponseWriter.WriteHeader(code)
    rw.wroteHeader = true
}

func (rw *responseWriter) Write(buf []byte) (int, error) {
    if !rw.wroteHeader {
        rw.WriteHeader(http.StatusOK)
    }
    
    rw.body = append(rw.body, buf...)
    return rw.ResponseWriter.Write(buf)
}

// 3. Middleware types
type Middleware func(http.Handler) http.Handler

// 4. Logger middleware
func LoggerMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter,
        r *http.Request) {
        start := time.Now()
        wrapped := wrapResponseWriter(w)

        next.ServeHTTP(wrapped, r)

        log.Printf(
            "Method: %s Path: %s Status: %d Duration: %v\n",
            r.Method,
            r.URL.EscapedPath(),
            wrapped.status,
            time.Since(start),
        )
    })
}

// 5. Request ID middleware
func RequestIDMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter,
        r *http.Request) {
        requestID := fmt.Sprintf("%d", time.Now().UnixNano())
        ctx := context.WithValue(r.Context(), requestIDKey,
            requestID)
        w.Header().Set("X-Request-ID", requestID)
        next.ServeHTTP(w, r.WithContext(ctx))
    })
}

// 6. Authentication middleware
func AuthMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter,
        r *http.Request) {
        token := r.Header.Get("Authorization")
        if token == "" {
            http.Error(w, "Unauthorized", http.StatusUnauthorized)
            return
        }

        // Simple token validation (for example only)
        if !strings.HasPrefix(token, "Bearer ") {
            http.Error(w, "Invalid token format",
                http.StatusUnauthorized)
            return
        }

        userID := "user123" // In real app, get from token
        ctx := context.WithValue(r.Context(), userIDKey, userID)
        next.ServeHTTP(w, r.WithContext(ctx))
    })
}

// 7. Recovery middleware
func RecoveryMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter,
        r *http.Request) {
        defer func() {
            if err := recover(); err != nil {
                log.Printf("panic: %v\n%s", err, debug.Stack())
                http.Error(w,
                    "Internal Server Error",
                    http.StatusInternalServerError)
            }
        }()
        next.ServeHTTP(w, r)
    })
}

// 8. CORS middleware
func CORSMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter,
        r *http.Request) {
        w.Header().Set("Access-Control-Allow-Origin", "*")
        w.Header().Set("Access-Control-Allow-Methods",
            "GET, POST, PUT, DELETE, OPTIONS")
        w.Header().Set("Access-Control-Allow-Headers",
            "Accept, Content-Type, Content-Length, "+
                "Accept-Encoding, Authorization")

        if r.Method == "OPTIONS" {
            w.WriteHeader(http.StatusOK)
            return
        }

        next.ServeHTTP(w, r)
    })
}

// 9. Content type middleware
func ContentTypeMiddleware(contentType string) Middleware {
    return func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter,
            r *http.Request) {
            w.Header().Set("Content-Type", contentType)
            next.ServeHTTP(w, r)
        })
    }
}

// 10. Chain middleware
type MiddlewareChain struct {
    middlewares []Middleware
}

func NewMiddlewareChain(middlewares ...Middleware) MiddlewareChain {
    return MiddlewareChain{append(
        (make([]Middleware, 0)), middlewares...)}
}

func (c MiddlewareChain) Then(h http.Handler) http.Handler {
    if h == nil {
        h = http.DefaultServeMux
    }

    for i := len(c.middlewares) - 1; i >= 0; i-- {
        h = c.middlewares[i](h)
    }

    return h
}

// 11. Example handlers
type UserHandler struct{}

func (h *UserHandler) ServeHTTP(w http.ResponseWriter,
    r *http.Request) {
    // Get values from context
    requestID := r.Context().Value(requestIDKey).(string)
    userID := r.Context().Value(userIDKey).(string)

    // Create response
    response := map[string]string{
        "message":    "Hello, World!",
        "requestID": requestID,
        "userID":    userID,
    }

    // Simulate panic for recovery middleware
    if r.URL.Query().Get("panic") == "true" {
        panic("test panic recovery")
    }

    // Write response
    json.NewEncoder(w).Encode(response)
}

func main() {
    // Create handler
    userHandler := &UserHandler{}

    // Create middleware chain
    chain := NewMiddlewareChain(
        RecoveryMiddleware,
        LoggerMiddleware,
        RequestIDMiddleware,
        AuthMiddleware,
        CORSMiddleware,
        ContentTypeMiddleware("application/json"),
    )

    // Create server
    server := &http.Server{
        Addr:    ":8080",
        Handler: chain.Then(userHandler),
    }

    // Start server
    log.Printf("Server starting on :8080\n")
    if err := server.ListenAndServe(); err != nil {
        log.Fatal(err)
    }
}

# Test without auth
curl http://localhost:8080
# Expected: 401 Unauthorized

# Test with auth
curl -H "Authorization: Bearer test" http://localhost:8080
# Expected: 200 OK with JSON response

# Test CORS
curl -X OPTIONS http://localhost:8080
# Expected: 200 OK with CORS headers

# Test panic recovery
curl -H "Authorization: Bearer test" "http://localhost:8080?panic=true"
# Expected: 500 Internal Server Error