breakpilot-compliance/breakpilot-compliance-sdk/services/api-gateway/internal/middleware/auth.go

// Package middleware provides HTTP middleware for the API Gateway
package middleware

import (
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/golang-jwt/jwt/v5"
)

// Claims represents JWT claims
type Claims struct {
	TenantID string   `json:"tenant_id"`
	UserID   string   `json:"user_id"`
	Scopes   []string `json:"scopes"`
	jwt.RegisteredClaims
}

// Auth middleware validates JWT tokens or API keys
func Auth(jwtSecret string) gin.HandlerFunc {
	return func(c *gin.Context) {
		authHeader := c.GetHeader("Authorization")
		if authHeader == "" {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"error": "Missing authorization header",
			})
			return
		}

		// Check for Bearer token (JWT)
		if strings.HasPrefix(authHeader, "Bearer ") {
			tokenString := strings.TrimPrefix(authHeader, "Bearer ")

			// Check if it's an API key (starts with pk_ or sk_)
			if strings.HasPrefix(tokenString, "pk_") || strings.HasPrefix(tokenString, "sk_") {
				// API Key authentication
				if err := validateAPIKey(c, tokenString); err != nil {
					c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
						"error": "Invalid API key",
					})
					return
				}
			} else {
				// JWT authentication
				claims, err := validateJWT(tokenString, jwtSecret)
				if err != nil {
					c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
						"error": "Invalid token",
					})
					return
				}

				// Store claims in context
				c.Set("tenant_id", claims.TenantID)
				c.Set("user_id", claims.UserID)
				c.Set("scopes", claims.Scopes)
			}
		} else {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"error": "Invalid authorization format",
			})
			return
		}

		c.Next()
	}
}

func validateJWT(tokenString, secret string) (*Claims, error) {
	token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (interface{}, error) {
		return []byte(secret), nil
	})

	if err != nil {
		return nil, err
	}

	if claims, ok := token.Claims.(*Claims); ok && token.Valid {
		return claims, nil
	}

	return nil, jwt.ErrTokenInvalidClaims
}

func validateAPIKey(c *gin.Context, apiKey string) error {
	// In production, this would validate against a database of API keys
	// For now, we extract tenant info from the X-Tenant-ID header
	tenantID := c.GetHeader("X-Tenant-ID")
	if tenantID == "" {
		tenantID = "default"
	}

	c.Set("tenant_id", tenantID)
	c.Set("user_id", "api-key-user")
	c.Set("scopes", []string{"read", "write"})

	return nil
}

// TenantIsolation ensures requests only access their own tenant's data
func TenantIsolation() gin.HandlerFunc {
	return func(c *gin.Context) {
		tenantID, exists := c.Get("tenant_id")
		if !exists || tenantID == "" {
			c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
				"error": "Tenant ID required",
			})
			return
		}

		// Add tenant ID to all database queries (handled by handlers)
		c.Set("tenant_filter", tenantID)

		c.Next()
	}
}