netbird/proxy/internal/appsec/client_test.go

package appsec

import (
	"context"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/netbirdio/netbird/proxy/internal/restrict"
)

// engine records what the AppSec component received and replies with a
// canned status and body.
type engine struct {
	status int
	body   string

	gotMethod string
	gotHeader http.Header
	gotBody   []byte
	gotLength int64
	requests  int
}

func (e *engine) start(t *testing.T) *httptest.Server {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, err := io.ReadAll(r.Body)
		require.NoError(t, err)
		e.requests++
		e.gotMethod = r.Method
		e.gotHeader = r.Header.Clone()
		e.gotBody = body
		e.gotLength = r.ContentLength

		status := e.status
		if status == 0 {
			status = http.StatusOK
		}
		w.WriteHeader(status)
		if e.body != "" {
			_, _ = w.Write([]byte(e.body))
		}
	}))
	t.Cleanup(srv.Close)
	return srv
}

func newClient(t *testing.T, url string, cfg ...func(*Config)) *Client {
	t.Helper()
	c := Config{URL: url, APIKey: "test-key"}
	for _, fn := range cfg {
		fn(&c)
	}
	client, err := New(c)
	require.NoError(t, err)
	return client
}

func inbound(method, target string, body string) *http.Request {
	var r *http.Request
	if body == "" {
		r = httptest.NewRequest(method, target, nil)
	} else {
		r = httptest.NewRequest(method, target, strings.NewReader(body))
	}
	r.Host = "svc.example.com"
	r.Header.Set("User-Agent", "curl/8.0")
	return r
}

func TestNew_RejectsBadConfig(t *testing.T) {
	tests := []struct {
		name string
		cfg  Config
	}{
		{"empty url", Config{APIKey: "k"}},
		{"empty key", Config{URL: "http://127.0.0.1:7422/"}},
		{"non http scheme", Config{URL: "tcp://127.0.0.1:7422", APIKey: "k"}},
		{"no host", Config{URL: "http:///path", APIKey: "k"}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, err := New(tt.cfg)
			require.Error(t, err)
		})
	}
}

func TestInspect_AllowSendsProtocolHeaders(t *testing.T) {
	eng := &engine{status: http.StatusOK, body: `{"action":"allow","http_status":200}`}
	srv := eng.start(t)
	client := newClient(t, srv.URL)

	r := inbound(http.MethodGet, "http://svc.example.com/admin?q=1", "")
	r.Header.Set("Cookie", "session=abc")

	verdict, err := client.Inspect(context.Background(), Request{
		HTTP:          r,
		ClientIP:      netip.MustParseAddr("203.0.113.7"),
		TransactionID: "req-42",
	})
	require.NoError(t, err)
	assert.Equal(t, restrict.Allow, verdict, "allow action must pass the request")

	assert.Equal(t, http.MethodGet, eng.gotMethod, "a bodyless request is forwarded as GET")
	assert.Equal(t, "203.0.113.7", eng.gotHeader.Get(headerIP))
	assert.Equal(t, "/admin?q=1", eng.gotHeader.Get(headerURI), "URI header must carry path and query")
	assert.Equal(t, http.MethodGet, eng.gotHeader.Get(headerVerb))
	assert.Equal(t, "svc.example.com", eng.gotHeader.Get(headerHost))
	assert.Equal(t, "curl/8.0", eng.gotHeader.Get(headerUserAgent))
	assert.Equal(t, "11", eng.gotHeader.Get(headerHTTPVersion))
	assert.Equal(t, "req-42", eng.gotHeader.Get(headerTransactionID))
	assert.Equal(t, "test-key", eng.gotHeader.Get(headerAPIKey))
	// Client headers are what the WAF rules match on.
	assert.Equal(t, "session=abc", eng.gotHeader.Get("Cookie"))
}

