easyai-ai-gateway/apps/api/internal/securityevents/service.go

package securityevents

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/sha256"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

type StateRepository interface {
	EnsureSecurityEventStreamState(context.Context, string, string, string) error
	BeginSecurityEventVerification(context.Context, string, string, []byte, time.Time) error
	RecordSecurityEventHeartbeatFailure(context.Context, string, string, string) error
	AdvanceSecurityEventStreamState(context.Context, string, string, time.Time, time.Duration) error
	EvaluateOIDCSecurityEvent(context.Context, string, string, string, string, string, time.Time, time.Time, time.Duration) (store.SecurityEventEvaluation, error)
}

type ServiceConfig struct {
	Issuer             string
	Audience           string
	StreamID           string
	ManagementTokenURL string
	ManagementClientID string
	ManagementSecret   string
	HeartbeatInterval  time.Duration
	StaleAfter         time.Duration
	HTTPClient         *http.Client
}

type Service struct {
	repository StateRepository
	config     ServiceConfig
	client     *http.Client
	clock      func() time.Time
	tokenMu    sync.Mutex
	token      string
	tokenUntil time.Time
	metrics    *Metrics
}

func (s *Service) SetMetrics(metrics *Metrics) { s.metrics = metrics }

func NewService(repository StateRepository, config ServiceConfig) (*Service, error) {
	if repository == nil || config.Issuer == "" || config.Audience == "" || config.StreamID == "" ||
		config.ManagementTokenURL == "" || config.ManagementClientID == "" || config.ManagementSecret == "" {
		return nil, errors.New("SSF heartbeat configuration is incomplete")
	}
	if config.HeartbeatInterval <= 0 || config.StaleAfter < 2*config.HeartbeatInterval {
		return nil, errors.New("SSF heartbeat timing is invalid")
	}
	client := config.HTTPClient
	if client == nil {
		transport := http.DefaultTransport.(*http.Transport).Clone()
		transport.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS12}
		transport.DisableCompression = true
		client = &http.Client{Timeout: 10 * time.Second,
			Transport:     transport,
			CheckRedirect: func(_ *http.Request, _ []*http.Request) error { return http.ErrUseLastResponse }}
	}
	return &Service{repository: repository, config: config, client: client, clock: time.Now}, nil
}

func (s *Service) Start(ctx context.Context) error {
	if err := s.Initialize(ctx); err != nil {
		return err
	}
	go s.Run(ctx)
	return nil
}

func (s *Service) Initialize(ctx context.Context) error {
	return s.repository.EnsureSecurityEventStreamState(ctx, s.config.Issuer, s.config.Audience, s.config.StreamID)
}

func (s *Service) Run(ctx context.Context) { s.run(ctx) }

// RequestVerification asks for an immediate verification without changing the
// remote stream status. In particular, this never re-enables a stream paused
// by an authentication-center administrator.
func (s *Service) RequestVerification(ctx context.Context) { s.sendVerification(ctx) }

func (s *Service) Evaluate(ctx context.Context, identity auth.OIDCSecurityEventIdentity) (auth.OIDCSecurityEventEvaluation, error) {
	result, err := s.repository.EvaluateOIDCSecurityEvent(
		ctx, s.config.Issuer, s.config.Audience, identity.Issuer, identity.TenantID, identity.Subject,
		identity.IssuedAt, s.clock().UTC(), s.config.StaleAfter,
	)
	if err != nil {
		return auth.OIDCSecurityEventEvaluation{}, err
	}
	if result.Revoked && s.metrics != nil {
		s.metrics.ObserveWatermarkRejection()
	}
	return auth.OIDCSecurityEventEvaluation{Revoked: result.Revoked, RequireIntrospection: result.RequireIntrospection}, nil
}

func (s *Service) run(ctx context.Context) {
	s.sendVerification(ctx)
	ticker := time.NewTicker(s.config.HeartbeatInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			s.sendVerification(ctx)
		}
	}
}

