browser/platform/service/control.go

package service

import (
	"crypto/hmac"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"browser.local/platform/domain"
	"browser.local/platform/repo"
	"browser.local/platform/validator"
)

const (
	defaultHeartbeatIntervalSeconds = 15
	defaultRunSessionTTL            = 24 * time.Hour
	maxRunRequestClockSkew          = 5 * time.Minute
	maxRunRequestNonces             = 8192
)

func (svc *CoreService) RegisterRunHello(hello domain.RunControlHello) (domain.RunControlHelloResult, error) {
	hello = domain.CopyRunControlHello(hello)
	if err := validator.ValidateRunControlHello(hello); err != nil {
		return domain.RunControlHelloResult{}, err
	}
	if hasComponentAuthIdentity(hello) {
		if err := svc.validateDedicatedRunHello(hello); err != nil {
			return domain.RunControlHelloResult{}, err
		}
		auth, err := svc.AuthenticateComponent(domain.ComponentAuthenticationRequest{
			ServerInstanceID: hello.ServerInstanceID,
			ComponentKind:    hello.ComponentKind,
			ComponentKey:     hello.ComponentKey,
			Generation:       hello.KeyGeneration,
			Key:              hello.RegistrationToken,
		})
		if err != nil {
			return domain.RunControlHelloResult{}, err
		}
		if !auth.Allowed {
			return domain.CopyRunControlHelloResult(domain.RunControlHelloResult{
				Accepted:                 false,
				RunEndpointID:            hello.RunEndpointID,
				ServerTime:               svc.now(),
				HeartbeatIntervalSeconds: defaultHeartbeatIntervalSeconds,
				FeatureFlags:             []string{"runtime-key.auth.denied"},
			}), nil
		}
	}

	stamp := svc.now()
	endpoint := domain.RunEndpoint{
		ID:              hello.RunEndpointID,
		DisplayName:     hello.DisplayName,
		Version:         hello.Version,
		Platform:        hello.Platform,
		Architecture:    hello.Architecture,
		Status:          domain.RunEndpointStatusOnline,
		Capabilities:    domain.CopyStringSlice(hello.CapabilityReport.Capabilities),
		Capacity:        hello.Capacity,
		LastHeartbeatAt: stamp,
	}
	if err := validator.ValidateRunEndpoint(endpoint); err != nil {
		return domain.RunControlHelloResult{}, err
	}

	svc.controlMu.Lock()
	defer svc.controlMu.Unlock()

	if err := svc.upsertRunEndpoint(endpoint); err != nil {
		return domain.RunControlHelloResult{}, err
	}
	previous, previousErr := svc.store.RunControlSessions().Get(hello.RunEndpointID)
	generation := 1
	if previousErr == nil {
		generation = previous.Generation + 1
	} else if !errors.Is(previousErr, repo.ErrNotFound) {
		return domain.RunControlHelloResult{}, previousErr
	}
	sessionToken, err := svc.nextSessionToken()
	if err != nil {
		return domain.RunControlHelloResult{}, err
	}
	session := domain.RunControlSession{
		RunEndpointID:            hello.RunEndpointID,
		SessionToken:             sessionToken,
		SessionTokenHash:         tokenHash(sessionToken),
		Status:                   domain.AuthSessionStatusActive,
		Generation:               generation,
		CapabilityFingerprint:    hello.CapabilityReport.Fingerprint,
		HeartbeatIntervalSeconds: defaultHeartbeatIntervalSeconds,
		CreatedAt:                stamp,
		UpdatedAt:                stamp,
		ExpiresAt:                stamp.Add(defaultRunSessionTTL),
		RequireSignedRequests:    hasComponentAuthIdentity(hello),
	}
	if err := validator.ValidateRunControlSession(session); err != nil {
		return domain.RunControlHelloResult{}, err
	}
	if previousErr == nil {
		if err := svc.store.RunControlSessions().Update(session); err != nil {
			return domain.RunControlHelloResult{}, err
		}
	} else if err := svc.store.RunControlSessions().Create(session); err != nil {
		return domain.RunControlHelloResult{}, err
	}
	svc.runSessions[hello.RunEndpointID] = session
	featureFlags := []string{"control.hello", "control.heartbeat", "signed-envelope.v1.optional"}
	if session.RequireSignedRequests {
		featureFlags[2] = "signed-envelope.v1.required"
	}
	return domain.CopyRunControlHelloResult(domain.RunControlHelloResult{
		Accepted:                 true,
		RunEndpointID:            hello.RunEndpointID,
		SessionToken:             sessionToken,
		ServerTime:               stamp,
		HeartbeatIntervalSeconds: defaultHeartbeatIntervalSeconds,
		SessionExpiresAt:         session.ExpiresAt,
		FeatureFlags:             featureFlags,
	}), nil
}