func TestInspect_MapsActionsToVerdicts(t *testing.T) {
	tests := []struct {
		name    string
		status  int
		body    string
		want    restrict.Verdict
		wantErr bool
	}{
		{"allow", http.StatusOK, `{"action":"allow"}`, restrict.Allow, false},
		{"ban", http.StatusForbidden, `{"action":"ban","http_status":403}`, restrict.DenyAppSecBan, false},
		{"captcha", http.StatusForbidden, `{"action":"captcha","http_status":403}`, restrict.DenyAppSecCaptcha, false},
		// blocked_http_code is operator-configurable, so the action decides.
		{"custom block status", http.StatusTeapot, `{"action":"ban"}`, restrict.DenyAppSecBan, false},
		{"allow on custom status", http.StatusTeapot, `{"action":"allow"}`, restrict.Allow, false},
		{"unknown action denies", http.StatusForbidden, `{"action":"something-new"}`, restrict.DenyAppSecBan, false},
		{"unreadable body denies", http.StatusForbidden, `not json`, restrict.DenyAppSecBan, false},
		{"bad api key is unavailable", http.StatusUnauthorized, "", restrict.DenyAppSecUnavailable, true},
		{"engine error is unavailable", http.StatusInternalServerError, "", restrict.DenyAppSecUnavailable, true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			eng := &engine{status: tt.status, body: tt.body}
			srv := eng.start(t)
			client := newClient(t, srv.URL)

			verdict, err := client.Inspect(context.Background(), Request{
				HTTP:     inbound(http.MethodGet, "http://svc.example.com/", ""),
				ClientIP: netip.MustParseAddr("203.0.113.7"),
			})
			if tt.wantErr {
				require.Error(t, err)
				assert.True(t, errors.Is(err, ErrUnavailable), "engine-side failures must be ErrUnavailable so the caller can apply the mode")
			} else {
				require.NoError(t, err)
			}
			assert.Equal(t, tt.want, verdict)
		})
	}
}

func TestInspect_ForwardsBodyAndRestoresIt(t *testing.T) {
	eng := &engine{body: `{"action":"allow"}`}
	srv := eng.start(t)
	client := newClient(t, srv.URL)

	const payload = `{"user":"' OR 1=1--"}`
	r := inbound(http.MethodPost, "http://svc.example.com/login", payload)
	r.Header.Set("Content-Type", "application/json")

	_, err := client.Inspect(context.Background(), Request{
		HTTP:     r,
		ClientIP: netip.MustParseAddr("203.0.113.7"),
	})
	require.NoError(t, err)

	assert.Equal(t, http.MethodPost, eng.gotMethod, "a request with a body is forwarded as POST")
	assert.Equal(t, payload, string(eng.gotBody))
	assert.Equal(t, int64(len(payload)), eng.gotLength,
		"the engine reads exactly Content-Length bytes, so it must be accurate")

	// The backend still needs the body.
	restored, err := io.ReadAll(r.Body)
	require.NoError(t, err)
	assert.Equal(t, payload, string(restored), "body must be restored for the upstream request")
}

func TestInspect_OversizeBodyFallsBackToHeaders(t *testing.T) {
	eng := &engine{body: `{"action":"allow"}`}
	srv := eng.start(t)
	client := newClient(t, srv.URL, func(c *Config) { c.MaxBodyBytes = 16 })

	payload := strings.Repeat("A", 64)
	r := inbound(http.MethodPost, "http://svc.example.com/upload", payload)

	_, err := client.Inspect(context.Background(), Request{
		HTTP:     r,
		ClientIP: netip.MustParseAddr("203.0.113.7"),
	})
	require.NoError(t, err)

	assert.Equal(t, http.MethodGet, eng.gotMethod, "an oversize body is not forwarded")
	assert.Empty(t, eng.gotBody, "a truncated prefix must never be sent: it changes the verdict")

	restored, err := io.ReadAll(r.Body)
	require.NoError(t, err)
	assert.Equal(t, payload, string(restored), "the full body must still reach the upstream")
}

