pub

pkg/middleware/logger.go

@@ -0,0 +1,41 @@
+package middleware
+
+import (
+	"log"
+	"net/http"
+	"time"
+)
+
+// responseWriter is a custom wrapper to capture the status code
+type responseWriter struct {
+	http.ResponseWriter
+	status int
+}
+
+func (rw *responseWriter) WriteHeader(code int) {
+	rw.status = code
+	rw.ResponseWriter.WriteHeader(code)
+}
+
+// Logger is a middleware that logs each incoming HTTP request
+func Logger() Middleware {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			start := time.Now()
+			
+			// Default status is 200 OK
+			rw := &responseWriter{ResponseWriter: w, status: http.StatusOK}
+			
+			next.ServeHTTP(rw, r)
+			
+			log.Printf(
+				"[%s] %s %s | %d | %v",
+				r.Method,
+				r.Host,
+				r.URL.Path,
+				rw.status,
+				time.Since(start),
+			)
+		})
+	}
+}

pkg/middleware/middleware.go

@@ -0,0 +1,15 @@
+package middleware
+
+import "net/http"
+
+// Middleware defines a function that wraps an http.Handler
+type Middleware func(http.Handler) http.Handler
+
+// Chain constructs a new http.Handler by chaining multiple Middlewares around a final http.Handler
+func Chain(handler http.Handler, middlewares ...Middleware) http.Handler {
+	// Loop in reverse to preserve the order of execution
+	for i := len(middlewares) - 1; i >= 0; i-- {
+		handler = middlewares[i](handler)
+	}
+	return handler
+}

pkg/middleware/ratelimit.go

@@ -0,0 +1,92 @@
+package middleware
+
+import (
+	"net"
+	"net/http"
+	"sync"
+	"time"
+
+	"golang.org/x/time/rate"
+)
+
+// visitor defines the rate limiter for a single IP
+type visitor struct {
+	limiter  *rate.Limiter
+	lastSeen time.Time
+}
+
+// IPRateLimiter holds a map of rate limiters per IP
+type IPRateLimiter struct {
+	mu       sync.Mutex
+	visitors map[string]*visitor
+	rate     rate.Limit
+	burst    int
+}
+
+func NewIPRateLimiter(r rate.Limit, b int) *IPRateLimiter {
+	i := &IPRateLimiter{
+		visitors: make(map[string]*visitor),
+		rate:     r,
+		burst:    b,
+	}
+
+	// Background routine to clean up old visitors
+	go func() {
+		for {
+			time.Sleep(time.Minute)
+			i.cleanupVisitors()
+		}
+	}()
+
+	return i
+}
+
+func (i *IPRateLimiter) getLimiter(ip string) *rate.Limiter {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
+	v, exists := i.visitors[ip]
+	if !exists {
+		limiter := rate.NewLimiter(i.rate, i.burst)
+		i.visitors[ip] = &visitor{
+			limiter:  limiter,
+			lastSeen: time.Now(),
+		}
+		return limiter
+	}
+
+	v.lastSeen = time.Now()
+	return v.limiter
+}
+
+func (i *IPRateLimiter) cleanupVisitors() {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+	for ip, v := range i.visitors {
+		if time.Since(v.lastSeen) > 3*time.Minute {
+			delete(i.visitors, ip)
+		}
+	}
+}
+
+// RateLimit is a middleware that limits requests per IP
+func RateLimit(requestsPerSecond float64, burst int) Middleware {
+	limiter := NewIPRateLimiter(rate.Limit(requestsPerSecond), burst)
+
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			// Extract IP
+			ip, _, err := net.SplitHostPort(r.RemoteAddr)
+			if err != nil {
+				ip = r.RemoteAddr
+			}
+
+			if !limiter.getLimiter(ip).Allow() {
+				http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
+				return
+			}
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}

pkg/middleware/recovery.go

@@ -0,0 +1,22 @@
+package middleware
+
+import (
+	"log"
+	"net/http"
+)
+
+// Recovery is a middleware that recovers from any panics and writes a 500 if there was one
+func Recovery() Middleware {
+	return func(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 RECOVERED: %v", err)
+					http.Error(w, "Internal Server Error", http.StatusInternalServerError)
+				}
+			}()
+			
+			next.ServeHTTP(w, r)
+		})
+	}
+}