proxy-pool/internal/checker/probe/probe.go

// Package probe executes one bounded Checker task and returns only a health
// fact. It has no Controller, Redis, or PostgreSQL dependency.
package probe

import (
	"context"
	"errors"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	controlplanev1 "proxy-pool/gen/controlplane/v1"
)

const (
	FailureInvalidTask      = "INVALID_TASK"
	FailureDeadlineExceeded = "DEADLINE_EXCEEDED"
	FailureProxyRequest     = "PROXY_REQUEST"
	FailureProxyAuth        = "PROXY_AUTH"
	FailureTargetHTTPStatus = "TARGET_HTTP_STATUS"
	FailureUnsupportedProxy = "UNSUPPORTED_PROXY_PROTOCOL"
	basicHandshakeProbeURL  = "http://example.invalid/"
)

var errUnsupportedProtocol = errors.New("unsupported proxy protocol")

type Result struct {
	Success      bool
	FailureClass string
	Latency      time.Duration
}

// Executor constructs a short-lived HTTP transport for each task. That keeps
// credentials task-scoped and avoids keeping stale short-TTL proxy sessions
// alive after a task completes.
type Executor struct{}

func NewExecutor() *Executor { return &Executor{} }

// Execute reports protocol/authorization reachability for BASIC and requires
// a successful target response for TARGET. Every path returns one fact so a
// Checker can report failures without treating ordinary probe failures as
// control-plane errors.
func (executor *Executor) Execute(ctx context.Context, task *controlplanev1.CheckTask) Result {
	started := time.Now()
	if ctx == nil || executor == nil {
		return Result{FailureClass: FailureInvalidTask}
	}
	deadline, target, proxyURL, err := prepare(task, started.UTC())
	if err != nil {
		if errors.Is(err, errUnsupportedProtocol) {
			return Result{FailureClass: FailureUnsupportedProxy, Latency: time.Since(started)}
		}
		return Result{FailureClass: FailureInvalidTask, Latency: time.Since(started)}
	}
	probeContext, cancel := context.WithDeadline(ctx, deadline)
	defer cancel()
	transport := &http.Transport{
		Proxy:                 http.ProxyURL(proxyURL),
		DialContext:           (&net.Dialer{}).DialContext,
		ForceAttemptHTTP2:     false,
		TLSHandshakeTimeout:   minDuration(time.Until(deadline), 10*time.Second),
		ResponseHeaderTimeout: minDuration(time.Until(deadline), 10*time.Second),
	}
	defer transport.CloseIdleConnections()
	client := &http.Client{
		Transport:     transport,
		CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
	}
	request, err := http.NewRequestWithContext(probeContext, http.MethodGet, target, nil)
	if err != nil {
		return Result{FailureClass: FailureInvalidTask, Latency: time.Since(started)}
	}
	response, err := client.Do(request)
	latency := time.Since(started)
	if err != nil {
		if errors.Is(probeContext.Err(), context.DeadlineExceeded) || errors.Is(ctx.Err(), context.DeadlineExceeded) {
			return Result{FailureClass: FailureDeadlineExceeded, Latency: latency}
		}
		return Result{FailureClass: FailureProxyRequest, Latency: latency}
	}
	_ = response.Body.Close()
	if response.StatusCode == http.StatusProxyAuthRequired {
		return Result{FailureClass: FailureProxyAuth, Latency: latency}
	}
	if task.GetLevel() == controlplanev1.CheckLevel_CHECK_LEVEL_BASIC {
		return Result{Success: true, Latency: latency}
	}
	if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusBadRequest {
		return Result{FailureClass: FailureTargetHTTPStatus, Latency: latency}
	}
	return Result{Success: true, Latency: latency}
}

func prepare(task *controlplanev1.CheckTask, now time.Time) (time.Time, string, *url.URL, error) {
	if task == nil || task.GetTaskId() == "" || task.GetProxyId() == "" || task.GetHost() == "" || task.GetPort() == 0 ||
		task.GetTimeout() == nil || task.GetTimeout().CheckValid() != nil || task.GetTimeout().AsDuration() <= 0 ||
		task.GetDeadline() == nil || task.GetDeadline().CheckValid() != nil || task.GetDeadline().AsTime().UTC().Compare(now) <= 0 ||
		task.GetAttempt() == 0 || task.GetMaxAttempts() == 0 || task.GetAttempt() > task.GetMaxAttempts() ||
		(task.GetSecretRef() == "") != (task.GetCredentialVersion() == "") {
		return time.Time{}, "", nil, errors.New("invalid checker task")
	}
	proxyScheme := ""
	switch task.GetProtocol() {
	case controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTP:
		proxyScheme = "http"
	case controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTPS:
		proxyScheme = "https"
	default:
		return time.Time{}, "", nil, errUnsupportedProtocol
	}
	deadline := task.GetDeadline().AsTime().UTC()
	timeoutDeadline := now.Add(task.GetTimeout().AsDuration())
	if timeoutDeadline.Before(deadline) {
		deadline = timeoutDeadline
	}
	if !deadline.After(now) {
		return time.Time{}, "", nil, errors.New("expired checker task")
	}
	proxyURL := &url.URL{Scheme: proxyScheme, Host: net.JoinHostPort(task.GetHost(), strconv.FormatUint(uint64(task.GetPort()), 10))}
	if task.GetUsername() != "" || task.GetPassword() != "" {
		proxyURL.User = url.UserPassword(task.GetUsername(), task.GetPassword())
	}
	target := basicHandshakeProbeURL
	switch task.GetLevel() {
	case controlplanev1.CheckLevel_CHECK_LEVEL_BASIC:
		if task.GetRoutingName() != "" || task.GetTargetUrl() != "" {
			return time.Time{}, "", nil, errors.New("invalid basic task")
		}
	case controlplanev1.CheckLevel_CHECK_LEVEL_TARGET:
		if task.GetRoutingName() == "" {
			return time.Time{}, "", nil, errors.New("target routing is required")
		}
		target = task.GetTargetUrl()
	case controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS:
		target = task.GetTargetUrl()
	default:
		return time.Time{}, "", nil, errors.New("unsupported task level")
	}
	parsedTarget, err := url.Parse(target)
	if err != nil || parsedTarget.Scheme == "" || parsedTarget.Host == "" || parsedTarget.User != nil || parsedTarget.Fragment != "" ||
		(parsedTarget.Scheme != "http" && parsedTarget.Scheme != "https") || strings.TrimSpace(task.GetHost()) != task.GetHost() {
		return time.Time{}, "", nil, errors.New("invalid probe target")
	}
	return deadline, parsedTarget.String(), proxyURL, nil
}

func minDuration(left, right time.Duration) time.Duration {
	if left <= 0 || left < right {
		return left
	}
	return right
}