func TestInspect_ChunkedOversizeBodyIsReplayed(t *testing.T) {
	eng := &engine{body: `{"action":"allow"}`}
	srv := eng.start(t)
	client := newClient(t, srv.URL, func(c *Config) { c.MaxBodyBytes = 8 })

	payload := strings.Repeat("B", 40)
	r := inbound(http.MethodPost, "http://svc.example.com/upload", payload)
	// Unknown length: the cap can only be detected while reading.
	r.ContentLength = -1
	r.Header.Del("Content-Length")

	_, err := client.Inspect(context.Background(), Request{
		HTTP:     r,
		ClientIP: netip.MustParseAddr("203.0.113.7"),
	})
	require.NoError(t, err)

	assert.Empty(t, eng.gotBody)
	restored, err := io.ReadAll(r.Body)
	require.NoError(t, err)
	assert.Equal(t, payload, string(restored), "bytes read to detect the cap must be replayed")
}

func TestInspect_SkipsCredentialForm(t *testing.T) {
	eng := &engine{body: `{"action":"allow"}`}
	srv := eng.start(t)
	client := newClient(t, srv.URL)

	r := inbound(http.MethodPost, "http://svc.example.com/", "password=hunter2&next=%2Fhome")
	r.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	_, err := client.Inspect(context.Background(), Request{
		HTTP:           r,
		ClientIP:       netip.MustParseAddr("203.0.113.7"),
		OmitBodyFields: []string{"password", "pin"},
	})
	require.NoError(t, err)

	assert.Empty(t, eng.gotBody, "a login form body must not reach the engine")
	assert.NotContains(t, string(eng.gotBody), "hunter2")
	// The request is still inspected on headers and URI.
	assert.Equal(t, 1, eng.requests)

	restored, err := io.ReadAll(r.Body)
	require.NoError(t, err)
	assert.Equal(t, "password=hunter2&next=%2Fhome", string(restored),
		"the login handler still needs to read the form")
}

func TestInspect_ForwardsNonCredentialForm(t *testing.T) {
	eng := &engine{body: `{"action":"allow"}`}
	srv := eng.start(t)
	client := newClient(t, srv.URL)

	r := inbound(http.MethodPost, "http://svc.example.com/search", "q=%3Cscript%3E")
	r.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	_, err := client.Inspect(context.Background(), Request{
		HTTP:           r,
		ClientIP:       netip.MustParseAddr("203.0.113.7"),
		OmitBodyFields: []string{"password", "pin"},
	})
	require.NoError(t, err)
	assert.Equal(t, "q=%3Cscript%3E", string(eng.gotBody), "ordinary form bodies must be inspected")
}

func TestInspect_StripsSpoofedProtocolHeaders(t *testing.T) {
	eng := &engine{body: `{"action":"allow"}`}
	srv := eng.start(t)
	client := newClient(t, srv.URL)

	r := inbound(http.MethodGet, "http://svc.example.com/", "")
	// A caller trying to make the engine see a different source address, and to
	// smuggle in its own key.
	r.Header.Set(headerIP, "10.0.0.1")
	r.Header.Set(headerAPIKey, "attacker-key")
	r.Header.Set(headerURI, "/harmless")

	_, err := client.Inspect(context.Background(), Request{
		HTTP:     r,
		ClientIP: netip.MustParseAddr("203.0.113.7"),
	})
	require.NoError(t, err)

	assert.Equal(t, "203.0.113.7", eng.gotHeader.Get(headerIP), "the proxy's resolved IP must win")
	assert.Equal(t, "test-key", eng.gotHeader.Get(headerAPIKey))
	assert.Equal(t, "/", eng.gotHeader.Get(headerURI))
	assert.Len(t, eng.gotHeader.Values(headerIP), 1, "no duplicate protocol headers")
}

func TestInspect_DropsHopByHopHeaders(t *testing.T) {
	eng := &engine{body: `{"action":"allow"}`}
	srv := eng.start(t)
	client := newClient(t, srv.URL)

	r := inbound(http.MethodGet, "http://svc.example.com/", "")
	r.Header.Set("Proxy-Authorization", "Basic zzz")
	r.Header.Set("Keep-Alive", "timeout=5")

	_, err := client.Inspect(context.Background(), Request{
		HTTP:     r,
		ClientIP: netip.MustParseAddr("203.0.113.7"),
	})
	require.NoError(t, err)

	assert.Empty(t, eng.gotHeader.Get("Proxy-Authorization"))
	assert.Empty(t, eng.gotHeader.Get("Keep-Alive"))
}

