v2

2025-08-22 18:57:40 -04:00
parent d648a55c0c
commit df567983c1
20 changed files with 4519 additions and 8 deletions
--- a/internal/handlers/saml.go
+++ b/internal/handlers/saml.go
@ -0,0 +1,352 @@
+package handlers
+
+import (
+	"encoding/json"
+	"net/http"
+	"time"
+
+	"github.com/gorilla/mux"
+	"go.uber.org/zap"
+
+	"github.com/kms/api-key-service/internal/auth"
+	"github.com/kms/api-key-service/internal/config"
+	"github.com/kms/api-key-service/internal/domain"
+	"github.com/kms/api-key-service/internal/errors"
+	"github.com/kms/api-key-service/internal/services"
+)
+
+// SAMLHandler handles SAML authentication endpoints
+type SAMLHandler struct {
+	samlProvider   *auth.SAMLProvider
+	sessionService services.SessionService
+	authService    services.AuthenticationService
+	tokenService   services.TokenService
+	config         config.ConfigProvider
+	logger         *zap.Logger
+}
+
+// NewSAMLHandler creates a new SAML handler
+func NewSAMLHandler(
+	config config.ConfigProvider,
+	sessionService services.SessionService,
+	authService services.AuthenticationService,
+	tokenService services.TokenService,
+	logger *zap.Logger,
+) (*SAMLHandler, error) {
+	samlProvider, err := auth.NewSAMLProvider(config, logger)
+	if err != nil {
+		return nil, err
+	}
+
+	return &SAMLHandler{
+		samlProvider:   samlProvider,
+		sessionService: sessionService,
+		authService:    authService,
+		config:         config,
+		logger:         logger,
+	}, nil
+}
+
+// RegisterRoutes registers SAML routes
+func (h *SAMLHandler) RegisterRoutes(router *mux.Router) {
+	// SAML endpoints
+	router.HandleFunc("/auth/saml/login", h.InitiateSAMLLogin).Methods("GET")
+	router.HandleFunc("/auth/saml/acs", h.HandleSAMLResponse).Methods("POST")
+	router.HandleFunc("/auth/saml/metadata", h.GetServiceProviderMetadata).Methods("GET")
+	router.HandleFunc("/auth/saml/slo", h.HandleSingleLogout).Methods("GET", "POST")
+}
+
+// InitiateSAMLLogin initiates SAML authentication
+func (h *SAMLHandler) InitiateSAMLLogin(w http.ResponseWriter, r *http.Request) {
+	if !h.config.GetBool("SAML_ENABLED") {
+		h.writeErrorResponse(w, errors.NewConfigurationError("SAML authentication is not enabled"))
+		return
+	}
+
+	// Get query parameters
+	appID := r.URL.Query().Get("app_id")
+	redirectURL := r.URL.Query().Get("redirect_url")
+	
+	if appID == "" {
+		h.writeErrorResponse(w, errors.NewValidationError("app_id parameter is required"))
+		return
+	}
+
+	// Generate relay state with app_id and redirect_url
+	relayState := appID
+	if redirectURL != "" {
+		relayState += "|" + redirectURL
+	}
+
+	h.logger.Debug("Initiating SAML login", 
+		zap.String("app_id", appID),
+		zap.String("redirect_url", redirectURL))
+
+	// Generate SAML authentication request
+	authURL, requestID, err := h.samlProvider.GenerateAuthRequest(r.Context(), relayState)
+	if err != nil {
+		h.logger.Error("Failed to generate SAML auth request", zap.Error(err))
+		h.writeErrorResponse(w, err)
+		return
+	}
+
+	// Store request ID in session/cache for validation
+	// In production, you should store this securely
+	h.logger.Debug("Generated SAML auth request", 
+		zap.String("request_id", requestID),
+		zap.String("auth_url", authURL))
+
+	// Redirect to IdP
+	http.Redirect(w, r, authURL, http.StatusFound)
+}
+
+// HandleSAMLResponse handles SAML assertion consumer service (ACS)
+func (h *SAMLHandler) HandleSAMLResponse(w http.ResponseWriter, r *http.Request) {
+	if !h.config.GetBool("SAML_ENABLED") {
+		h.writeErrorResponse(w, errors.NewConfigurationError("SAML authentication is not enabled"))
+		return
+	}
+
+	h.logger.Debug("Handling SAML response")
+
+	// Parse form data
+	if err := r.ParseForm(); err != nil {
+		h.writeErrorResponse(w, errors.NewValidationError("Failed to parse form data").WithInternal(err))
+		return
+	}
+
+	samlResponse := r.FormValue("SAMLResponse")
+	relayState := r.FormValue("RelayState")
+
+	if samlResponse == "" {
+		h.writeErrorResponse(w, errors.NewValidationError("SAMLResponse is required"))
+		return
+	}
+
+	h.logger.Debug("Processing SAML response", zap.String("relay_state", relayState))
+
+	// Process SAML response
+	// In production, you should retrieve and validate the original request ID
+	authContext, err := h.samlProvider.ProcessSAMLResponse(r.Context(), samlResponse, "")
+	if err != nil {
+		h.logger.Error("Failed to process SAML response", zap.Error(err))
+		h.writeErrorResponse(w, err)
+		return
+	}
+
+	// Parse relay state to get app_id and redirect_url
+	appID, redirectURL := h.parseRelayState(relayState)
+	if appID == "" {
+		h.writeErrorResponse(w, errors.NewValidationError("Invalid relay state: missing app_id"))
+		return
+	}
+
+	// Create user session
+	sessionReq := &domain.CreateSessionRequest{
+		UserID:      authContext.UserID,
+		AppID:       appID,
+		SessionType: domain.SessionTypeWeb,
+		IPAddress:   h.getClientIP(r),
+		UserAgent:   r.UserAgent(),
+		ExpiresAt:   time.Now().Add(8 * time.Hour), // 8 hour session
+		Permissions: authContext.Permissions,
+		Claims:      authContext.Claims,
+	}
+
+	session, err := h.sessionService.CreateSession(r.Context(), sessionReq)
+	if err != nil {
+		h.logger.Error("Failed to create session", zap.Error(err))
+		h.writeErrorResponse(w, err)
+		return
+	}
+
+	// Generate JWT token for the session using the existing token service
+	userToken := &domain.UserToken{
+		AppID:       appID,
+		UserID:      authContext.UserID,
+		Permissions: authContext.Permissions,
+		IssuedAt:    time.Now(),
+		ExpiresAt:   session.ExpiresAt,
+		MaxValidAt:  session.ExpiresAt,
+		TokenType:   domain.TokenTypeUser,
+		Claims:      authContext.Claims,
+	}
+
+	tokenString, err := h.authService.GenerateJWTToken(r.Context(), userToken)
+	if err != nil {
+		h.logger.Error("Failed to create JWT token", zap.Error(err))
+		h.writeErrorResponse(w, err)
+		return
+	}
+
+	h.logger.Debug("SAML authentication successful", 
+		zap.String("user_id", authContext.UserID),
+		zap.String("session_id", session.ID.String()))
+
+	// If redirect URL is provided, redirect with token
+	if redirectURL != "" {
+		// Add token as query parameter or fragment
+		redirectURL += "?token=" + tokenString
+		http.Redirect(w, r, redirectURL, http.StatusFound)
+		return
+	}
+
+	// Otherwise, return JSON response
+	response := map[string]interface{}{
+		"success": true,
+		"token":   tokenString,
+		"user": map[string]interface{}{
+			"id":    authContext.UserID,
+			"email": authContext.Claims["email"],
+			"name":  authContext.Claims["name"],
+		},
+		"session_id": session.ID.String(),
+		"expires_at": session.ExpiresAt,
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(response)
+}
+
+// GetServiceProviderMetadata returns SP metadata XML
+func (h *SAMLHandler) GetServiceProviderMetadata(w http.ResponseWriter, r *http.Request) {
+	if !h.config.GetBool("SAML_ENABLED") {
+		h.writeErrorResponse(w, errors.NewConfigurationError("SAML authentication is not enabled"))
+		return
+	}
+
+	h.logger.Debug("Generating SP metadata")
+
+	metadata, err := h.samlProvider.GenerateServiceProviderMetadata()
+	if err != nil {
+		h.logger.Error("Failed to generate SP metadata", zap.Error(err))
+		h.writeErrorResponse(w, err)
+		return
+	}
+
+	w.Header().Set("Content-Type", "application/xml")
+	w.Write([]byte(metadata))
+}
+
+// HandleSingleLogout handles SAML single logout
+func (h *SAMLHandler) HandleSingleLogout(w http.ResponseWriter, r *http.Request) {
+	if !h.config.GetBool("SAML_ENABLED") {
+		h.writeErrorResponse(w, errors.NewConfigurationError("SAML authentication is not enabled"))
+		return
+	}
+
+	h.logger.Debug("Handling SAML single logout")
+
+	// Get session ID from query parameter or form
+	sessionID := r.URL.Query().Get("session_id")
+	if sessionID == "" && r.Method == "POST" {
+		r.ParseForm()
+		sessionID = r.FormValue("session_id")
+	}
+
+	if sessionID != "" {
+		// Revoke specific session
+		h.logger.Debug("Revoking session", zap.String("session_id", sessionID))
+		// Implementation would depend on how you store session IDs
+		// For now, we'll just log it
+	}
+
+	// In a full implementation, you would:
+	// 1. Parse the SAML LogoutRequest
+	// 2. Validate the request
+	// 3. Revoke the user's sessions
+	// 4. Generate a LogoutResponse
+	// 5. Redirect back to the IdP
+
+	// For now, return a simple success response
+	response := map[string]interface{}{
+		"success": true,
+		"message": "Logout successful",
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(response)
+}
+
+// parseRelayState parses the relay state to extract app_id and redirect_url
+func (h *SAMLHandler) parseRelayState(relayState string) (appID, redirectURL string) {
+	if relayState == "" {
+		return "", ""
+	}
+
+	// RelayState format: "app_id|redirect_url" or just "app_id"
+	parts := []string{relayState}
+	if len(relayState) > 0 && relayState[0] != '|' {
+		// Split on first pipe character
+		for i, char := range relayState {
+			if char == '|' {
+				parts = []string{relayState[:i], relayState[i+1:]}
+				break
+			}
+		}
+	}
+
+	appID = parts[0]
+	if len(parts) > 1 {
+		redirectURL = parts[1]
+	}
+
+	return appID, redirectURL
+}
+
+// getClientIP extracts the client IP address from the request
+func (h *SAMLHandler) getClientIP(r *http.Request) string {
+	// Check X-Forwarded-For header first
+	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
+		// Take the first IP if multiple are present
+		if idx := len(xff); idx > 0 {
+			for i, char := range xff {
+				if char == ',' {
+					return xff[:i]
+				}
+			}
+			return xff
+		}
+	}
+
+	// Check X-Real-IP header
+	if xri := r.Header.Get("X-Real-IP"); xri != "" {
+		return xri
+	}
+
+	// Fall back to RemoteAddr
+	return r.RemoteAddr
+}
+
+// writeErrorResponse writes an error response
+func (h *SAMLHandler) writeErrorResponse(w http.ResponseWriter, err error) {
+	var statusCode int
+	var errorCode string
+
+	switch {
+	case errors.IsValidationError(err):
+		statusCode = http.StatusBadRequest
+		errorCode = "VALIDATION_ERROR"
+	case errors.IsAuthenticationError(err):
+		statusCode = http.StatusUnauthorized
+		errorCode = "AUTHENTICATION_ERROR"
+	case errors.IsConfigurationError(err):
+		statusCode = http.StatusServiceUnavailable
+		errorCode = "CONFIGURATION_ERROR"
+	default:
+		statusCode = http.StatusInternalServerError
+		errorCode = "INTERNAL_ERROR"
+	}
+
+	response := map[string]interface{}{
+		"success": false,
+		"error": map[string]interface{}{
+			"code":    errorCode,
+			"message": err.Error(),
+		},
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(statusCode)
+	json.NewEncoder(w).Encode(response)
+}