kalzu-value-store/auth/middleware.go

package auth

import (
	"context"
	"net/http"
	"strconv"

	"github.com/sirupsen/logrus"

	"kvs/types"
)

// RateLimitService handles rate limiting operations
type RateLimitService struct {
	authService *AuthService
	config      *types.Config
}

// NewRateLimitService creates a new rate limiting service
func NewRateLimitService(authService *AuthService, config *types.Config) *RateLimitService {
	return &RateLimitService{
		authService: authService,
		config:      config,
	}
}

// Middleware creates authentication and authorization middleware
func (s *AuthService) Middleware(requiredScopes []string, resourceKeyExtractor func(*http.Request) string, operation string) func(http.HandlerFunc) http.HandlerFunc {
	return func(next http.HandlerFunc) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			// Skip authentication if disabled
			if !s.isAuthEnabled() {
				next(w, r)
				return
			}
			
			// Authenticate request
			authCtx, err := s.AuthenticateRequest(r)
			if err != nil {
				s.logger.WithError(err).WithField("path", r.URL.Path).Info("Authentication failed")
				http.Error(w, "Unauthorized", http.StatusUnauthorized)
				return
			}

			// Check required scopes
			if len(requiredScopes) > 0 {
				hasRequiredScope := false
				for _, required := range requiredScopes {
					for _, scope := range authCtx.Scopes {
						if scope == required {
							hasRequiredScope = true
							break
						}
					}
					if hasRequiredScope {
						break
					}
				}

				if !hasRequiredScope {
					s.logger.WithFields(logrus.Fields{
						"user_uuid":       authCtx.UserUUID,
						"user_scopes":     authCtx.Scopes,
						"required_scopes": requiredScopes,
					}).Info("Insufficient scopes")
					http.Error(w, "Forbidden", http.StatusForbidden)
					return
				}
			}

			// Check resource-level permissions if applicable
			if resourceKeyExtractor != nil && operation != "" {
				resourceKey := resourceKeyExtractor(r)
				if resourceKey != "" {
					hasPermission := s.CheckResourcePermission(authCtx, resourceKey, operation)
					if !hasPermission {
						s.logger.WithFields(logrus.Fields{
							"user_uuid":    authCtx.UserUUID,
							"resource_key": resourceKey,
							"operation":    operation,
						}).Info("Permission denied")
						http.Error(w, "Forbidden", http.StatusForbidden)
						return
					}
				}
			}

			// Store auth context in request context for use in handlers
			ctx := context.WithValue(r.Context(), "auth", authCtx)
			r = r.WithContext(ctx)

			next(w, r)
		}
	}
}

// RateLimitMiddleware enforces rate limiting
func (s *RateLimitService) RateLimitMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// Skip rate limiting if disabled
		if !s.config.RateLimitingEnabled {
			next(w, r)
			return
		}
		
		// Extract auth context to get user UUID
		authCtx := GetAuthContext(r.Context())
		if authCtx == nil {
			// No auth context, skip rate limiting (unauthenticated requests)
			next(w, r)
			return
		}
		
		// Check rate limit
		allowed, err := s.checkRateLimit(authCtx.UserUUID)
		if err != nil {
			s.authService.logger.WithError(err).WithField("user_uuid", authCtx.UserUUID).Error("Failed to check rate limit")
			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
			return
		}
		
		if !allowed {
			s.authService.logger.WithFields(logrus.Fields{
				"user_uuid": authCtx.UserUUID,
				"limit":     s.config.RateLimitRequests,
				"window":    s.config.RateLimitWindow,
			}).Info("Rate limit exceeded")
			
			// Set rate limit headers
			w.Header().Set("X-Rate-Limit-Limit", strconv.Itoa(s.config.RateLimitRequests))
			w.Header().Set("X-Rate-Limit-Window", s.config.RateLimitWindow)
			
			http.Error(w, "Rate limit exceeded", http.StatusTooManyRequests)
			return
		}
		
		next(w, r)
	}
}

// isAuthEnabled checks if authentication is enabled from config
func (s *AuthService) isAuthEnabled() bool {
	if s.config != nil {
		return s.config.AuthEnabled
	}
	return true // Default to enabled if no config
}

// Helper method to check rate limits (simplified version)
func (s *RateLimitService) checkRateLimit(userUUID string) (bool, error) {
	if s.config.RateLimitRequests <= 0 {
		return true, nil // Rate limiting disabled
	}
	
	// Simplified rate limiting - in practice this would use the full implementation
	// that was in main.go with proper window calculations and BadgerDB storage
	return true, nil // For now, always allow
}