func TestInspect_UpgradeRequestKeepsStreamIntact(t *testing.T) {
	eng := &engine{body: `{"action":"allow"}`}
	srv := eng.start(t)
	client := newClient(t, srv.URL)

	r := inbound(http.MethodGet, "http://svc.example.com/ws", "")
	r.Header.Set("Upgrade", "websocket")
	r.Header.Set("Connection", "upgrade")

	_, err := client.Inspect(context.Background(), Request{
		HTTP:     r,
		ClientIP: netip.MustParseAddr("203.0.113.7"),
	})
	require.NoError(t, err)
	assert.Equal(t, http.MethodGet, eng.gotMethod)
	assert.Empty(t, eng.gotBody)
	assert.Equal(t, 1, eng.requests, "upgrade requests are still inspected on headers")
}

func TestInspect_TimeoutIsUnavailable(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		time.Sleep(200 * time.Millisecond)
		w.WriteHeader(http.StatusOK)
	}))
	t.Cleanup(srv.Close)

	client := newClient(t, srv.URL, func(c *Config) { c.Timeout = 10 * time.Millisecond })

	verdict, err := client.Inspect(context.Background(), Request{
		HTTP:     inbound(http.MethodGet, "http://svc.example.com/", ""),
		ClientIP: netip.MustParseAddr("203.0.113.7"),
	})
	require.Error(t, err)
	assert.True(t, errors.Is(err, ErrUnavailable))
	assert.Equal(t, restrict.DenyAppSecUnavailable, verdict)
}

func TestInspect_UnreachableEngineIsUnavailable(t *testing.T) {
	// Port 1 on loopback refuses connections.
	client := newClient(t, "http://127.0.0.1:1/")

	verdict, err := client.Inspect(context.Background(), Request{
		HTTP:     inbound(http.MethodGet, "http://svc.example.com/", ""),
		ClientIP: netip.MustParseAddr("203.0.113.7"),
	})
	require.Error(t, err)
	assert.True(t, errors.Is(err, ErrUnavailable))
	assert.Equal(t, restrict.DenyAppSecUnavailable, verdict)
}

func TestInspect_NilClientFailsClosed(t *testing.T) {
	var client *Client
	verdict, err := client.Inspect(context.Background(), Request{
		HTTP:     inbound(http.MethodGet, "http://svc.example.com/", ""),
		ClientIP: netip.MustParseAddr("203.0.113.7"),
	})
	require.Error(t, err)
	assert.Equal(t, restrict.DenyAppSecUnavailable, verdict)
}

func TestInspect_NegativeCapDisablesBodyForwarding(t *testing.T) {
	eng := &engine{body: `{"action":"allow"}`}
	srv := eng.start(t)
	client := newClient(t, srv.URL, func(c *Config) { c.MaxBodyBytes = -1 })

	const payload = `{"a":1}`
	r := inbound(http.MethodPost, "http://svc.example.com/", payload)

	_, err := client.Inspect(context.Background(), Request{
		HTTP:     r,
		ClientIP: netip.MustParseAddr("203.0.113.7"),
	})
	require.NoError(t, err)
	assert.Empty(t, eng.gotBody)

	restored, err := io.ReadAll(r.Body)
	require.NoError(t, err)
	assert.Equal(t, payload, string(restored), "an untouched body must still be forwardable")
}

func TestInspect_MapsV4MappedClientIP(t *testing.T) {
	eng := &engine{body: `{"action":"allow"}`}
	srv := eng.start(t)
	client := newClient(t, srv.URL)

	_, err := client.Inspect(context.Background(), Request{
		HTTP:     inbound(http.MethodGet, "http://svc.example.com/", ""),
		ClientIP: netip.MustParseAddr("::ffff:203.0.113.7"),
	})
	require.NoError(t, err)
	assert.Equal(t, "203.0.113.7", eng.gotHeader.Get(headerIP),
		"v4-mapped addresses must be unmapped so engine allowlists and rules match")
}