func (svc *CoreService) validateDedicatedRunHello(hello domain.RunControlHello) error {
	if hello.ComponentKind != domain.DistributionComponentRun {
		return validationError("component-authenticated run hello must use the run component")
	}
	instance, err := svc.store.ServerInstances().Get(hello.ServerInstanceID)
	if err != nil {
		return err
	}
	if strings.TrimSpace(instance.DeploymentTargetID) == "" {
		return nil // legacy Run registrations keep their historical endpoint contract.
	}
	if hello.PluginID != instance.PluginID || hello.RunEndpointID != instance.RunEndpointID {
		return validationError("run endpoint identity does not match the server binding")
	}
	instances, err := svc.store.ServerInstances().List(domain.ServerInstanceFilter{RunEndpointID: hello.RunEndpointID})
	if err != nil {
		return err
	}
	for _, candidate := range instances {
		if candidate.ID != instance.ID && candidate.State != domain.ServerInstanceStateDeleted {
			return validationError("run endpoint is already bound to another server")
		}
	}
	return nil
}

func hasComponentAuthIdentity(hello domain.RunControlHello) bool {
	return hello.ServerInstanceID != "" || hello.PluginID != "" || hello.ComponentKind != "" || hello.ComponentKey != "" || hello.KeyGeneration != 0
}

func (svc *CoreService) AcceptRunHeartbeat(heartbeat domain.RunControlHeartbeat) (domain.RunControlHeartbeatResult, error) {
	heartbeat = domain.CopyRunControlHeartbeat(heartbeat)
	if err := validator.ValidateRunControlHeartbeat(heartbeat); err != nil {
		return domain.RunControlHeartbeatResult{}, err
	}

	stamp := svc.now()

	svc.controlMu.Lock()
	defer svc.controlMu.Unlock()

	session, err := svc.currentRunSession(heartbeat.RunEndpointID, heartbeat.SessionToken)
	if err != nil {
		return domain.RunControlHeartbeatResult{}, err
	}

	endpoint, err := svc.store.RunEndpoints().Get(heartbeat.RunEndpointID)
	if err != nil {
		return domain.RunControlHeartbeatResult{}, err
	}
	endpoint.Version = heartbeat.Version
	endpoint.Status = heartbeat.Status
	endpoint.Capacity = heartbeat.Capacity
	endpoint.LastHeartbeatAt = stamp
	if err := validator.ValidateRunEndpoint(endpoint); err != nil {
		return domain.RunControlHeartbeatResult{}, err
	}
	if err := svc.store.RunEndpoints().Update(endpoint); err != nil {
		return domain.RunControlHeartbeatResult{}, err
	}

	refreshCapabilities := session.CapabilityFingerprint != heartbeat.CapabilityFingerprint
	session.CapabilityFingerprint = heartbeat.CapabilityFingerprint
	session.UpdatedAt = stamp
	if err := svc.store.RunControlSessions().Update(session); err != nil {
		return domain.RunControlHeartbeatResult{}, err
	}
	svc.runSessions[heartbeat.RunEndpointID] = session

	return domain.CopyRunControlHeartbeatResult(domain.RunControlHeartbeatResult{
		Accepted:             true,
		RunEndpointID:        heartbeat.RunEndpointID,
		NextHeartbeatSeconds: session.HeartbeatIntervalSeconds,
		RefreshCapabilities:  refreshCapabilities,
		ServerTime:           stamp,
	}), nil
}

func (svc *CoreService) upsertRunEndpoint(endpoint domain.RunEndpoint) error {
	if _, err := svc.store.RunEndpoints().Get(endpoint.ID); err != nil {
		if errors.Is(err, repo.ErrNotFound) {
			return svc.store.RunEndpoints().Create(endpoint)
		}
		return err
	}
	return svc.store.RunEndpoints().Update(endpoint)
}

func (svc *CoreService) nextSessionToken() (string, error) {
	return randomToken()
}

func (svc *CoreService) currentRunSession(runEndpointID string, sessionToken string) (domain.RunControlSession, error) {
	session, exists := svc.runSessions[runEndpointID]
	if !exists {
		stored, err := svc.store.RunControlSessions().Get(runEndpointID)
		if err != nil {
			return domain.RunControlSession{}, runAuthenticationError(true)
		}
		session = stored
	}
	presentedHash := tokenHash(strings.TrimSpace(sessionToken))
	if strings.TrimSpace(sessionToken) == "" || session.Status != domain.AuthSessionStatusActive || !session.RevokedAt.IsZero() || !svc.now().Before(session.ExpiresAt) || subtle.ConstantTimeCompare([]byte(session.SessionTokenHash), []byte(presentedHash)) != 1 {
		if session.Status == domain.AuthSessionStatusActive && !svc.now().Before(session.ExpiresAt) {
			session.Status = domain.AuthSessionStatusRevoked
			session.RevokedAt = svc.now()
			session.UpdatedAt = session.RevokedAt
			_ = svc.store.RunControlSessions().Update(session)
		}
		return domain.RunControlSession{}, runAuthenticationError(session.RequireSignedRequests)
	}
	return domain.CopyRunControlSession(session), nil
}

