crowdsec-bouncer-traefik-plugin/bouncer_test.go

package crowdsec_bouncer_traefik_plugin //nolint:revive,stylecheck

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"reflect"
	"strings"
	"sync"
	"testing"
	"text/template"
	"time"

	cache "github.com/maxlerebourg/crowdsec-bouncer-traefik-plugin/pkg/cache"
	configuration "github.com/maxlerebourg/crowdsec-bouncer-traefik-plugin/pkg/configuration"
	ip "github.com/maxlerebourg/crowdsec-bouncer-traefik-plugin/pkg/ip"
	logger "github.com/maxlerebourg/crowdsec-bouncer-traefik-plugin/pkg/logger"
)

func TestServeHTTP(t *testing.T) {
	cfg := CreateConfig()
	cfg.CrowdsecLapiKey = "test"

	ctx := context.Background()
	next := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {})

	handler, err := New(ctx, next, cfg, "demo-plugin")
	if err != nil {
		t.Fatal(err)
	}

	recorder := httptest.NewRecorder()
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://localhost", nil)
	if err != nil {
		t.Fatal(err)
	}

	handler.ServeHTTP(recorder, req)
}

func TestNew(t *testing.T) {
	type args struct {
		ctx    context.Context //nolint:containedctx
		next   http.Handler
		config *configuration.Config
		name   string
	}
	tests := []struct {
		name    string
		args    args
		want    http.Handler
		wantErr bool
	}{
		// TODO: Add test cases.
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := New(tt.args.ctx, tt.args.next, tt.args.config, tt.args.name)
			if (err != nil) != tt.wantErr {
				t.Errorf("New() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("New() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestBouncer_ServeHTTP(t *testing.T) {
	type fields struct {
		next                   http.Handler
		name                   string
		template               *template.Template
		enabled                bool
		crowdsecScheme         string
		crowdsecHost           string
		crowdsecKey            string
		crowdsecMode           string
		updateInterval         int64
		defaultDecisionTimeout int64
		forwardedCustomHeader  string
		clientPoolStrategy     *ip.PoolStrategy
		serverPoolStrategy     *ip.PoolStrategy
		httpClient             *http.Client
		cacheClient            *cache.Client
	}
	type args struct {
		rw  http.ResponseWriter
		req *http.Request
	}
	tests := []struct {
		name   string
		fields fields
		args   args
	}{
		// TODO: Add test cases.
	}
	for _, tt := range tests {
		t.Run(tt.name, func(_ *testing.T) {
			bouncer := &Bouncer{
				next:                   tt.fields.next,
				name:                   tt.fields.name,
				template:               tt.fields.template,
				enabled:                tt.fields.enabled,
				crowdsecScheme:         tt.fields.crowdsecScheme,
				crowdsecHost:           tt.fields.crowdsecHost,
				crowdsecKey:            tt.fields.crowdsecKey,
				crowdsecMode:           tt.fields.crowdsecMode,
				updateInterval:         tt.fields.updateInterval,
				defaultDecisionTimeout: tt.fields.defaultDecisionTimeout,
				forwardedCustomHeader:  tt.fields.forwardedCustomHeader,
				clientPoolStrategy:     tt.fields.clientPoolStrategy,
				serverPoolStrategy:     tt.fields.serverPoolStrategy,
				httpClient:             tt.fields.httpClient,
				cacheClient:            tt.fields.cacheClient,
			}
			bouncer.ServeHTTP(tt.args.rw, tt.args.req)
		})
	}
}

func Test_handleNoStreamCache(t *testing.T) {
	type args struct {
		bouncer  *Bouncer
		remoteIP string
	}
	tests := []struct {
		name    string
		args    args
		wantErr bool
	}{
		// TODO: Add test cases.
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if _, err := handleNoStreamCache(tt.args.bouncer, tt.args.remoteIP); (err != nil) != tt.wantErr {
				t.Errorf("handleNoStreamCache() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}

func Test_handleStreamCache(t *testing.T) {
	type args struct {
		bouncer *Bouncer
	}
	tests := []struct {
		name    string
		args    args
		wantErr bool
	}{
		// TODO: Add test cases.
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := handleStreamCache(tt.args.bouncer)
			if (err != nil) != tt.wantErr {
				t.Errorf("handleStreamCache() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
		})
	}
}

func Test_crowdsecQuery(t *testing.T) {
	type args struct {
		bouncer   *Bouncer
		stringURL string
		data      []byte
	}
	tests := []struct {
		name    string
		args    args
		want    []byte
		wantErr bool
	}{
		// TODO: Add test cases.
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := crowdsecQuery(tt.args.bouncer, tt.args.stringURL, tt.args.data)
			if (err != nil) != tt.wantErr {
				t.Errorf("crowdsecQuery() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("crowdsecQuery() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestHandleBanServeHTTPWithDifferentMethods(t *testing.T) {
	html := "<html>You are banned</html>"
	banTemplate, _ := template.New("html").Delims("{{", "}}").Parse(html)
	tests := []struct {
		name              string
		method            string
		banTemplate       *template.Template
		expectBodyContent bool
	}{
		{
			name:              "GET request should have body with template",
			method:            http.MethodGet,
			banTemplate:       banTemplate,
			expectBodyContent: true,
		},
		{
			name:              "HEAD request should NOT have body even with template",
			method:            http.MethodHead,
			banTemplate:       banTemplate,
			expectBodyContent: false,
		},
		{
			name:              "POST request should have body with template",
			method:            http.MethodPost,
			banTemplate:       banTemplate,
			expectBodyContent: true,
		},
		{
			name:              "PUT request should have body with template",
			method:            http.MethodPut,
			banTemplate:       banTemplate,
			expectBodyContent: true,
		},
		{
			name:              "DELETE request should have body with template",
			method:            http.MethodDelete,
			banTemplate:       banTemplate,
			expectBodyContent: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			bouncer := &Bouncer{
				remediationStatusCode:   http.StatusForbidden,
				remediationCustomHeader: "X-Test-Remediation",
				banTemplate:             tt.banTemplate,
				banTemplateContentType:  "text/html; charset=utf-8",
			}

			rw := httptest.NewRecorder()
			req := &http.Request{Method: tt.method}
			bouncer.handleBanServeHTTP(rw, req, "0.0.0.0", "TEST")

			// Check status code
			if rw.Code != http.StatusForbidden {
				t.Errorf("Expected status code 403, got %d", rw.Code)
			}

			// Check custom header
			headerValue := rw.Header().Get("X-Test-Remediation")
			if headerValue != "ban" {
				t.Errorf("Expected header X-Test-Remediation to be 'ban', got %s", headerValue)
			}

			// Check body content
			body := rw.Body.String()
			hasBodyContent := len(body) > 0

			if hasBodyContent != tt.expectBodyContent {
				t.Errorf("Method %s: expected body content: %v, got body content: %v (body: %q)",
					tt.method, tt.expectBodyContent, hasBodyContent, body)
			}

			// If we expect body content, verify it matches template
			if tt.expectBodyContent && body != html {
				t.Errorf("Expected body %q, got %q", html, body)
			}
		})
	}
}

func TestHandleBanServeHTTPContentType(t *testing.T) {
	html := "<html>You are banned</html>"
	banTemplate, _ := template.New("html").Delims("{{", "}}").Parse(html)
	tests := []struct {
		name                   string
		banTemplate            *template.Template
		banTemplateContentType string
	}{
		{
			name:                   "Default HTML content type",
			banTemplate:            banTemplate,
			banTemplateContentType: "text/html; charset=utf-8",
		},
		{
			name:                   "Custom JSON content type",
			banTemplate:            banTemplate,
			banTemplateContentType: "application/json",
		},
		{
			name:                   "Content type set even when banTemplate is nil",
			banTemplate:            nil,
			banTemplateContentType: "application/json",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			bouncer := &Bouncer{
				remediationStatusCode:  http.StatusForbidden,
				banTemplate:            tt.banTemplate,
				banTemplateContentType: tt.banTemplateContentType,
			}

			rw := httptest.NewRecorder()
			req := &http.Request{Method: http.MethodGet}
			bouncer.handleBanServeHTTP(rw, req, "0.0.0.0", "TEST")

			if got := rw.Header().Get("Content-Type"); got != tt.banTemplateContentType {
				t.Errorf("Expected Content-Type %q, got %q", tt.banTemplateContentType, got)
			}
		})
	}
}

func TestCaptchaMethodBasedLogic(t *testing.T) {
	tests := []struct {
		name              string
		method            string
		remediation       string
		expectBanFallback bool
	}{
		{
			name:              "GET with captcha remediation should allow captcha",
			method:            http.MethodGet,
			remediation:       cache.CaptchaValue,
			expectBanFallback: false,
		},
		{
			name:              "HEAD with captcha remediation should fallback to ban",
			method:            http.MethodHead,
			remediation:       cache.CaptchaValue,
			expectBanFallback: true,
		},
		{
			name:              "POST with captcha remediation should allow captcha",
			method:            http.MethodPost,
			remediation:       cache.CaptchaValue,
			expectBanFallback: false,
		},
		{
			name:              "PUT with captcha remediation should allow captcha",
			method:            http.MethodPut,
			remediation:       cache.CaptchaValue,
			expectBanFallback: false,
		},
		{
			name:              "DELETE with captcha remediation should allow captcha",
			method:            http.MethodDelete,
			remediation:       cache.CaptchaValue,
			expectBanFallback: false,
		},
		{
			name:              "PATCH with captcha remediation should allow captcha",
			method:            http.MethodPatch,
			remediation:       cache.CaptchaValue,
			expectBanFallback: false,
		},
		{
			name:              "OPTIONS with captcha remediation should allow captcha",
			method:            http.MethodOptions,
			remediation:       cache.CaptchaValue,
			expectBanFallback: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Test the core logic: captcha is served for all methods except HEAD
			shouldUseCaptcha := tt.remediation == cache.CaptchaValue && tt.method != http.MethodHead

			if shouldUseCaptcha == tt.expectBanFallback {
				t.Errorf("Method %s with %s remediation: expected ban fallback %v, but logic would use captcha %v",
					tt.method, tt.remediation, tt.expectBanFallback, shouldUseCaptcha)
			}
		})
	}
}

// blockingBody simulates a request body that never reaches EOF, like a
// bidirectional gRPC stream that keeps its body open for the whole life of
// the connection. Reading from it blocks until the test is done.
type blockingBody struct {
	done <-chan struct{}
}

func (b blockingBody) Read(_ []byte) (int, error) {
	<-b.done
	return 0, io.EOF
}

func (blockingBody) Close() error { return nil }

func Test_isBodyUnreadable(t *testing.T) {
	realBody := func() io.ReadCloser { return io.NopCloser(strings.NewReader("data")) }
	tests := []struct {
		name          string
		protoMajor    int
		contentLength int64
		body          io.ReadCloser
		want          bool
	}{
		{name: "http2 grpc stream without content-length", protoMajor: 2, contentLength: -1, body: realBody(), want: true},
		{name: "http3 stream without content-length", protoMajor: 3, contentLength: -1, body: realBody(), want: true},
		{name: "http2 with content-length", protoMajor: 2, contentLength: 42, body: realBody(), want: false},
		{name: "http1.1 chunked without content-length", protoMajor: 1, contentLength: -1, body: realBody(), want: false},
		{name: "http2 without body", protoMajor: 2, contentLength: -1, body: nil, want: false},
		{name: "http2 with http.NoBody", protoMajor: 2, contentLength: -1, body: http.NoBody, want: false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req, _ := http.NewRequest(http.MethodPost, "http://localhost", nil)
			req.ProtoMajor = tt.protoMajor
			req.ContentLength = tt.contentLength
			req.Body = tt.body
			if got := isBodyUnreadable(req); got != tt.want {
				t.Errorf("isBodyUnreadable() = %v, want %v", got, tt.want)
			}
		})
	}
}

// newStreamingRequest builds an HTTP/2 request whose body never reaches EOF,
// like a bidirectional gRPC stream (issue #323).
func newStreamingRequest(done <-chan struct{}) *http.Request {
	req, _ := http.NewRequest(http.MethodPost, "http://localhost/signalexchange.SignalExchange/ConnectStream", blockingBody{done: done})
	req.Header.Set("Content-Type", "application/grpc")
	req.ProtoMajor = 2
	req.ContentLength = -1
	return req
}

// Test_appsecQuery_streamingDoesNotBlock is a regression test for issue #323:
// a gRPC streaming request whose body never reaches EOF must not be buffered
// (io.ReadAll would block until timeout and wrongly produce a 403). The appsec
// query must complete promptly, inspecting headers only.
func Test_appsecQuery_streamingDoesNotBlock(t *testing.T) {
	appsecServer := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, _ *http.Request) {
		rw.WriteHeader(http.StatusOK)
	}))
	defer appsecServer.Close()

	appsecURL, _ := url.Parse(appsecServer.URL)
	bouncer := &Bouncer{
		appsecScheme:           appsecURL.Scheme,
		appsecHost:             appsecURL.Host,
		appsecPath:             "/",
		appsecBodyLimit:        10485760,
		appsecUnreachableBlock: true,
		appsecFailureBlock:     true,
		httpAppsecClient:       appsecServer.Client(),
		log:                    logger.New("INFO", ""),
	}

	done := make(chan struct{})
	defer close(done)

	finished := make(chan error, 1)
	go func() {
		finished <- appsecQuery(bouncer, "1.2.3.4", newStreamingRequest(done))
	}()

	select {
	case err := <-finished:
		if err != nil {
			t.Errorf("appsecQuery() on streaming request returned error: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("appsecQuery() blocked on a streaming request body (issue #323 regression)")
	}
}

// Test_appsecQuery_dropUnreadableBody verifies that, when configured to do so,
// a request with an unreadable body is dropped (blocked) instead of forwarded
// without its body, mirroring the reference APPSEC_DROP_UNREADABLE_BODY option.
func Test_appsecQuery_dropUnreadableBody(t *testing.T) {
	appsecServer := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, _ *http.Request) {
		rw.WriteHeader(http.StatusOK)
	}))
	defer appsecServer.Close()

	appsecURL, _ := url.Parse(appsecServer.URL)
	bouncer := &Bouncer{
		appsecScheme:              appsecURL.Scheme,
		appsecHost:                appsecURL.Host,
		appsecPath:                "/",
		appsecBodyLimit:           10485760,
		appsecUnreadableBodyBlock: true,
		httpAppsecClient:          appsecServer.Client(),
		log:                       logger.New("INFO", ""),
	}

	done := make(chan struct{})
	defer close(done)

	finished := make(chan error, 1)
	go func() {
		finished <- appsecQuery(bouncer, "1.2.3.4", newStreamingRequest(done))
	}()

	select {
	case err := <-finished:
		if err == nil {
			t.Error("appsecQuery() expected an error to block the request, got nil")
		}
	case <-time.After(2 * time.Second):
		t.Fatal("appsecQuery() blocked on a streaming request body (issue #323 regression)")
	}
}

func newUnreadableGetRequest(done <-chan struct{}) *http.Request {
	req, _ := http.NewRequest(http.MethodGet, "http://localhost/", blockingBody{done: done})
	req.ProtoMajor = 3
	req.ContentLength = -1
	return req
}

// Test_appsecQuery_unreadableBodyGetNotDropped is a regression test for issue #351
func Test_appsecQuery_unreadableBodyGetNotDropped(t *testing.T) {
	appsecServer := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, _ *http.Request) {
		rw.WriteHeader(http.StatusOK)
	}))
	defer appsecServer.Close()

	appsecURL, _ := url.Parse(appsecServer.URL)
	bouncer := &Bouncer{
		appsecScheme:              appsecURL.Scheme,
		appsecHost:                appsecURL.Host,
		appsecPath:                "/",
		appsecBodyLimit:           10485760,
		appsecUnreadableBodyBlock: true,
		httpAppsecClient:          appsecServer.Client(),
		log:                       logger.New("INFO", ""),
	}

	done := make(chan struct{})
	defer close(done)

	finished := make(chan error, 1)
	go func() {
		finished <- appsecQuery(bouncer, "1.2.3.4", newUnreadableGetRequest(done))
	}()

	select {
	case err := <-finished:
		if err != nil {
			t.Errorf("appsecQuery() on an HTTP/3 GET without content-length returned error: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("appsecQuery() blocked on an HTTP/3 GET request body (issue #351 regression)")
	}
}

// Test_appsecQuery_reusesConnection is a regression test for issue #384: the appsec response
// body must be drained for non-200 status response so net/http can return the conn to the idle pool.
func Test_appsecQuery_reusesConnection(t *testing.T) {
	for _, status := range []int{http.StatusOK, http.StatusForbidden, http.StatusInternalServerError} {
		t.Run(http.StatusText(status), func(t *testing.T) {
			var mu sync.Mutex
			conns := map[string]bool{}
			appsecServer := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
				mu.Lock()
				conns[r.RemoteAddr] = true
				mu.Unlock()
				rw.WriteHeader(status)
				if _, errWrite := rw.Write([]byte(`{"action":"allow"}`)); errWrite != nil {
					t.Errorf("appsec stub write: %v", errWrite)
				}
			}))
			defer appsecServer.Close()

			appsecURL, _ := url.Parse(appsecServer.URL)
			bouncer := &Bouncer{
				appsecScheme:       appsecURL.Scheme,
				appsecHost:         appsecURL.Host,
				appsecPath:         "/",
				appsecBodyLimit:    10485760,
				appsecFailureBlock: false,
				httpAppsecClient:   appsecServer.Client(),
				log:                logger.New("INFO", ""),
			}

			const calls = 10
			for i := 0; i < calls; i++ { //nolint:intrange
				req, _ := http.NewRequest(http.MethodGet, "http://localhost/", nil)
				_ = appsecQuery(bouncer, "1.2.3.4", req)
			}

			mu.Lock()
			defer mu.Unlock()
			if len(conns) != 1 {
				t.Errorf("appsecQuery() opened %d connections for %d calls, want 1 (response body not drained?)", len(conns), calls)
			}
		})
	}
}

func newUnreadableRequest(method string, done <-chan struct{}) *http.Request {
	req, _ := http.NewRequest(method, "http://localhost/api/admin/reservations/8fff14a2", blockingBody{done: done})
	req.ProtoMajor = 3
	req.ContentLength = -1
	return req
}

func Test_appsecQuery_unreadableBodyMethods(t *testing.T) {
	appsecServer := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, _ *http.Request) {
		rw.WriteHeader(http.StatusOK)
	}))
	defer appsecServer.Close()

	appsecURL, _ := url.Parse(appsecServer.URL)

	tests := []struct {
		method      string
		wantDropped bool
	}{
		{method: http.MethodGet, wantDropped: false},
		{method: http.MethodHead, wantDropped: false},
		{method: http.MethodOptions, wantDropped: false},
		{method: http.MethodDelete, wantDropped: false},
		{method: http.MethodPost, wantDropped: true},
		{method: http.MethodPut, wantDropped: true},
		{method: http.MethodPatch, wantDropped: true},
	}

	for _, tt := range tests {
		t.Run(tt.method, func(t *testing.T) {
			bouncer := &Bouncer{
				appsecScheme:              appsecURL.Scheme,
				appsecHost:                appsecURL.Host,
				appsecPath:                "/",
				appsecBodyLimit:           10485760,
				appsecUnreadableBodyBlock: true,
				httpAppsecClient:          appsecServer.Client(),
				log:                       logger.New("INFO", ""),
			}

			done := make(chan struct{})
			defer close(done)

			finished := make(chan error, 1)
			go func() {
				finished <- appsecQuery(bouncer, "1.2.3.4", newUnreadableRequest(tt.method, done))
			}()

			select {
			case err := <-finished:
				if tt.wantDropped && err == nil {
					t.Errorf("appsecQuery() on an unreadable-body %s: expected the request to be dropped, got nil", tt.method)
				}
				if !tt.wantDropped && err != nil {
					t.Errorf("appsecQuery() on a bodyless HTTP/3 %s returned error: %v", tt.method, err)
				}
			case <-time.After(2 * time.Second):
				t.Fatalf("appsecQuery() blocked on an unreadable %s body", tt.method)
			}
		})
	}
}