488 lines
14 KiB
Go
488 lines
14 KiB
Go
package health
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"errors"
|
|
"fmt"
|
|
"reflect"
|
|
"sort"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
healthDomain "proxy-pool/internal/domain/health"
|
|
proxyDomain "proxy-pool/internal/domain/proxy"
|
|
)
|
|
|
|
var (
|
|
ErrInvalidTaskBroker = errors.New("invalid health task broker")
|
|
ErrInvalidTaskClaim = errors.New("invalid health task claim")
|
|
ErrInvalidLeasedTask = errors.New("invalid leased health task")
|
|
ErrTaskNotFound = errors.New("health task not found")
|
|
ErrTaskLeaseExpired = errors.New("health task lease expired")
|
|
ErrTaskLeaseNotOwned = errors.New("health task lease is not owned by checker")
|
|
ErrTaskObservation = errors.New("health observation does not match task")
|
|
)
|
|
|
|
const defaultMaxTasksPerClaim = 128
|
|
|
|
// TaskMaterial is the short-lived proxy connection material needed to execute
|
|
// one probe. It crosses only the authenticated Checker control stream and is
|
|
// deliberately redacted from formatted values and persistence APIs.
|
|
type TaskMaterial struct {
|
|
Protocol proxyDomain.Scheme
|
|
Host string
|
|
Port uint16
|
|
SecretRef string
|
|
CredentialVersion string
|
|
Username string
|
|
Password string
|
|
}
|
|
|
|
func (TaskMaterial) Format(state fmt.State, _ rune) {
|
|
_, _ = state.Write([]byte("health.TaskMaterial{Credentials:<redacted>}"))
|
|
}
|
|
|
|
// TaskMaterialResolver resolves a task's endpoint and credential material at
|
|
// claim time. A production implementation reads the Controller-local
|
|
// credential store; Checkers never read Redis or PostgreSQL.
|
|
type TaskMaterialResolver interface {
|
|
ResolveCheckTask(context.Context, Candidate) (TaskMaterial, error)
|
|
}
|
|
|
|
type TaskMaterialResolverFunc func(context.Context, Candidate) (TaskMaterial, error)
|
|
|
|
func (resolver TaskMaterialResolverFunc) ResolveCheckTask(ctx context.Context, candidate Candidate) (TaskMaterial, error) {
|
|
return resolver(ctx, candidate)
|
|
}
|
|
|
|
// TaskClaim is one bounded pull request from a Checker process. MaxInFlight
|
|
// applies across all streams for the checker ID, so reconnecting cannot grow
|
|
// its local work window.
|
|
type TaskClaim struct {
|
|
CheckerID string
|
|
InstanceID string
|
|
MaxInFlight int
|
|
SupportedLevels []healthDomain.Level
|
|
}
|
|
|
|
// LeasedTask is a Controller-assigned task. It holds material only in memory
|
|
// for the duration of a task lease and must not be logged.
|
|
type LeasedTask struct {
|
|
TaskID string
|
|
ProxyID string
|
|
Protocol proxyDomain.Scheme
|
|
Host string
|
|
Port uint16
|
|
SecretRef string
|
|
CredentialVersion string
|
|
Username string
|
|
Password string
|
|
Level healthDomain.Level
|
|
RoutingName string
|
|
TargetURL string
|
|
Deadline time.Time
|
|
Attempts int
|
|
}
|
|
|
|
func (LeasedTask) Format(state fmt.State, _ rune) {
|
|
_, _ = state.Write([]byte("health.LeasedTask{Credentials:<redacted>}"))
|
|
}
|
|
|
|
// TaskBroker is the shared task lease boundary. The Scheduler uses Offer,
|
|
// Checkers use Claim, and ReportObservations fences facts against the lease.
|
|
// The production Redis implementation will use this exact contract.
|
|
type TaskBroker interface {
|
|
TaskSink
|
|
Claim(context.Context, TaskClaim) ([]LeasedTask, error)
|
|
AuthorizeObservation(context.Context, string, healthDomain.Observation, time.Time) error
|
|
CompleteObservation(context.Context, string, healthDomain.Observation, time.Time) error
|
|
}
|
|
|
|
type MemoryTaskBrokerOptions struct {
|
|
LeaseTTL time.Duration
|
|
MaxPerClaim int
|
|
Now func() time.Time
|
|
Material TaskMaterialResolver
|
|
}
|
|
|
|
// MemoryTaskBroker is a deterministic contract implementation for unit tests
|
|
// and single-process development. Production Controller instances use Redis so
|
|
// leases and limits remain shared across replicas.
|
|
type MemoryTaskBroker struct {
|
|
mu sync.Mutex
|
|
options MemoryTaskBrokerOptions
|
|
tasks map[string]brokerTask
|
|
leases map[string]taskLease
|
|
finished map[string]taskCompletion
|
|
}
|
|
|
|
type brokerTask struct {
|
|
plan PlannedTask
|
|
}
|
|
|
|
type taskLease struct {
|
|
checkerID string
|
|
instanceID string
|
|
expiresAt time.Time
|
|
}
|
|
|
|
type taskCompletion struct {
|
|
checkerID string
|
|
expiresAt time.Time
|
|
}
|
|
|
|
func NewMemoryTaskBroker(options MemoryTaskBrokerOptions) (*MemoryTaskBroker, error) {
|
|
if options.LeaseTTL <= 0 || options.Now == nil || nilTaskMaterialResolver(options.Material) {
|
|
return nil, ErrInvalidTaskBroker
|
|
}
|
|
if options.MaxPerClaim == 0 {
|
|
options.MaxPerClaim = defaultMaxTasksPerClaim
|
|
}
|
|
if options.MaxPerClaim <= 0 {
|
|
return nil, ErrInvalidTaskBroker
|
|
}
|
|
return &MemoryTaskBroker{
|
|
options: options, tasks: make(map[string]brokerTask), leases: make(map[string]taskLease),
|
|
finished: make(map[string]taskCompletion),
|
|
}, nil
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) Offer(ctx context.Context, plans []PlannedTask) (int, error) {
|
|
if ctx == nil || broker == nil || broker.options.Now == nil {
|
|
return 0, ErrInvalidTaskBroker
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return 0, err
|
|
}
|
|
if len(plans) == 0 {
|
|
return 0, nil
|
|
}
|
|
now := broker.options.Now().UTC()
|
|
if now.IsZero() {
|
|
return 0, ErrInvalidTaskBroker
|
|
}
|
|
for _, plan := range plans {
|
|
if err := validatePlannedTask(plan, now); err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
broker.mu.Lock()
|
|
defer broker.mu.Unlock()
|
|
broker.reapLocked(now)
|
|
offered := 0
|
|
for _, plan := range plans {
|
|
taskID := deterministicTaskID(plan)
|
|
if _, exists := broker.tasks[taskID]; exists {
|
|
continue
|
|
}
|
|
broker.tasks[taskID] = brokerTask{plan: plan}
|
|
offered++
|
|
}
|
|
return offered, nil
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) Claim(ctx context.Context, claim TaskClaim) ([]LeasedTask, error) {
|
|
if ctx == nil || broker == nil || broker.options.Now == nil || nilTaskMaterialResolver(broker.options.Material) {
|
|
return nil, ErrInvalidTaskBroker
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
levels, err := normalizeTaskClaim(claim, broker.options.MaxPerClaim)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
now := broker.options.Now().UTC()
|
|
if now.IsZero() {
|
|
return nil, ErrInvalidTaskBroker
|
|
}
|
|
broker.mu.Lock()
|
|
defer broker.mu.Unlock()
|
|
broker.reapLocked(now)
|
|
capacity := claim.MaxInFlight - broker.activeForCheckerLocked(claim.CheckerID)
|
|
if capacity <= 0 {
|
|
return []LeasedTask{}, nil
|
|
}
|
|
if capacity > broker.options.MaxPerClaim {
|
|
capacity = broker.options.MaxPerClaim
|
|
}
|
|
taskIDs := broker.claimableTaskIDsLocked(levels)
|
|
if len(taskIDs) > capacity {
|
|
taskIDs = taskIDs[:capacity]
|
|
}
|
|
if len(taskIDs) == 0 {
|
|
return []LeasedTask{}, nil
|
|
}
|
|
materials := make([]TaskMaterial, len(taskIDs))
|
|
for index, taskID := range taskIDs {
|
|
material, resolveErr := broker.options.Material.ResolveCheckTask(ctx, broker.tasks[taskID].plan.Candidate)
|
|
if resolveErr != nil {
|
|
return nil, resolveErr
|
|
}
|
|
if validateTaskMaterial(material) != nil {
|
|
return nil, ErrInvalidLeasedTask
|
|
}
|
|
materials[index] = material
|
|
}
|
|
result := make([]LeasedTask, len(taskIDs))
|
|
for index, taskID := range taskIDs {
|
|
plan := broker.tasks[taskID].plan
|
|
leaseUntil := now.Add(broker.options.LeaseTTL)
|
|
if plan.Deadline.Before(leaseUntil) {
|
|
leaseUntil = plan.Deadline
|
|
}
|
|
broker.leases[taskID] = taskLease{checkerID: claim.CheckerID, instanceID: claim.InstanceID, expiresAt: leaseUntil}
|
|
result[index] = leasedTask(taskID, plan, materials[index])
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) AuthorizeObservation(
|
|
ctx context.Context,
|
|
checkerID string,
|
|
observation healthDomain.Observation,
|
|
now time.Time,
|
|
) error {
|
|
if ctx == nil || broker == nil || !validTaskIdentifier(checkerID) || now.IsZero() {
|
|
return ErrInvalidTaskBroker
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
normalized, err := healthDomain.NormalizeObservation(observation)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
broker.mu.Lock()
|
|
defer broker.mu.Unlock()
|
|
return broker.authorizeObservationLocked(checkerID, normalized, now.UTC())
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) CompleteObservation(
|
|
ctx context.Context,
|
|
checkerID string,
|
|
observation healthDomain.Observation,
|
|
now time.Time,
|
|
) error {
|
|
if ctx == nil || broker == nil || !validTaskIdentifier(checkerID) || now.IsZero() {
|
|
return ErrInvalidTaskBroker
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
normalized, err := healthDomain.NormalizeObservation(observation)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
now = now.UTC()
|
|
broker.mu.Lock()
|
|
defer broker.mu.Unlock()
|
|
if err := broker.authorizeObservationLocked(checkerID, normalized, now); err != nil {
|
|
return err
|
|
}
|
|
if _, complete := broker.finished[normalized.TaskID]; complete {
|
|
return nil
|
|
}
|
|
task := broker.tasks[normalized.TaskID]
|
|
delete(broker.leases, normalized.TaskID)
|
|
broker.finished[normalized.TaskID] = taskCompletion{checkerID: checkerID, expiresAt: task.plan.Deadline}
|
|
return nil
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) authorizeObservationLocked(
|
|
checkerID string,
|
|
observation healthDomain.Observation,
|
|
now time.Time,
|
|
) error {
|
|
broker.reapLocked(now)
|
|
task, exists := broker.tasks[observation.TaskID]
|
|
if !exists {
|
|
return ErrTaskNotFound
|
|
}
|
|
if !taskMatchesObservation(task.plan, observation) {
|
|
return ErrTaskObservation
|
|
}
|
|
if finished, complete := broker.finished[observation.TaskID]; complete {
|
|
if finished.checkerID != checkerID {
|
|
return ErrTaskLeaseNotOwned
|
|
}
|
|
return nil
|
|
}
|
|
lease, leased := broker.leases[observation.TaskID]
|
|
if !leased || !lease.expiresAt.After(now) {
|
|
return ErrTaskLeaseExpired
|
|
}
|
|
if lease.checkerID != checkerID {
|
|
return ErrTaskLeaseNotOwned
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) reapLocked(now time.Time) {
|
|
for taskID, task := range broker.tasks {
|
|
if !task.plan.Deadline.After(now) {
|
|
delete(broker.tasks, taskID)
|
|
delete(broker.leases, taskID)
|
|
delete(broker.finished, taskID)
|
|
}
|
|
}
|
|
for taskID, lease := range broker.leases {
|
|
if !lease.expiresAt.After(now) {
|
|
delete(broker.leases, taskID)
|
|
}
|
|
}
|
|
for taskID, complete := range broker.finished {
|
|
if !complete.expiresAt.After(now) {
|
|
delete(broker.finished, taskID)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) activeForCheckerLocked(checkerID string) int {
|
|
active := 0
|
|
for _, lease := range broker.leases {
|
|
if lease.checkerID == checkerID {
|
|
active++
|
|
}
|
|
}
|
|
return active
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) claimableTaskIDsLocked(levels map[healthDomain.Level]struct{}) []string {
|
|
result := make([]string, 0, len(broker.tasks))
|
|
for taskID, task := range broker.tasks {
|
|
if _, leased := broker.leases[taskID]; leased {
|
|
continue
|
|
}
|
|
if _, complete := broker.finished[taskID]; complete {
|
|
continue
|
|
}
|
|
if _, supported := levels[task.plan.Candidate.Level]; supported {
|
|
result = append(result, taskID)
|
|
}
|
|
}
|
|
sort.Slice(result, func(left, right int) bool {
|
|
leftTask, rightTask := broker.tasks[result[left]].plan, broker.tasks[result[right]].plan
|
|
if leftTask.Priority != rightTask.Priority {
|
|
return leftTask.Priority < rightTask.Priority
|
|
}
|
|
if !leftTask.Candidate.DueAt.Equal(rightTask.Candidate.DueAt) {
|
|
return leftTask.Candidate.DueAt.Before(rightTask.Candidate.DueAt)
|
|
}
|
|
return result[left] < result[right]
|
|
})
|
|
return result
|
|
}
|
|
|
|
func normalizeTaskClaim(claim TaskClaim, maximum int) (map[healthDomain.Level]struct{}, error) {
|
|
if !validTaskIdentifier(claim.CheckerID) || !validTaskIdentifier(claim.InstanceID) || claim.MaxInFlight <= 0 ||
|
|
claim.MaxInFlight > maximum || len(claim.SupportedLevels) == 0 {
|
|
return nil, ErrInvalidTaskClaim
|
|
}
|
|
levels := make(map[healthDomain.Level]struct{}, len(claim.SupportedLevels))
|
|
for _, level := range claim.SupportedLevels {
|
|
switch level {
|
|
case healthDomain.LevelBasic, healthDomain.LevelEgress, healthDomain.LevelTarget:
|
|
default:
|
|
return nil, ErrInvalidTaskClaim
|
|
}
|
|
if _, duplicate := levels[level]; duplicate {
|
|
return nil, ErrInvalidTaskClaim
|
|
}
|
|
levels[level] = struct{}{}
|
|
}
|
|
return levels, nil
|
|
}
|
|
|
|
func validatePlannedTask(task PlannedTask, now time.Time) error {
|
|
if task.Deadline.IsZero() || !task.Deadline.After(now) || task.Attempts <= 0 || !validTaskIdentifier(task.Candidate.ProxyID) ||
|
|
task.Candidate.DueAt.IsZero() {
|
|
return ErrInvalidLeasedTask
|
|
}
|
|
switch task.Candidate.State {
|
|
case proxyDomain.StateFetched, proxyDomain.StateSuspect, proxyDomain.StateUnhealthy, proxyDomain.StateAvailable:
|
|
default:
|
|
return ErrInvalidLeasedTask
|
|
}
|
|
switch task.Candidate.Level {
|
|
case healthDomain.LevelBasic, healthDomain.LevelEgress:
|
|
if task.Candidate.RoutingName != "" || task.Candidate.TargetURL != "" {
|
|
return ErrInvalidLeasedTask
|
|
}
|
|
case healthDomain.LevelTarget:
|
|
if _, err := healthDomain.NormalizeTargetProfile(healthDomain.TargetProfile{
|
|
RoutingName: task.Candidate.RoutingName, TargetURL: task.Candidate.TargetURL,
|
|
}); err != nil {
|
|
return ErrInvalidLeasedTask
|
|
}
|
|
default:
|
|
return ErrInvalidLeasedTask
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func validateTaskMaterial(material TaskMaterial) error {
|
|
if strings.TrimSpace(material.Host) != material.Host || material.Host == "" || material.Port == 0 ||
|
|
(material.SecretRef == "") != (material.CredentialVersion == "") {
|
|
return ErrInvalidLeasedTask
|
|
}
|
|
switch material.Protocol {
|
|
case proxyDomain.SchemeHTTP, proxyDomain.SchemeHTTPS, proxyDomain.SchemeSOCKS5:
|
|
return nil
|
|
default:
|
|
return ErrInvalidLeasedTask
|
|
}
|
|
}
|
|
|
|
func deterministicTaskID(task PlannedTask) string {
|
|
payload := candidateIdentity(task.Candidate) + "\x00" + task.Deadline.UTC().Format(time.RFC3339Nano) + "\x00" +
|
|
strconv.Itoa(task.Attempts)
|
|
digest := sha256.Sum256([]byte(payload))
|
|
return "check_" + hex.EncodeToString(digest[:])
|
|
}
|
|
|
|
func leasedTask(taskID string, plan PlannedTask, material TaskMaterial) LeasedTask {
|
|
return LeasedTask{
|
|
TaskID: taskID, ProxyID: plan.Candidate.ProxyID, Protocol: material.Protocol, Host: material.Host, Port: material.Port,
|
|
SecretRef: material.SecretRef, CredentialVersion: material.CredentialVersion, Username: material.Username, Password: material.Password,
|
|
Level: plan.Candidate.Level, RoutingName: plan.Candidate.RoutingName, TargetURL: plan.Candidate.TargetURL,
|
|
Deadline: plan.Deadline.UTC(), Attempts: plan.Attempts,
|
|
}
|
|
}
|
|
|
|
func taskMatchesObservation(task PlannedTask, observation healthDomain.Observation) bool {
|
|
return task.Candidate.ProxyID == observation.ProxyID && task.Candidate.Level == observation.Level &&
|
|
task.Candidate.RoutingName == observation.RoutingName && task.Candidate.TargetURL == observation.TargetURL
|
|
}
|
|
|
|
func validTaskIdentifier(value string) bool {
|
|
if value == "" || len(value) > 256 || strings.TrimSpace(value) != value {
|
|
return false
|
|
}
|
|
for _, character := range value {
|
|
if character <= ' ' || character == '\x7f' {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
func nilTaskMaterialResolver(value TaskMaterialResolver) bool {
|
|
if value == nil {
|
|
return true
|
|
}
|
|
reflected := reflect.ValueOf(value)
|
|
switch reflected.Kind() {
|
|
case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice:
|
|
return reflected.IsNil()
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
var _ TaskBroker = (*MemoryTaskBroker)(nil)
|