traefik/pkg/proxy/smart_builder_test.go
2024-10-17 09:12:04 +02:00

113 lines
3.1 KiB
Go

package proxy
import (
"encoding/pem"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/traefik/traefik/v3/pkg/config/dynamic"
"github.com/traefik/traefik/v3/pkg/config/static"
"github.com/traefik/traefik/v3/pkg/proxy/httputil"
"github.com/traefik/traefik/v3/pkg/server/service"
"github.com/traefik/traefik/v3/pkg/testhelpers"
"github.com/traefik/traefik/v3/pkg/types"
"golang.org/x/net/http2"
"golang.org/x/net/http2/h2c"
)
func TestSmartBuilder_Build(t *testing.T) {
tests := []struct {
desc string
serversTransport dynamic.ServersTransport
fastProxyConfig static.FastProxyConfig
https bool
h2c bool
wantFastProxy bool
}{
{
desc: "fastproxy",
fastProxyConfig: static.FastProxyConfig{Debug: true},
wantFastProxy: true,
},
{
desc: "fastproxy with https and without DisableHTTP2",
https: true,
fastProxyConfig: static.FastProxyConfig{Debug: true},
wantFastProxy: false,
},
{
desc: "fastproxy with https and DisableHTTP2",
https: true,
serversTransport: dynamic.ServersTransport{DisableHTTP2: true},
fastProxyConfig: static.FastProxyConfig{Debug: true},
wantFastProxy: true,
},
{
desc: "fastproxy with h2c",
h2c: true,
fastProxyConfig: static.FastProxyConfig{Debug: true},
wantFastProxy: false,
},
}
for _, test := range tests {
t.Run(test.desc, func(t *testing.T) {
t.Parallel()
var callCount int
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
callCount++
if test.wantFastProxy {
assert.Contains(t, r.Header, "X-Traefik-Fast-Proxy")
} else {
assert.NotContains(t, r.Header, "X-Traefik-Fast-Proxy")
}
})
var server *httptest.Server
if test.https {
server = httptest.NewUnstartedServer(handler)
server.EnableHTTP2 = false
server.StartTLS()
certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: server.TLS.Certificates[0].Certificate[0]})
test.serversTransport.RootCAs = []types.FileOrContent{
types.FileOrContent(certPEM),
}
} else {
server = httptest.NewServer(h2c.NewHandler(handler, &http2.Server{}))
}
t.Cleanup(func() {
server.Close()
})
targetURL := testhelpers.MustParseURL(server.URL)
if test.h2c {
targetURL.Scheme = "h2c"
}
serversTransports := map[string]*dynamic.ServersTransport{
"test": &test.serversTransport,
}
transportManager := service.NewTransportManager(nil)
transportManager.Update(serversTransports)
httpProxyBuilder := httputil.NewProxyBuilder(transportManager, nil)
proxyBuilder := NewSmartBuilder(transportManager, httpProxyBuilder, test.fastProxyConfig)
proxyHandler, err := proxyBuilder.Build("test", targetURL, false, false, false, time.Second)
require.NoError(t, err)
rw := httptest.NewRecorder()
proxyHandler.ServeHTTP(rw, httptest.NewRequest(http.MethodGet, "/", http.NoBody))
assert.Equal(t, 1, callCount)
})
}
}