OrchestrAD/backend/internal/api/middleware.go

// Package api - Middleware for authentication and authorization
package api

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

	"github.com/Grace-Solutions/OrchestrAD/internal/auth"
	"github.com/Grace-Solutions/OrchestrAD/internal/models"
)

type contextKey string

const userContextKey contextKey = "user"

// AuthMiddleware validates session tokens
func AuthMiddleware(authService *auth.Service) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			token := extractToken(r)
			if token == "" {
				WriteError(w, http.StatusUnauthorized, ErrCodeUnauthorized, "Authentication required")
				return
			}

			user, err := authService.ValidateSession(token)
			if err != nil {
				WriteError(w, http.StatusUnauthorized, ErrCodeUnauthorized, "Invalid or expired token")
				return
			}

			ctx := context.WithValue(r.Context(), userContextKey, user)
			next.ServeHTTP(w, r.WithContext(ctx))
		})
	}
}

// OptionalAuthMiddleware adds user to context if authenticated, but doesn't require it
func OptionalAuthMiddleware(authService *auth.Service) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			token := extractToken(r)
			if token != "" {
				user, err := authService.ValidateSession(token)
				if err == nil {
					ctx := context.WithValue(r.Context(), userContextKey, user)
					r = r.WithContext(ctx)
				}
			}
			next.ServeHTTP(w, r)
		})
	}
}

// RequireRole creates middleware that requires a specific role
func RequireRole(roles ...string) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			user := GetUserFromContext(r.Context())
			if user == nil {
				WriteError(w, http.StatusUnauthorized, ErrCodeUnauthorized, "Authentication required")
				return
			}

			if !hasAnyRole(user, roles) {
				WriteError(w, http.StatusForbidden, ErrCodeForbidden, "Insufficient permissions")
				return
			}

			next.ServeHTTP(w, r)
		})
	}
}

// RequireAdmin is a convenience middleware for admin-only routes
func RequireAdmin() func(http.Handler) http.Handler {
	return RequireRole("SuperAdmin", "Admin")
}

// RequireOperator is a convenience middleware for operator-level routes
func RequireOperator() func(http.Handler) http.Handler {
	return RequireRole("SuperAdmin", "Admin", "Operator")
}

// GetUserFromContext retrieves the authenticated user from context
func GetUserFromContext(ctx context.Context) *models.User {
	user, ok := ctx.Value(userContextKey).(*models.User)
	if !ok {
		return nil
	}
	return user
}

func extractToken(r *http.Request) string {
	// Check Authorization header
	auth := r.Header.Get("Authorization")
	if strings.HasPrefix(auth, "Bearer ") {
		return auth[7:]
	}
	
	// Check X-API-Key header for API key auth
	if apiKey := r.Header.Get("X-API-Key"); apiKey != "" {
		return apiKey
	}
	
	// Check cookie (for browser sessions)
	cookie, err := r.Cookie("session")
	if err == nil && cookie.Value != "" {
		return cookie.Value
	}
	
	return ""
}

func hasAnyRole(user *models.User, roles []string) bool {
	for _, userRole := range user.Roles {
		for _, required := range roles {
			if userRole.Name == required {
				return true
			}
		}
	}
	return false
}

// CSRFMiddleware validates CSRF tokens for mutation requests
func CSRFMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Skip CSRF for GET, HEAD, OPTIONS
		if r.Method == "GET" || r.Method == "HEAD" || r.Method == "OPTIONS" {
			next.ServeHTTP(w, r)
			return
		}
		
		// Skip CSRF for API key auth
		if r.Header.Get("X-API-Key") != "" {
			next.ServeHTTP(w, r)
			return
		}
		
		// Skip CSRF for Bearer token auth (typically from NextAuth)
		if strings.HasPrefix(r.Header.Get("Authorization"), "Bearer ") {
			next.ServeHTTP(w, r)
			return
		}
		
		// TODO: Validate CSRF token from header/body
		csrfToken := r.Header.Get("X-CSRF-Token")
		if csrfToken == "" {
			WriteError(w, http.StatusForbidden, ErrCodeForbidden, "CSRF token required")
			return
		}
		
		// TODO: Validate the token against stored session
		next.ServeHTTP(w, r)
	})
}