func (svc *CoreService) revokeRunControlSessionForInstance(instance domain.ServerInstance) error {
	if strings.TrimSpace(instance.RunEndpointID) == "" {
		return nil
	}
	instances, err := svc.store.ServerInstances().List(domain.ServerInstanceFilter{RunEndpointID: instance.RunEndpointID})
	if err != nil {
		return err
	}
	for _, candidate := range instances {
		if candidate.ID != instance.ID && candidate.State != domain.ServerInstanceStateDeleted {
			return nil
		}
	}
	svc.controlMu.Lock()
	defer svc.controlMu.Unlock()

	session, err := svc.store.RunControlSessions().Get(instance.RunEndpointID)
	if err != nil {
		if errors.Is(err, repo.ErrNotFound) {
			delete(svc.runSessions, instance.RunEndpointID)
			return nil
		}
		return err
	}
	if session.Status == domain.AuthSessionStatusActive {
		stamp := svc.now()
		session.Status = domain.AuthSessionStatusRevoked
		session.RevokedAt = stamp
		session.UpdatedAt = stamp
		if err := validator.ValidateRunControlSession(session); err != nil {
			return err
		}
		if err := svc.store.RunControlSessions().Update(session); err != nil {
			return err
		}
	}
	delete(svc.runSessions, instance.RunEndpointID)

	endpoint, err := svc.store.RunEndpoints().Get(instance.RunEndpointID)
	if err != nil {
		if errors.Is(err, repo.ErrNotFound) {
			return nil
		}
		return err
	}
	endpoint.Status = domain.RunEndpointStatusOffline
	if err := validator.ValidateRunEndpoint(endpoint); err != nil {
		return err
	}
	return svc.store.RunEndpoints().Update(endpoint)
}

func runAuthenticationError(requireSigned bool) error {
	if !requireSigned {
		return validationError("sessionToken is invalid")
	}
	return fmt.Errorf("sessionToken is invalid: %w", ErrUnauthorized)
}

func (svc *CoreService) AuthorizeRunRequestSignature(request domain.RunRequestSignature) error {
	svc.controlMu.Lock()
	defer svc.controlMu.Unlock()

	session, err := svc.currentRunSession(request.RunEndpointID, request.SessionToken)
	if err != nil {
		return err
	}
	if strings.TrimSpace(request.Signature) == "" && !session.RequireSignedRequests {
		return nil
	}
	if strings.TrimSpace(request.Timestamp) == "" || strings.TrimSpace(request.Nonce) == "" || strings.TrimSpace(request.Signature) == "" {
		return runAuthenticationError(true)
	}
	unixSeconds, err := strconv.ParseInt(request.Timestamp, 10, 64)
	if err != nil {
		return runAuthenticationError(true)
	}
	stamp := time.Unix(unixSeconds, 0).UTC()
	delta := svc.now().Sub(stamp)
	if delta < -maxRunRequestClockSkew || delta > maxRunRequestClockSkew {
		return runAuthenticationError(true)
	}
	session.UsedNonces = activeRunNonces(session.UsedNonces, svc.now().Add(-maxRunRequestClockSkew))
	if len(request.Nonce) > 128 || runNonceSeen(session.UsedNonces, request.Nonce) || len(session.UsedNonces) >= maxRunRequestNonces {
		return runAuthenticationError(true)
	}
	canonical := strings.Join([]string{request.Method, request.Path, request.Timestamp, request.Nonce, request.BodyHash}, "\n")
	mac := hmac.New(sha256.New, []byte(request.SessionToken))
	_, _ = mac.Write([]byte(canonical))
	expected := hex.EncodeToString(mac.Sum(nil))
	provided, err := hex.DecodeString(request.Signature)
	if err != nil || subtle.ConstantTimeCompare([]byte(expected), []byte(hex.EncodeToString(provided))) != 1 {
		return runAuthenticationError(true)
	}
	session.UsedNonces = append(session.UsedNonces, request.Timestamp+":"+request.Nonce)
	session.UpdatedAt = svc.now()
	if err := svc.store.RunControlSessions().Update(session); err != nil {
		return err
	}
	svc.runSessions[request.RunEndpointID] = session
	return nil
}

func activeRunNonces(entries []string, cutoff time.Time) []string {
	active := make([]string, 0, len(entries))
	for _, entry := range entries {
		timestamp, _, ok := strings.Cut(entry, ":")
		if !ok {
			active = append(active, entry)
			continue
		}
		unixSeconds, err := strconv.ParseInt(timestamp, 10, 64)
		if err == nil && !time.Unix(unixSeconds, 0).Before(cutoff) {
			active = append(active, entry)
		}
	}
	return active
}

func runNonceSeen(entries []string, nonce string) bool {
	for _, entry := range entries {
		_, storedNonce, ok := strings.Cut(entry, ":")
		if (ok && storedNonce == nonce) || (!ok && entry == nonce) {
			return true
		}
	}
	return false
}