proxy-pool/internal/domain/workerruntime/validation.go

package workerruntime

import (
	"crypto/sha256"
	"encoding/json"
	"regexp"
	"sort"
	"strings"
)

const (
	maximumLabels          = 32
	maximumLabelKeyBytes   = 64
	maximumLabelValueBytes = 256
	maximumLabelTotalBytes = 4 << 10
)

var identifierPattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$`)

// ValidIdentifier accepts stable Worker, Session and Proxy identifiers.
func ValidIdentifier(value string) bool {
	return identifierPattern.MatchString(value)
}

// NormalizeLabels validates and deep-copies the bounded Worker label set.
func NormalizeLabels(labels map[string]string) (map[string]string, error) {
	if len(labels) > maximumLabels {
		return nil, ErrInvalidSession
	}
	normalized := make(map[string]string, len(labels))
	total := 0
	for key, value := range labels {
		if !ValidIdentifier(key) || len(key) > maximumLabelKeyBytes || value == "" ||
			strings.TrimSpace(value) != value || strings.IndexByte(value, 0) >= 0 ||
			len(value) > maximumLabelValueBytes {
			return nil, ErrInvalidSession
		}
		total += len(key) + len(value)
		if total > maximumLabelTotalBytes {
			return nil, ErrInvalidSession
		}
		normalized[key] = value
	}
	return normalized, nil
}

// NormalizeSession prepares a new, not-yet-acknowledged session for storage.
func NormalizeSession(session Session) (Session, error) {
	if !ValidIdentifier(session.WorkerID) || !ValidIdentifier(session.InstanceID) ||
		!ValidIdentifier(session.SessionID) || !ValidIdentifier(session.Zone) ||
		session.ProtocolVersion == 0 || session.AckedSnapshotVersion != 0 ||
		session.AckedOwnershipEpoch != 0 || !checksumIsZero(session.AckedChecksum) ||
		session.RuntimeEnabled {
		return Session{}, ErrInvalidSession
	}
	labels, err := NormalizeLabels(session.Labels)
	if err != nil {
		return Session{}, err
	}
	session.Labels = labels
	return session, nil
}

func NormalizeSnapshotReference(reference SnapshotReference) (SnapshotReference, error) {
	if !ValidIdentifier(reference.WorkerID) || reference.Version == 0 || reference.OwnershipEpoch == 0 ||
		checksumIsZero(reference.Checksum) {
		return SnapshotReference{}, ErrInvalidSnapshotReference
	}
	return reference, nil
}

func NormalizeAcknowledgement(acknowledgement SnapshotAcknowledgement) (SnapshotAcknowledgement, error) {
	if !ValidIdentifier(acknowledgement.WorkerID) || !ValidIdentifier(acknowledgement.SessionID) ||
		(acknowledgement.ErrorCode != "" && !ValidIdentifier(acknowledgement.ErrorCode)) {
		return SnapshotAcknowledgement{}, ErrInvalidAcknowledgement
	}
	reference, err := NormalizeSnapshotReference(acknowledgement.Reference)
	if err != nil || reference.WorkerID != acknowledgement.WorkerID {
		return SnapshotAcknowledgement{}, ErrInvalidAcknowledgement
	}
	acknowledgement.Reference = reference
	return acknowledgement, nil
}

func NormalizeOutcomeReference(reference OutcomeReference) (OutcomeReference, error) {
	if !ValidIdentifier(reference.WorkerID) || !ValidIdentifier(reference.SessionID) || reference.Sequence == 0 ||
		checksumIsZero(reference.Digest) {
		return OutcomeReference{}, ErrInvalidOutcome
	}
	return reference, nil
}

// NormalizeReport returns the canonical sparse replacement and its digest.
func NormalizeReport(report Report) (Report, [sha256.Size]byte, error) {
	if !ValidIdentifier(report.WorkerID) || !ValidIdentifier(report.SessionID) || report.Sequence == 0 ||
		report.SnapshotVersion == 0 || report.OwnershipEpoch == 0 || report.ObservedAt.IsZero() {
		return Report{}, [sha256.Size]byte{}, ErrInvalidReport
	}
	normalized := report
	normalized.ObservedAt = report.ObservedAt.UTC()
	normalized.Counters = append([]Counter(nil), report.Counters...)
	sort.Slice(normalized.Counters, func(left, right int) bool {
		return normalized.Counters[left].ProxyID < normalized.Counters[right].ProxyID
	})
	seen := make(map[string]struct{}, len(normalized.Counters))
	for _, counter := range normalized.Counters {
		if !ValidIdentifier(counter.ProxyID) || counter.Active < 0 || counter.Reserved < 0 {
			return Report{}, [sha256.Size]byte{}, ErrInvalidReport
		}
		if _, exists := seen[counter.ProxyID]; exists {
			return Report{}, [sha256.Size]byte{}, ErrInvalidReport
		}
		seen[counter.ProxyID] = struct{}{}
	}
	payload, err := json.Marshal(normalized)
	if err != nil {
		return Report{}, [sha256.Size]byte{}, ErrInvalidReport
	}
	return normalized, sha256.Sum256(payload), nil
}

func checksumIsZero(checksum [sha256.Size]byte) bool {
	return checksum == [sha256.Size]byte{}
}

func compareSnapshotTuple(left, right SnapshotReference) int {
	switch {
	case left.OwnershipEpoch < right.OwnershipEpoch:
		return -1
	case left.OwnershipEpoch > right.OwnershipEpoch:
		return 1
	case left.Version < right.Version:
		return -1
	case left.Version > right.Version:
		return 1
	default:
		return 0
	}
}

func sessionReference(session Session) SnapshotReference {
	return SnapshotReference{
		WorkerID: session.WorkerID, Version: session.AckedSnapshotVersion,
		OwnershipEpoch: session.AckedOwnershipEpoch, Checksum: session.AckedChecksum,
	}
}