gochat/internal/middleware/account_scope.go

package middleware

// Reference: P2E §1.3 — AccountScope middleware
// Extracts account_id from JWT claims and injects it plus the PolicyContext
// into the Gin context for downstream RBAC enforcement.
//
// IMPORTANT: This middleware package does NOT import internal/model or internal/service
// to avoid potential circular dependencies. Instead, it defines its own interface
// (RBACLookup) that the service layer implements. This follows Go's interface segregation
// pattern — the middleware defines what it needs, and the service provides it.

import (
	"net/http"
	"strconv"

	"github.com/gin-gonic/gin"

	"github.com/gochat/gochat/internal/auth"
	"github.com/gochat/gochat/pkg/response"
)

// AccountScope validates and injects account context from JWT claims or
// X-Account-ID header. It also constructs the PolicyContext from the
// AccountUser association so that downstream PolicyMiddleware can check permissions.
//
// Basic version (without service lookup) uses JWT claims directly.
// For production, use AccountScopeWithService() which looks up AccountUser from DB.
//
// Usage:
//   router.Use(AuthRequired(jwtSvc), AccountScope())
func AccountScope() gin.HandlerFunc {
	return func(c *gin.Context) {
		// Step 1: Get user_id from JWT claims (set by AuthRequired)
		userID, exists := c.Get("user_id")
		if !exists {
			response.AbortWithStatusError(c, http.StatusUnauthorized, response.ErrUnauthorized,
				"User ID not found in context — ensure AuthRequired middleware is applied first")
			return
		}

		// Step 2: Determine account_id
		// Priority: X-Account-ID header > JWT claims account_id
		accountID := getAccountID(c)
		if accountID == 0 {
			response.AbortWithStatusError(c, http.StatusBadRequest, response.ErrBadRequest,
				"Account ID required — provide via X-Account-ID header or JWT claims")
			return
		}

		// Step 3: Override JWT claims account_id if header was provided
		c.Set("account_id", accountID)

		// Step 4: Get role from JWT claims
		roleStr := "agent" // default fallback
		role, roleExists := c.Get("user_role")
		if roleExists {
			if s, ok := role.(string); ok {
				roleStr = s
			}
		}

		// Step 5: Get custom_role_id from JWT claims
		customRoleID := uint(0)
		claims, claimsExists := c.Get("auth_claims")
		if claimsExists {
			if jwtClaims, ok := claims.(*auth.Claims); ok {
				if jwtClaims.CustomRoleID != 0 {
					customRoleID = jwtClaims.CustomRoleID
				}
				if jwtClaims.Role != "" {
					roleStr = jwtClaims.Role
				}
			}
		}

		// Step 6: Build PolicyContext from JWT claims data
		// For custom_role, permissions will be nil (loaded by service layer)
		var permissions auth.PermissionMatrixMap
		policyCtx := auth.NewPolicyContext(
			userID.(uint),
			accountID,
			roleStr,
			customRoleID,
			permissions,
		)

		c.Set("policy_context", policyCtx)
		c.Next()
	}
}

// AccountScopeWithService creates an AccountScope middleware that uses
// a RBACLookup service to look up the user's AccountUser record and CustomRole
// permissions from the database. This is the production-grade version.
//
// Usage:
//   rbacSvc := service.NewRBACService(db)
//   router.Use(AuthRequired(jwtSvc), AccountScopeWithService(rbacSvc))
func AccountScopeWithService(lookup RBACLookup) gin.HandlerFunc {
	return func(c *gin.Context) {
		userID, exists := c.Get("user_id")
		if !exists {
			response.AbortWithStatusError(c, http.StatusUnauthorized, response.ErrUnauthorized,
				"User ID not found in context — ensure AuthRequired middleware is applied first")
			return
		}

		accountID := getAccountID(c)
		if accountID == 0 {
			response.AbortWithStatusError(c, http.StatusBadRequest, response.ErrBadRequest,
				"Account ID required — provide via X-Account-ID header or JWT claims")
			return
		}

		c.Set("account_id", accountID)

		// Look up AccountUser to get role and CustomRoleID
		accountUser, err := lookup.GetAccountUser(userID.(uint), accountID)
		if err != nil {
			response.AbortWithStatusError(c, http.StatusForbidden, response.ErrForbidden,
				"User does not belong to this account")
			return
		}

		// Build permissions matrix based on role
		permissions := auth.PermissionMatrixMap{}
		if accountUser.Role == "custom_role" && accountUser.CustomRoleID > 0 {
			pm, err := lookup.GetCustomRolePermissions(accountUser.CustomRoleID)
			if err != nil {
				// Fallback to agent defaults if custom role not found
				permissions = auth.AgentDefaultPermissions
			} else {
				permissions = pm
			}
		}

		policyCtx := auth.NewPolicyContext(
			userID.(uint),
			accountID,
			accountUser.Role,
			accountUser.CustomRoleID,
			permissions,
		)

		c.Set("policy_context", policyCtx)
		c.Next()
	}
}

// getAccountID extracts account ID from the request.
// Priority: X-Account-ID header > JWT claims account_id
func getAccountID(c *gin.Context) uint {
	// Try X-Account-ID header first (for account switching)
	headerAccountID := c.GetHeader("X-Account-ID")
	if headerAccountID != "" {
		id, err := strconv.ParseUint(headerAccountID, 10, 32)
		if err == nil && id != 0 {
			return uint(id)
		}
	}

	// Fall back to JWT claims (set by AuthRequired middleware)
	if accountID, exists := c.Get("account_id"); exists {
		switch v := accountID.(type) {
		case uint:
			if v != 0 {
				return v
			}
		case float64:
			if v != 0 {
				return uint(v)
			}
		case int:
			if v != 0 {
				return uint(v)
			}
		}
	}

	return 0
}

// --- RBACLookup Interface ---
// Defines what the middleware needs from the service layer.
// The service package implements this interface, avoiding circular imports.

// RBACLookup is the interface that the RBAC service must implement
// for use with AccountScopeWithService middleware.
type RBACLookup interface {
	GetAccountUser(userID, accountID uint) (*AccountUserRole, error)
	GetCustomRolePermissions(customRoleID uint) (auth.PermissionMatrixMap, error)
}

// AccountUserRole holds the essential AccountUser data needed by middleware.
// This is a lightweight DTO that the service populates from the model.AccountUser.
type AccountUserRole struct {
	UserID       uint
	AccountID    uint
	Role         string
	CustomRoleID uint
	Availability string
}