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

package securityevents

import (
	"context"
	"crypto/sha256"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

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

type fakeStateRepository struct {
	mutex        sync.Mutex
	stateHash    []byte
	beginCount   int
	beginSignal  chan struct{}
	verification chan struct{}
	evaluation   store.SecurityEventEvaluation
}

func (*fakeStateRepository) EnsureSecurityEventStreamState(context.Context, string, string, string) error {
	return nil
}
func (f *fakeStateRepository) BeginSecurityEventVerification(_ context.Context, _, _ string, hash []byte, _ time.Time) error {
	f.mutex.Lock()
	f.stateHash = append([]byte(nil), hash...)
	f.beginCount++
	if f.beginSignal != nil {
		select {
		case f.beginSignal <- struct{}{}:
		default:
		}
	}
	f.mutex.Unlock()
	return nil
}
func (*fakeStateRepository) RecordSecurityEventHeartbeatFailure(context.Context, string, string, string) error {
	return nil
}
func (*fakeStateRepository) AdvanceSecurityEventStreamState(context.Context, string, string, time.Time, time.Duration) error {
	return nil
}
func (f *fakeStateRepository) EvaluateOIDCSecurityEvent(context.Context, string, string, string, string, string, string, time.Time, time.Time, time.Duration) (store.SecurityEventEvaluation, error) {
	return f.evaluation, nil
}

type deniedHeartbeatLeadershipRepository struct {
	*fakeStateRepository
}

func (*deniedHeartbeatLeadershipRepository) TryAcquireSecurityEventHeartbeatLeadership(context.Context) (store.Leadership, bool, error) {
	return nil, false, nil
}

type fakeHeartbeatLeadership struct {
	released chan struct{}
	once     sync.Once
}

func (*fakeHeartbeatLeadership) KeepAlive(context.Context) error { return nil }
func (leadership *fakeHeartbeatLeadership) Release() {
	leadership.once.Do(func() { close(leadership.released) })
}

type grantedHeartbeatLeadershipRepository struct {
	*fakeStateRepository
	leadership *fakeHeartbeatLeadership
}

func (repository *grantedHeartbeatLeadershipRepository) TryAcquireSecurityEventHeartbeatLeadership(context.Context) (store.Leadership, bool, error) {
	return repository.leadership, true, nil
}

func TestServiceFollowerDoesNotRunVerificationHeartbeat(t *testing.T) {
	repository := &deniedHeartbeatLeadershipRepository{fakeStateRepository: &fakeStateRepository{}}
	service, err := NewService(repository, ServiceConfig{
		Issuer: "https://auth.test.example/ssf", Audience: "urn:easyai:ssf:receiver:app",
		StreamID:           "00000000-0000-4000-8000-000000000001",
		ManagementTokenURL: "https://auth.test.example/token",
		ManagementClientID: "gateway-ssf", ManagementSecret: "management-secret",
		HeartbeatInterval: time.Hour, StaleAfter: 2 * time.Hour,
	})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() {
		service.Run(ctx)
		close(done)
	}()
	time.Sleep(100 * time.Millisecond)
	cancel()
	<-done
	repository.mutex.Lock()
	defer repository.mutex.Unlock()
	if repository.beginCount != 0 {
		t.Fatalf("follower persisted %d verification heartbeats", repository.beginCount)
	}
}

func TestServiceLeaderRunsHeartbeatAndReleasesLeadership(t *testing.T) {
	repository := &grantedHeartbeatLeadershipRepository{
		fakeStateRepository: &fakeStateRepository{beginSignal: make(chan struct{}, 1)},
		leadership:          &fakeHeartbeatLeadership{released: make(chan struct{})},
	}
	tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/token" {
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"access_token":"management-token","token_type":"Bearer","expires_in":300}`))
			return
		}
		w.WriteHeader(http.StatusNoContent)
	}))
	defer tokenServer.Close()
	service, err := NewService(repository, ServiceConfig{
		Issuer: tokenServer.URL, Audience: "urn:easyai:ssf:receiver:app",
		StreamID:           "00000000-0000-4000-8000-000000000001",
		ManagementTokenURL: tokenServer.URL + "/token",
		ManagementClientID: "gateway-ssf", ManagementSecret: "management-secret",
		HeartbeatInterval: time.Hour, StaleAfter: 2 * time.Hour, HTTPClient: tokenServer.Client(),
	})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() {
		service.Run(ctx)
		close(done)
	}()
	select {
	case <-repository.beginSignal:
	case <-time.After(2 * time.Second):
		t.Fatal("leader did not persist a verification heartbeat")
	}
	cancel()
	<-done
	select {
	case <-repository.leadership.released:
	default:
		t.Fatal("leader did not release PostgreSQL leadership")
	}
}

func TestServiceRequestsVerificationWithMachineTokenAndPersistsStateFirst(t *testing.T) {
	repository := &fakeStateRepository{verification: make(chan struct{}, 1)}
	var server *httptest.Server
	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/token":
			clientID, secret, ok := r.BasicAuth()
			if !ok || clientID != "gateway-ssf" || secret != "management-secret" {
				t.Fatal("machine client authentication is missing")
			}
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"access_token":"management-token","token_type":"Bearer","expires_in":300}`))
		case "/ssf/v1/verify":
			if r.Header.Get("Authorization") != "Bearer management-token" {
				t.Fatal("verification bearer is missing")
			}
			repository.mutex.Lock()
			persisted := len(repository.stateHash) == sha256.Size
			repository.mutex.Unlock()
			if !persisted {
				t.Fatal("verification state was not persisted before request")
			}
			w.WriteHeader(http.StatusNoContent)
			repository.verification <- struct{}{}
		}
	}))
	defer server.Close()
	service, err := NewService(repository, ServiceConfig{
		Issuer: server.URL + "/ssf", Audience: "urn:easyai:ssf:receiver:app", StreamID: "00000000-0000-4000-8000-000000000001",
		ManagementTokenURL: server.URL + "/token", ManagementClientID: "gateway-ssf", ManagementSecret: "management-secret",
		HeartbeatInterval: time.Hour, StaleAfter: 2 * time.Hour, HTTPClient: server.Client(),
	})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	if err := service.Start(ctx); err != nil {
		t.Fatal(err)
	}
	select {
	case <-repository.verification:
	case <-time.After(2 * time.Second):
		t.Fatal("verification was not requested")
	}
}