303 lines
9.9 KiB
Go
303 lines
9.9 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 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
|
|
}
|