/backend/internal/features/verification/agents/service.go

package verification_agents

import (
	"crypto/sha256"
	"crypto/subtle"
	"fmt"
	"log/slog"
	"strings"
	"time"

	"github.com/google/uuid"

	audit_logs "databasus-backend/internal/features/audit_logs"
	users_models "databasus-backend/internal/features/users/models"
	cache_utils "databasus-backend/internal/util/cache"
)

type AgentService struct {
	agentRepository *AgentRepository
	auditLogService *audit_logs.AuditLogService
	rateLimiter     *cache_utils.RateLimiter
	logger          *slog.Logger

	heartbeatListeners []AgentHeartbeatedListener
}

func (s *AgentService) AddAgentHeartbeatedListener(listener AgentHeartbeatedListener) {
	s.heartbeatListeners = append(s.heartbeatListeners, listener)
}

func (s *AgentService) CreateAgent(
	user *users_models.User,
	req *CreateAgentRequest,
) (*CreatedAgentResponse, error) {
	plainToken := newPlainToken()
	tokenHash := hashAgentToken(plainToken)

	agent := &Agent{
		ID:        uuid.New(),
		Name:      req.Name,
		TokenHash: tokenHash,
		CreatedAt: time.Now().UTC(),
	}

	if err := s.agentRepository.Create(agent); err != nil {
		return nil, err
	}

	s.auditLogService.WriteAuditLog(
		fmt.Sprintf("verification agent created: %s", agent.Name),
		&user.ID,
		nil,
	)

	return &CreatedAgentResponse{
		Agent: agent,
		Token: plainToken,
	}, nil
}

func (s *AgentService) ListAgents() ([]*Agent, error) {
	return s.agentRepository.FindAll()
}

func (s *AgentService) CountLiveAgents() (int64, error) {
	return s.agentRepository.CountLive()
}

func (s *AgentService) GetStaleAgents(threshold time.Duration) ([]*Agent, error) {
	return s.agentRepository.FindStale(time.Now().UTC().Add(-threshold))
}

func (s *AgentService) RotateToken(
	user *users_models.User,
	agentID uuid.UUID,
) (string, error) {
	agent, err := s.agentRepository.FindByID(agentID)
	if err != nil {
		return "", err
	}
	if agent == nil || agent.IsDeleted() {
		return "", ErrAgentNotFound
	}

	plainToken := newPlainToken()
	tokenHash := hashAgentToken(plainToken)

	if err := s.agentRepository.UpdateTokenHash(agentID, tokenHash); err != nil {
		return "", err
	}

	s.auditLogService.WriteAuditLog(
		fmt.Sprintf("verification agent token rotated: %s", agent.Name),
		&user.ID,
		nil,
	)

	return plainToken, nil
}

func (s *AgentService) DeleteAgent(user *users_models.User, agentID uuid.UUID) error {
	agent, err := s.agentRepository.FindByID(agentID)
	if err != nil {
		return err
	}
	if agent == nil || agent.IsDeleted() {
		return ErrAgentNotFound
	}

	if err := s.agentRepository.SoftDelete(agentID, time.Now().UTC()); err != nil {
		return err
	}

	s.auditLogService.WriteAuditLog(
		fmt.Sprintf("verification agent deleted: %s", agent.Name),
		&user.ID,
		nil,
	)

	return nil
}

func (s *AgentService) VerifyAgentCredentials(
	agentID uuid.UUID,
	rawToken string,
) (*Agent, error) {
	if rawToken == "" {
		return nil, ErrInvalidAgentToken
	}

	agent, err := s.agentRepository.FindByID(agentID)
	if err != nil {
		return nil, err
	}
	if agent == nil || agent.IsDeleted() {
		return nil, ErrInvalidAgentToken
	}

	incomingHash := hashAgentToken(rawToken)
	if subtle.ConstantTimeCompare([]byte(agent.TokenHash), []byte(incomingHash)) != 1 {
		return nil, ErrInvalidAgentToken
	}

	return agent, nil
}

func (s *AgentService) GetAgentByID(id uuid.UUID) (*Agent, error) {
	return s.agentRepository.FindByID(id)
}

func (s *AgentService) Heartbeat(
	agent *Agent,
	req *HeartbeatRequest,
) (seenAt time.Time, abortVerificationIDs []uuid.UUID, err error) {
	now := time.Now().UTC()

	err = s.agentRepository.UpdateCapacityAndLastSeen(
		agent.ID,
		AgentCapacity{
			MaxCPU:            req.MaxCPU,
			MaxRAMGb:          req.MaxRAMGb,
			MaxDiskGb:         req.MaxDiskGb,
			MaxConcurrentJobs: req.MaxConcurrentJobs,
		},
		now,
	)
	if err != nil {
		return time.Time{}, nil, err
	}

	abortSet := make(map[uuid.UUID]struct{})
	for _, listener := range s.heartbeatListeners {
		listenerAborts, listenerErr := listener.OnAgentHeartbeated(agent, req.CurrentVerificationIDs)
		if listenerErr != nil {
			s.logger.Error(
				"agent heartbeated listener failed",
				"error", listenerErr,
				"agent_id", agent.ID,
			)

			continue
		}

		for _, id := range listenerAborts {
			abortSet[id] = struct{}{}
		}
	}

	abortVerificationIDs = make([]uuid.UUID, 0, len(abortSet))
	for id := range abortSet {
		abortVerificationIDs = append(abortVerificationIDs, id)
	}

	return now, abortVerificationIDs, nil
}

func newPlainToken() string {
	return strings.ReplaceAll(uuid.New().String(), "-", "")
}

func hashAgentToken(token string) string {
	hash := sha256.Sum256([]byte(token))
	return fmt.Sprintf("%x", hash)
}