proxy-pool/internal/checker/probe/probe_test.go
youfak 7740f6f14d
Some checks are pending
ci / proto (push) Waiting to run
ci / test (ubuntu-latest) (push) Waiting to run
ci / test (windows-latest) (push) Waiting to run
ci / race (push) Waiting to run
ci / integration (push) Waiting to run
feat: report observed egress addresses
2026-08-02 08:46:10 +08:00

353 lines
12 KiB
Go

package probe
import (
"context"
"encoding/binary"
"io"
"net"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"testing"
"time"
controlplanev1 "proxy-pool/gen/controlplane/v1"
"google.golang.org/protobuf/types/known/durationpb"
"google.golang.org/protobuf/types/known/timestamppb"
)
func TestExecutorBasicConfirmsProxyHandshakeEvenWhenProbeTargetFails(t *testing.T) {
proxy := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
if request.URL.String() != "http://example.invalid/" {
t.Errorf("proxy request URL = %q", request.URL)
}
if value := request.Header.Get("Proxy-Authorization"); value != "Basic dXNlcjpzZWNyZXQ=" {
t.Errorf("Proxy-Authorization = %q", value)
}
response.WriteHeader(http.StatusBadGateway)
}))
defer proxy.Close()
task := validTask(t, proxy.URL, controlplanev1.CheckLevel_CHECK_LEVEL_BASIC)
result := NewExecutor().Execute(context.Background(), task)
if !result.Success || result.FailureClass != "" || result.Latency <= 0 {
t.Fatalf("Execute(BASIC) = %+v", result)
}
}
func splitProxyAddress(t *testing.T, address string) (string, uint32) {
t.Helper()
host, rawPort, err := net.SplitHostPort(address[len("http://"):])
if err != nil {
t.Fatalf("net.SplitHostPort(%q): %v", address, err)
}
port, err := strconv.ParseUint(rawPort, 10, 16)
if err != nil {
t.Fatalf("strconv.ParseUint(%q): %v", rawPort, err)
}
return host, uint32(port)
}
func TestExecutorTargetReportsHTTPFailureAsFact(t *testing.T) {
proxy := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) {
response.WriteHeader(http.StatusServiceUnavailable)
}))
defer proxy.Close()
task := validTask(t, proxy.URL, controlplanev1.CheckLevel_CHECK_LEVEL_TARGET)
task.RoutingName = "route-a"
task.TargetUrl = "http://target.example/check"
result := NewExecutor().Execute(context.Background(), task)
if result.Success || result.FailureClass != FailureTargetHTTPStatus || result.Latency <= 0 {
t.Fatalf("Execute(TARGET) = %+v", result)
}
}
func TestExecutorRejectsEgressTaskWithRoutingProfile(t *testing.T) {
proxy := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) {
response.WriteHeader(http.StatusOK)
}))
defer proxy.Close()
task := validTask(t, proxy.URL, controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS)
task.RoutingName = "route-a"
task.TargetUrl = "https://egress.example/identity"
result := NewExecutor().Execute(context.Background(), task)
if result.Success || result.FailureClass != FailureInvalidTask {
t.Fatalf("Execute(EGRESS with routing) = %+v", result)
}
}
func TestExecutorEgressReportsObservedIP(t *testing.T) {
proxy := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
if request.URL.String() != "http://egress.example/identity" {
t.Errorf("proxy request URL = %q", request.URL)
}
response.Header().Set("Content-Type", "application/json")
_, _ = response.Write([]byte(`{"ip":"2001:db8::42"}`))
}))
defer proxy.Close()
task := validTask(t, proxy.URL, controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS)
task.TargetUrl = "http://egress.example/identity"
result := NewExecutor().Execute(context.Background(), task)
if !result.Success || result.FailureClass != "" || result.ObservedEgressIP != "2001:db8::42" || result.Latency <= 0 {
t.Fatalf("Execute(EGRESS) = %+v", result)
}
}
func TestExecutorEgressRejectsMissingIdentity(t *testing.T) {
proxy := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) {
_, _ = response.Write([]byte(`{"status":"ok"}`))
}))
defer proxy.Close()
task := validTask(t, proxy.URL, controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS)
task.TargetUrl = "http://egress.example/identity"
result := NewExecutor().Execute(context.Background(), task)
if result.Success || result.FailureClass != FailureEgressIdentity || result.ObservedEgressIP != "" || result.Latency <= 0 {
t.Fatalf("Execute(EGRESS missing identity) = %+v", result)
}
}
func TestParseEgressIP(t *testing.T) {
for _, test := range []struct {
name string
body string
want string
}{
{name: "plain IPv4", body: "198.51.100.7\n", want: "198.51.100.7"},
{name: "JSON query", body: `{"query":"2001:db8::7"}`, want: "2001:db8::7"},
{name: "JSON origin list", body: `{"origin":"198.51.100.8, 2001:db8::8"}`, want: "198.51.100.8"},
{name: "invalid", body: `{"ip":"unknown"}`},
} {
t.Run(test.name, func(t *testing.T) {
got, ok := parseEgressIP([]byte(test.body))
if got != test.want || ok != (test.want != "") {
t.Fatalf("parseEgressIP(%q) = (%q, %t), want (%q, %t)", test.body, got, ok, test.want, test.want != "")
}
})
}
}
func TestExecutorReportsUnsupportedProtocolAsFact(t *testing.T) {
now := time.Now().UTC()
result := NewExecutor().Execute(context.Background(), &controlplanev1.CheckTask{
TaskId: "task-a", ProxyId: "proxy-a", Protocol: controlplanev1.ProxyProtocol_PROXY_PROTOCOL_UNSPECIFIED,
Host: "proxy.example", Port: 1080, Level: controlplanev1.CheckLevel_CHECK_LEVEL_BASIC,
Timeout: durationpb.New(time.Second), Attempt: 1, MaxAttempts: 1, Deadline: timestamppb.New(now.Add(time.Second)),
})
if result.Success || result.FailureClass != FailureUnsupportedProxy {
t.Fatalf("Execute(SOCKS5) = %+v", result)
}
}
func TestExecutorBasicSupportsAuthenticatedSOCKS5Proxy(t *testing.T) {
target := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
if request.Host != "example.invalid" {
t.Errorf("request host = %q", request.Host)
}
response.WriteHeader(http.StatusBadGateway)
}))
defer target.Close()
proxy := newSOCKS5TestProxy(t, target.URL, "user", "secret")
task := validSOCKS5Task(t, proxy.address, "user", "secret")
result := NewExecutor().Execute(context.Background(), task)
if !result.Success || result.FailureClass != "" || result.Latency <= 0 {
t.Fatalf("Execute(SOCKS5 BASIC) = %+v", result)
}
select {
case address := <-proxy.requested:
if address != "example.invalid:80" {
t.Fatalf("SOCKS5 requested address = %q", address)
}
case <-time.After(time.Second):
t.Fatal("SOCKS5 proxy did not receive CONNECT request")
}
}
func TestExecutorReportsSOCKS5AuthenticationFailure(t *testing.T) {
target := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
defer target.Close()
proxy := newSOCKS5TestProxy(t, target.URL, "user", "secret")
result := NewExecutor().Execute(context.Background(), validSOCKS5Task(t, proxy.address, "user", "wrong"))
if result.Success || result.FailureClass != FailureProxyAuth || result.Latency <= 0 {
t.Fatalf("Execute(SOCKS5 auth failure) = %+v", result)
}
}
func validTask(t *testing.T, proxyAddress string, level controlplanev1.CheckLevel) *controlplanev1.CheckTask {
t.Helper()
host, port := splitProxyAddress(t, proxyAddress)
now := time.Now().UTC()
return &controlplanev1.CheckTask{
TaskId: "task-a", ProxyId: "proxy-a", Protocol: controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTP,
Host: host, Port: port, Level: level, Username: "user", Password: "secret",
Timeout: durationpb.New(time.Second), Attempt: 1, MaxAttempts: 1,
Deadline: timestamppb.New(now.Add(time.Second)),
}
}
func validSOCKS5Task(t *testing.T, address, username, password string) *controlplanev1.CheckTask {
t.Helper()
host, rawPort, err := net.SplitHostPort(address)
if err != nil {
t.Fatalf("net.SplitHostPort(%q): %v", address, err)
}
port, err := strconv.ParseUint(rawPort, 10, 16)
if err != nil {
t.Fatalf("strconv.ParseUint(%q): %v", rawPort, err)
}
now := time.Now().UTC()
return &controlplanev1.CheckTask{
TaskId: "task-socks", ProxyId: "proxy-socks", Protocol: controlplanev1.ProxyProtocol_PROXY_PROTOCOL_SOCKS5,
Host: host, Port: uint32(port), Level: controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, Username: username, Password: password,
Timeout: durationpb.New(time.Second), Attempt: 1, MaxAttempts: 1, Deadline: timestamppb.New(now.Add(time.Second)),
}
}
type socks5TestProxy struct {
address string
target string
username string
password string
requested chan string
listener net.Listener
}
func newSOCKS5TestProxy(t *testing.T, target, username, password string) *socks5TestProxy {
t.Helper()
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("net.Listen(): %v", err)
}
proxy := &socks5TestProxy{
address: listener.Addr().String(), target: target, username: username, password: password,
requested: make(chan string, 1), listener: listener,
}
go func() {
for {
connection, err := listener.Accept()
if err != nil {
return
}
go proxy.serve(connection)
}
}()
t.Cleanup(func() { _ = listener.Close() })
return proxy
}
func (proxy *socks5TestProxy) serve(connection net.Conn) {
defer connection.Close()
if _, err := readSOCKS5Greeting(connection); err != nil {
return
}
method := byte(0x00)
if proxy.username != "" || proxy.password != "" {
method = 0x02
}
if _, err := connection.Write([]byte{0x05, method}); err != nil {
return
}
if method == 0x02 {
username, password, err := readSOCKS5Credentials(connection)
if err != nil {
return
}
if username != proxy.username || password != proxy.password {
_, _ = connection.Write([]byte{0x01, 0x01})
return
}
if _, err := connection.Write([]byte{0x01, 0x00}); err != nil {
return
}
}
address, err := readSOCKS5ConnectRequest(connection)
if err != nil {
return
}
proxy.requested <- address
targetAddress := strings.TrimPrefix(proxy.target, "http://")
target, err := net.Dial("tcp", targetAddress)
if err != nil {
_, _ = connection.Write([]byte{0x05, 0x01, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
return
}
defer target.Close()
if _, err := connection.Write([]byte{0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0}); err != nil {
return
}
done := make(chan struct{}, 2)
go func() { _, _ = io.Copy(target, connection); done <- struct{}{} }()
go func() { _, _ = io.Copy(connection, target); done <- struct{}{} }()
<-done
}
func readSOCKS5Greeting(reader io.Reader) ([]byte, error) {
header := make([]byte, 2)
if _, err := io.ReadFull(reader, header); err != nil || header[0] != 0x05 || header[1] == 0 {
return nil, io.ErrUnexpectedEOF
}
methods := make([]byte, header[1])
_, err := io.ReadFull(reader, methods)
return methods, err
}
func readSOCKS5Credentials(reader io.Reader) (string, string, error) {
header := make([]byte, 2)
if _, err := io.ReadFull(reader, header); err != nil || header[0] != 0x01 || header[1] == 0 {
return "", "", io.ErrUnexpectedEOF
}
username := make([]byte, header[1])
if _, err := io.ReadFull(reader, username); err != nil {
return "", "", err
}
passwordLength := make([]byte, 1)
if _, err := io.ReadFull(reader, passwordLength); err != nil || passwordLength[0] == 0 {
return "", "", io.ErrUnexpectedEOF
}
password := make([]byte, passwordLength[0])
if _, err := io.ReadFull(reader, password); err != nil {
return "", "", err
}
return string(username), string(password), nil
}
func readSOCKS5ConnectRequest(reader io.Reader) (string, error) {
header := make([]byte, 4)
if _, err := io.ReadFull(reader, header); err != nil || header[0] != 0x05 || header[1] != 0x01 || header[2] != 0x00 {
return "", io.ErrUnexpectedEOF
}
var host string
switch header[3] {
case 0x01:
address := make([]byte, net.IPv4len)
if _, err := io.ReadFull(reader, address); err != nil {
return "", err
}
host = net.IP(address).String()
case 0x03:
length := make([]byte, 1)
if _, err := io.ReadFull(reader, length); err != nil || length[0] == 0 {
return "", io.ErrUnexpectedEOF
}
address := make([]byte, length[0])
if _, err := io.ReadFull(reader, address); err != nil {
return "", err
}
host = string(address)
case 0x04:
address := make([]byte, net.IPv6len)
if _, err := io.ReadFull(reader, address); err != nil {
return "", err
}
host = net.IP(address).String()
default:
return "", io.ErrUnexpectedEOF
}
port := make([]byte, 2)
if _, err := io.ReadFull(reader, port); err != nil {
return "", err
}
return net.JoinHostPort(host, strconv.Itoa(int(binary.BigEndian.Uint16(port)))), nil
}