func (s *Service) sendVerification(ctx context.Context) {
	outcome := "failed"
	defer func() {
		if s.metrics != nil {
			s.metrics.ObserveHeartbeat(outcome)
		}
	}()
	if err := s.repository.AdvanceSecurityEventStreamState(ctx, s.config.Issuer, s.config.Audience, s.clock().UTC(), s.config.StaleAfter); err != nil {
		_ = s.repository.RecordSecurityEventHeartbeatFailure(ctx, s.config.Issuer, s.config.Audience, "state_advance_failed")
	}
	state, err := randomVerificationState()
	if err != nil {
		_ = s.repository.RecordSecurityEventHeartbeatFailure(ctx, s.config.Issuer, s.config.Audience, "state_generation_failed")
		return
	}
	hash := sha256.Sum256([]byte(state))
	now := s.clock().UTC()
	if err := s.repository.BeginSecurityEventVerification(ctx, s.config.Issuer, s.config.Audience, hash[:], now); err != nil {
		_ = s.repository.RecordSecurityEventHeartbeatFailure(ctx, s.config.Issuer, s.config.Audience, "state_persistence_failed")
		return
	}
	token, err := s.managementToken(ctx)
	if err != nil {
		_ = s.repository.RecordSecurityEventHeartbeatFailure(ctx, s.config.Issuer, s.config.Audience, "management_token_failed")
		return
	}
	payload, _ := json.Marshal(map[string]string{"stream_id": s.config.StreamID, "state": state})
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(s.config.Issuer, "/")+"/v1/verify", bytes.NewReader(payload))
	if err != nil {
		return
	}
	request.Header.Set("Authorization", "Bearer "+token)
	request.Header.Set("Content-Type", "application/json")
	response, err := s.client.Do(request)
	if err != nil {
		_ = s.repository.RecordSecurityEventHeartbeatFailure(ctx, s.config.Issuer, s.config.Audience, "verification_request_failed")
		return
	}
	defer response.Body.Close()
	_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024))
	if response.StatusCode != http.StatusNoContent {
		_ = s.repository.RecordSecurityEventHeartbeatFailure(ctx, s.config.Issuer, s.config.Audience, fmt.Sprintf("verification_http_%d", response.StatusCode))
		return
	}
	outcome = "accepted"
}

func (s *Service) managementToken(ctx context.Context) (string, error) {
	s.tokenMu.Lock()
	defer s.tokenMu.Unlock()
	if s.token != "" && s.clock().Before(s.tokenUntil) {
		return s.token, nil
	}
	form := url.Values{"grant_type": {"client_credentials"}, "scope": {"ssf.stream.read ssf.stream.verify"}}
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, s.config.ManagementTokenURL, strings.NewReader(form.Encode()))
	if err != nil {
		return "", err
	}
	request.SetBasicAuth(s.config.ManagementClientID, s.config.ManagementSecret)
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	response, err := s.client.Do(request)
	if err != nil {
		return "", err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024))
		return "", fmt.Errorf("management token endpoint returned HTTP %d", response.StatusCode)
	}
	payload, err := io.ReadAll(io.LimitReader(response.Body, 64*1024+1))
	if err != nil || len(payload) > 64*1024 {
		return "", errors.New("management token response is too large")
	}
	var result struct {
		AccessToken string `json:"access_token"`
		TokenType   string `json:"token_type"`
		ExpiresIn   int64  `json:"expires_in"`
	}
	if json.Unmarshal(payload, &result) != nil || result.AccessToken == "" || !strings.EqualFold(result.TokenType, "Bearer") {
		return "", errors.New("management token response is invalid")
	}
	if result.ExpiresIn <= 0 {
		result.ExpiresIn = 60
	}
	cacheFor := time.Duration(result.ExpiresIn) * time.Second
	if cacheFor > 30*time.Second {
		cacheFor -= 30 * time.Second
	}
	s.token, s.tokenUntil = result.AccessToken, s.clock().Add(cacheFor)
	return s.token, nil
}

func randomVerificationState() (string, error) {
	payload := make([]byte, 32)
	if _, err := rand.Read(payload); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(payload), nil
}