databasus--databasus/backend/internal/features/verification/agents/middleware.go

package verification_agents

import (
	"net/http"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
)

const (
	agentContextKey = "verification_agent"

	rateLimitAgentEndpoint = "verification-agent-id"
	rateLimitAgentMax      = 10
	rateLimitAgentWindow   = time.Second

	genericAuthError = "invalid agent credentials"
)

func (s *AgentService) RequireAgentAuth() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		clientIP := ctx.ClientIP()

		agentID, err := uuid.Parse(ctx.Param("agentId"))
		if err != nil {
			s.logger.Warn("verification agent auth failure",
				"client_ip", clientIP, "reason", "invalid_uuid")
			ctx.AbortWithStatusJSON(http.StatusUnauthorized,
				gin.H{"error": genericAuthError})

			return
		}

		isAllowed, rateLimitErr := s.rateLimiter.CheckLimit(
			agentID.String(), rateLimitAgentEndpoint,
			rateLimitAgentMax, rateLimitAgentWindow,
		)
		if rateLimitErr != nil {
			s.logger.Error("verification agent rate limit check failed",
				"error", rateLimitErr, "agent_id", agentID, "client_ip", clientIP)
		} else if !isAllowed {
			s.logger.Warn("verification agent per-agent rate limit hit",
				"agent_id", agentID, "client_ip", clientIP)
			ctx.AbortWithStatusJSON(http.StatusTooManyRequests,
				gin.H{"error": "too many requests"})

			return
		}

		header := ctx.GetHeader("Authorization")
		token := strings.TrimPrefix(header, "Bearer ")
		if token == "" || token == header {
			s.logger.Warn("verification agent auth failure",
				"client_ip", clientIP, "agent_id", agentID, "reason", "missing_token")
			ctx.AbortWithStatusJSON(http.StatusUnauthorized,
				gin.H{"error": genericAuthError})

			return
		}

		agent, err := s.VerifyAgentCredentials(agentID, token)
		if err != nil {
			s.logger.Warn("verification agent auth failure",
				"client_ip", clientIP, "agent_id", agentID, "reason", "invalid_credentials")
			ctx.AbortWithStatusJSON(http.StatusUnauthorized,
				gin.H{"error": genericAuthError})

			return
		}

		ctx.Set(agentContextKey, agent)
		ctx.Next()
	}
}

func GetAgentFromContext(ctx *gin.Context) (*Agent, bool) {
	v, exists := ctx.Get(agentContextKey)
	if !exists {
		return nil, false
	}

	agent, ok := v.(*Agent)

	return agent, ok
}