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

package securityevents

import (
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/json"
	"errors"
	"io"
	"mime"
	"net/http"
	"strings"
	"time"

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

type Repository interface {
	ApplySessionRevoked(context.Context, store.ApplySessionRevokedInput) (store.ApplySecurityEventResult, error)
	ApplySecurityEventStreamUpdated(context.Context, string, string, string, string, string, string, time.Time) (bool, error)
	ConfirmSecurityEventVerification(context.Context, string, string, string, string, []byte, time.Time) (bool, error)
	RecordSecurityEventReceipt(context.Context, string, string, string, string, string) (bool, error)
}

type Handler struct {
	verifier     *Verifier
	repository   Repository
	issuer       string
	audience     string
	streamID     string
	bearerSecret string
	nextSecret   string
	clock        func() time.Time
	metrics      *Metrics
}

func (h *Handler) SetMetrics(metrics *Metrics) { h.metrics = metrics }

func NewHandler(verifier *Verifier, repository Repository, issuer, audience, streamID, bearerSecret, nextSecret string) (*Handler, error) {
	if verifier == nil || repository == nil || issuer == "" || audience == "" || streamID == "" || len(bearerSecret) < 16 {
		return nil, errors.New("SSF receiver dependencies and bearer secret are required")
	}
	return &Handler{verifier: verifier, repository: repository, issuer: issuer, audience: audience, streamID: streamID, bearerSecret: bearerSecret, nextSecret: nextSecret, clock: time.Now}, nil
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	started := h.clock()
	outcome := "rejected"
	var sessionsDeleted int64
	defer func() {
		if h.metrics != nil {
			h.metrics.ObserveReceipt(outcome, h.clock().Sub(started), sessionsDeleted)
		}
	}()
	w.Header().Set("Cache-Control", "no-store")
	if r.Method != http.MethodPost {
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}
	if !h.validAuthorization(r.Header.Get("Authorization")) {
		w.Header().Set("WWW-Authenticate", `Bearer realm="ssf-receiver"`)
		writeSETError(w, http.StatusUnauthorized, "authentication_failed")
		return
	}
	mediaType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err != nil || mediaType != "application/secevent+jwt" || r.Header.Get("Content-Encoding") != "" {
		writeSETError(w, http.StatusBadRequest, "invalid_request")
		return
	}
	r.Body = http.MaxBytesReader(w, r.Body, 64*1024)
	payload, err := io.ReadAll(r.Body)
	if err != nil || len(payload) == 0 {
		writeSETError(w, http.StatusBadRequest, "invalid_request")
		return
	}
	event, err := h.verifier.Verify(r.Context(), string(payload))
	if err != nil {
		var protocol *protocolError
		if errors.As(err, &protocol) {
			writeSETError(w, http.StatusBadRequest, protocol.Code)
		} else {
			writeSETError(w, http.StatusBadRequest, "invalid_request")
		}
		return
	}
	switch event.EventType {
	case SessionRevokedEventType:
		var result store.ApplySecurityEventResult
		result, err = h.repository.ApplySessionRevoked(r.Context(), store.ApplySessionRevokedInput{
			Issuer: event.Issuer, Audience: event.Audience, JTI: event.JTI, TransactionID: event.TransactionID,
			SubjectIssuer: event.SubjectIssuer, TenantID: event.TenantID, Subject: event.Subject, EventTimestamp: event.EventTimestamp,
			InitiatingEntity: event.InitiatingEntity,
		})
		if err == nil {
			sessionsDeleted = result.SessionsDeleted
			if result.Duplicate {
				outcome = "duplicate"
			} else {
				outcome = "accepted"
			}
		}
	case VerificationEventType:
		var confirmed bool
		if event.State == "" {
			confirmed, err = h.repository.RecordSecurityEventReceipt(r.Context(), h.issuer, h.audience, event.JTI, VerificationEventType, h.streamID)
		} else {
			stateHash := sha256.Sum256([]byte(event.State))
			confirmed, err = h.repository.ConfirmSecurityEventVerification(r.Context(), h.issuer, h.audience, h.streamID, event.JTI, stateHash[:], h.clock().UTC())
		}
		if err != nil && strings.Contains(err.Error(), "invalid verification state") {
			writeSETError(w, http.StatusBadRequest, "invalid_request")
			return
		}
		if err == nil {
			if confirmed {
				outcome = "accepted"
				if h.metrics != nil {
					h.metrics.ObserveVerificationAccepted()
				}
			} else {
				outcome = "duplicate"
			}
		}
	case StreamUpdatedEventType:
		var inserted bool
		inserted, err = h.repository.ApplySecurityEventStreamUpdated(
			r.Context(), h.issuer, h.audience, event.JTI, h.streamID, event.StreamStatus, event.Reason, h.clock().UTC(),
		)
		if err == nil {
			if inserted {
				outcome = "accepted"
			} else {
				outcome = "duplicate"
			}
		}
	}
	if err != nil {
		w.Header().Set("Retry-After", "1")
		w.WriteHeader(http.StatusServiceUnavailable)
		return
	}
	w.WriteHeader(http.StatusAccepted)
}

func (h *Handler) validAuthorization(value string) bool {
	if constantBearerCompare(value, h.bearerSecret) {
		return true
	}
	return h.nextSecret != "" && constantBearerCompare(value, h.nextSecret)
}

func constantBearerCompare(header, secret string) bool {
	providedHash := sha256.Sum256([]byte(header))
	expectedHash := sha256.Sum256([]byte("Bearer " + secret))
	return subtle.ConstantTimeCompare(providedHash[:], expectedHash[:]) == 1
}

func writeSETError(w http.ResponseWriter, status int, code string) {
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Content-Language", "en")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(map[string]string{
		"err":         code,
		"description": setErrorDescription(code),
	})
}

func setErrorDescription(code string) string {
	switch code {
	case "invalid_request":
		return "The SET request is malformed or contains invalid claims."
	case "invalid_key":
		return "The SET signing key or signature is invalid."
	case "invalid_issuer":
		return "The SET issuer is not accepted by this receiver."
	case "invalid_audience":
		return "The SET audience does not identify this receiver."
	case "authentication_failed":
		return "Receiver authentication failed."
	case "access_denied":
		return "The SET is not permitted for this receiver."
	default:
		return "The SET request could not be processed."
	}
}