Aether/backend/pkg/middleware/proxy.go

package middleware

import (
	"context"
	"net"
	"net/http"
	"strings"

	"aether/pkg/config"
)

type contextKey int

const (
	keyRealIP  contextKey = iota
	keyIsHTTPS contextKey = iota
)

// RealIP returns the real client IP stored in context by TrustedProxy middleware.
// Falls back to r.RemoteAddr if the middleware was not applied.
func RealIP(r *http.Request) string {
	if ip, ok := r.Context().Value(keyRealIP).(string); ok && ip != "" {
		return ip
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

// IsHTTPS reports whether the original client connection was HTTPS,
// either directly or via a trusted proxy's X-Forwarded-Proto header.
func IsHTTPS(r *http.Request) bool {
	if v, ok := r.Context().Value(keyIsHTTPS).(bool); ok {
		return v
	}
	return r.TLS != nil
}

// TrustedProxy extracts the real client IP and protocol from X-Forwarded-*
// headers, but only when the immediate peer is in the trusted proxy list.
// Untrusted peers never override RemoteAddr.
func TrustedProxy(cfg *config.Config) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			remoteHost, _, err := net.SplitHostPort(r.RemoteAddr)
			if err != nil {
				remoteHost = r.RemoteAddr
			}
			peerIP := net.ParseIP(remoteHost)

			realIP := remoteHost
			isHTTPS := r.TLS != nil

			if peerIP != nil && cfg.IsTrustedProxy(peerIP) {
				if xfp := r.Header.Get("X-Forwarded-Proto"); xfp != "" {
					isHTTPS = strings.EqualFold(strings.TrimSpace(xfp), "https")
				}
				// X-Forwarded-For may be a comma-separated chain; the leftmost
				// entry is the original client (set by the first proxy).
				if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
					parts := strings.SplitN(xff, ",", 2)
					if candidate := strings.TrimSpace(parts[0]); candidate != "" {
						realIP = candidate
					}
				}
			}

			ctx := context.WithValue(r.Context(), keyRealIP, realIP)
			ctx = context.WithValue(ctx, keyIsHTTPS, isHTTPS)
			next.ServeHTTP(w, r.WithContext(ctx))
		})
	}
}