Skip to content

Commit

Permalink
Set timeout for requests made by the EVP proxy (#26336)
Browse files Browse the repository at this point in the history
  • Loading branch information
albertvaka authored Jun 10, 2024
1 parent 1c424e7 commit 5194aa8
Show file tree
Hide file tree
Showing 4 changed files with 104 additions and 20 deletions.
13 changes: 9 additions & 4 deletions pkg/trace/api/api.go
Original file line number Diff line number Diff line change
Expand Up @@ -195,6 +195,14 @@ func replyWithVersion(hash string, version string, h http.Handler) http.Handler
})
}

func getConfiguredRequestTimeoutDuration(conf *config.AgentConfig) time.Duration {
timeout := 5 * time.Second
if conf.ReceiverTimeout > 0 {
timeout = time.Duration(conf.ReceiverTimeout) * time.Second
}
return timeout
}

// Start starts doing the HTTP server and is ready to receive traces
func (r *HTTPReceiver) Start() {
if r.conf.ReceiverPort == 0 &&
Expand All @@ -205,10 +213,7 @@ func (r *HTTPReceiver) Start() {
return
}

timeout := 5 * time.Second
if r.conf.ReceiverTimeout > 0 {
timeout = time.Duration(r.conf.ReceiverTimeout) * time.Second
}
timeout := getConfiguredRequestTimeoutDuration(r.conf)
httpLogger := log.NewThrottled(5, 10*time.Second) // limit to 5 messages every 10 seconds
r.server = &http.Server{
ReadTimeout: timeout,
Expand Down
8 changes: 8 additions & 0 deletions pkg/trace/api/evp_proxy.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ package api

import (
"bytes"
"context"
"fmt"
"io"
stdlog "log"
Expand Down Expand Up @@ -174,6 +175,13 @@ func (t *evpProxyTransport) RoundTrip(req *http.Request) (rresp *http.Response,
req.Header.Set("DD-APPLICATION-KEY", t.conf.EVPProxy.ApplicationKey)
}

// Timeout: Our outbound request(s) can't take longer than the WriteTimeout of the server
timeout := getConfiguredRequestTimeoutDuration(t.conf)
deadline := time.Now().Add(timeout)
ctx, ctxCancel := context.WithDeadline(req.Context(), deadline)
req = req.WithContext(ctx)
defer ctxCancel()

// Set target URL and API key header (per domain)
req.URL.Scheme = "https"
setTarget := func(r *http.Request, host, apiKey string) {
Expand Down
101 changes: 86 additions & 15 deletions pkg/trace/api/evp_proxy_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,9 @@ import (
"net/http"
"net/http/httptest"
"net/http/httputil"
"strings"
"testing"
"time"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
Expand All @@ -31,10 +33,11 @@ func (r roundTripperMock) RoundTrip(req *http.Request) (*http.Response, error) {
return r(req)
}

// sendRequestThroughForwarder sends a request through the evpProxyForwarder handler and returns the forwarded
// request(s), their response and the log output. The path for inReq shouldn't have the /evp_proxy/v1
// prefix since it is passed directly to the inner proxy handler and not the trace-agent API handler.
func sendRequestThroughForwarder(conf *config.AgentConfig, inReq *http.Request, statsd statsd.ClientInterface) (outReqs []*http.Request, resp *http.Response, logs string) {
// sendRequestThroughForwarderWithMockRoundTripper sends a request through the evpProxyForwarder
// handler and returns the forwarded request(s), their response and the log output. The path for
// inReq shouldn't have the /evp_proxy/v1 prefix since it is passed directly to the inner proxy
// handler and not the trace-agent API handler.
func sendRequestThroughForwarderWithMockRoundTripper(conf *config.AgentConfig, inReq *http.Request, statsd statsd.ClientInterface) (outReqs []*http.Request, resp *http.Response, logs string) {
mockRoundTripper := roundTripperMock(func(req *http.Request) (*http.Response, error) {
if req.Body != nil {
if _, err := io.ReadAll(req.Body); err != nil && err != io.EOF {
Expand All @@ -57,6 +60,26 @@ func sendRequestThroughForwarder(conf *config.AgentConfig, inReq *http.Request,
return outReqs, rec.Result(), loggerBuffer.String()
}

// sendRequestThroughForwarderAgainstDummyServer sends a request to the specified serverHost URL,
// which should be localhost at a port where you ran a dummy server, for E2E testing. The path for
// inReq shouldn't have the /evp_proxy/v1 prefix since it is passed directly to the inner proxy
// handler and not the trace-agent API handler.
func sendRequestThroughForwarderAgainstDummyServer(conf *config.AgentConfig, inReq *http.Request, statsd statsd.ClientInterface, serverHost string) (resp *http.Response, logs string) {
reqModifierRoundTripper := roundTripperMock(func(req *http.Request) (*http.Response, error) {
req.URL.Scheme = "http"
req.Host = serverHost
req.URL.Host = serverHost
return conf.NewHTTPTransport().RoundTrip(req)
})
handler := evpProxyForwarder(conf, statsd)
var loggerBuffer bytes.Buffer
handler.(*httputil.ReverseProxy).ErrorLog = log.New(io.Writer(&loggerBuffer), "", 0)
handler.(*httputil.ReverseProxy).Transport.(*evpProxyTransport).transport = reqModifierRoundTripper
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, inReq)
return rec.Result(), loggerBuffer.String()
}

func TestEVPProxyForwarder(t *testing.T) {
randBodyBuf := make([]byte, 1024)
rand.Read(randBodyBuf)
Expand All @@ -77,7 +100,7 @@ func TestEVPProxyForwarder(t *testing.T) {
req.Header.Set("User-Agent", "test_user_agent")
req.Header.Set("X-Datadog-EVP-Subdomain", "my.subdomain")
req.Header.Set("Content-Type", "text/json")
proxyreqs, resp, logs := sendRequestThroughForwarder(conf, req, stats)
proxyreqs, resp, logs := sendRequestThroughForwarderWithMockRoundTripper(conf, req, stats)

require.Equal(t, http.StatusOK, resp.StatusCode, "Got: ", fmt.Sprint(resp.StatusCode))
resp.Body.Close()
Expand Down Expand Up @@ -128,7 +151,7 @@ func TestEVPProxyForwarder(t *testing.T) {
req := httptest.NewRequest("POST", "/mypath/mysubpath?arg=test", bytes.NewReader(randBodyBuf))
req.Header.Set("X-Datadog-EVP-Subdomain", "my.subdomain")
req.Header.Set(header.ContainerID, "myid")
proxyreqs, resp, logs := sendRequestThroughForwarder(conf, req, stats)
proxyreqs, resp, logs := sendRequestThroughForwarderWithMockRoundTripper(conf, req, stats)

resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode, "Got: ", fmt.Sprint(resp.StatusCode))
Expand All @@ -147,7 +170,7 @@ func TestEVPProxyForwarder(t *testing.T) {
}
req := httptest.NewRequest("POST", "/mypath/mysubpath?arg=test", bytes.NewReader(randBodyBuf))
req.Header.Set("X-Datadog-EVP-Subdomain", "my.subdomain")
proxyreqs, resp, logs := sendRequestThroughForwarder(conf, req, stats)
proxyreqs, resp, logs := sendRequestThroughForwarderWithMockRoundTripper(conf, req, stats)

resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode, "Got: ", fmt.Sprint(resp.StatusCode))
Expand Down Expand Up @@ -180,7 +203,7 @@ func TestEVPProxyForwarder(t *testing.T) {
conf.Endpoints[0].APIKey = "test_api_key"

req := httptest.NewRequest("POST", "/mypath/mysubpath", bytes.NewReader(randBodyBuf))
proxyreqs, resp, logs := sendRequestThroughForwarder(conf, req, stats)
proxyreqs, resp, logs := sendRequestThroughForwarderWithMockRoundTripper(conf, req, stats)

resp.Body.Close()
require.Len(t, proxyreqs, 0)
Expand All @@ -204,7 +227,7 @@ func TestEVPProxyForwarder(t *testing.T) {

req := httptest.NewRequest("POST", "/mypath/mysubpath", bytes.NewReader(randBodyBuf))
req.Header.Set("X-Datadog-EVP-Subdomain", "/google.com%3Fattack=")
proxyreqs, resp, logs := sendRequestThroughForwarder(conf, req, stats)
proxyreqs, resp, logs := sendRequestThroughForwarderWithMockRoundTripper(conf, req, stats)

resp.Body.Close()
require.Len(t, proxyreqs, 0)
Expand All @@ -228,7 +251,7 @@ func TestEVPProxyForwarder(t *testing.T) {

req := httptest.NewRequest("POST", "/mypath/my%20subpath", bytes.NewReader(randBodyBuf))
req.Header.Set("X-Datadog-EVP-Subdomain", "my.subdomain")
proxyreqs, resp, logs := sendRequestThroughForwarder(conf, req, stats)
proxyreqs, resp, logs := sendRequestThroughForwarderWithMockRoundTripper(conf, req, stats)

resp.Body.Close()
require.Len(t, proxyreqs, 0)
Expand All @@ -255,7 +278,7 @@ func TestEVPProxyForwarder(t *testing.T) {

req := httptest.NewRequest("POST", "/mypath/mysubpath?test=bad%20arg", bytes.NewReader(randBodyBuf))
req.Header.Set("X-Datadog-EVP-Subdomain", "my.subdomain")
proxyreqs, resp, logs := sendRequestThroughForwarder(conf, req, stats)
proxyreqs, resp, logs := sendRequestThroughForwarderWithMockRoundTripper(conf, req, stats)

resp.Body.Close()
require.Len(t, proxyreqs, 0)
Expand Down Expand Up @@ -283,7 +306,7 @@ func TestEVPProxyForwarder(t *testing.T) {

req := httptest.NewRequest("POST", "/mypath/mysubpath", bytes.NewReader(randBodyBuf))
req.Header.Set("X-Datadog-EVP-Subdomain", "my.subdomain")
proxyreqs, resp, logs := sendRequestThroughForwarder(conf, req, stats)
proxyreqs, resp, logs := sendRequestThroughForwarderWithMockRoundTripper(conf, req, stats)

resp.Body.Close()
require.Len(t, proxyreqs, 0)
Expand Down Expand Up @@ -324,7 +347,7 @@ func TestEVPProxyForwarder(t *testing.T) {
req := httptest.NewRequest("POST", "/mypath/mysubpath", bytes.NewReader(randBodyBuf))
req.Header.Set("X-Datadog-EVP-Subdomain", "my.subdomain")
req.Header.Set("X-Datadog-NeedsAppKey", "true")
proxyreqs, resp, logs := sendRequestThroughForwarder(conf, req, stats)
proxyreqs, resp, logs := sendRequestThroughForwarderWithMockRoundTripper(conf, req, stats)

resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode, "Got: ", fmt.Sprint(resp.StatusCode))
Expand All @@ -343,7 +366,7 @@ func TestEVPProxyForwarder(t *testing.T) {
req := httptest.NewRequest("POST", "/mypath/mysubpath", bytes.NewReader(randBodyBuf))
req.Header.Set("X-Datadog-EVP-Subdomain", "my.subdomain")
req.Header.Set("X-Datadog-NeedsAppKey", "true")
proxyreqs, resp, logs := sendRequestThroughForwarder(conf, req, stats)
proxyreqs, resp, logs := sendRequestThroughForwarderWithMockRoundTripper(conf, req, stats)

resp.Body.Close()
require.Len(t, proxyreqs, 0)
Expand Down Expand Up @@ -375,7 +398,7 @@ func TestEVPProxyForwarder(t *testing.T) {
req.Header.Set("Content-Encoding", "gzip")
req.Header.Set("Unexpected-Header", "To-Be-Discarded")
req.Header.Set("DD-CI-PROVIDER-NAME", "Allowed-Header")
proxyreqs, resp, logs := sendRequestThroughForwarder(conf, req, stats)
proxyreqs, resp, logs := sendRequestThroughForwarderWithMockRoundTripper(conf, req, stats)

require.Equal(t, http.StatusOK, resp.StatusCode, "Got: ", fmt.Sprint(resp.StatusCode))
resp.Body.Close()
Expand Down Expand Up @@ -414,3 +437,51 @@ func TestEVPProxyHandler(t *testing.T) {
require.Equal(t, http.StatusMethodNotAllowed, rec.Code)
})
}

func TestE2E(t *testing.T) {
randBodyBuf := make([]byte, 1024)
rand.Read(randBodyBuf)

stats := &teststatsd.Client{}

t.Run("ok", func(t *testing.T) {
conf := newTestReceiverConfig()
conf.Site = "us3.datadoghq.com"
conf.Endpoints[0].APIKey = "test_api_key"

server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Write([]byte(`OK`))
}))

req := httptest.NewRequest("POST", "/mypath/mysubpath?arg=test", bytes.NewReader(randBodyBuf))
req.Header.Set("X-Datadog-EVP-Subdomain", "my.subdomain")
resp, logs := sendRequestThroughForwarderAgainstDummyServer(conf, req, stats, strings.TrimPrefix(server.URL, "http://"))

resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode, "Got: ", fmt.Sprint(resp.StatusCode))
assert.Equal(t, "", logs)
body, err := io.ReadAll(resp.Body)
assert.NoError(t, err)
assert.Equal(t, "OK", string(body))
})

t.Run("timeout", func(t *testing.T) {
conf := newTestReceiverConfig()
conf.Site = "us3.datadoghq.com"
conf.Endpoints[0].APIKey = "test_api_key"
conf.ReceiverTimeout = 1 // in seconds

server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
time.Sleep(2 * time.Second)
w.Write([]byte(`OK`))
}))

req := httptest.NewRequest("POST", "/mypath/mysubpath?arg=test", bytes.NewReader(randBodyBuf))
req.Header.Set("X-Datadog-EVP-Subdomain", "my.subdomain")
resp, logs := sendRequestThroughForwarderAgainstDummyServer(conf, req, stats, strings.TrimPrefix(server.URL, "http://"))

resp.Body.Close()
require.Equal(t, http.StatusBadGateway, resp.StatusCode, "Got: ", fmt.Sprint(resp.StatusCode))
assert.Equal(t, "http: proxy error: context deadline exceeded\n", logs)
})
}
2 changes: 1 addition & 1 deletion tasks/go_deps.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ def test_list(
filename = os.path.join(ctx.cwd, f"dependencies_{goos}_{goarch}.txt")
if not os.path.isfile(filename):
print(
f"File {filename} does not exist. To execute the dependencies list check for the {binary} binary, please run the task `inv -e go-deps.generate --binaries {binary}"
f"File {filename} does not exist. To execute the dependencies list check for the {binary} binary, please run the task `inv -e go-deps.generate"
)
continue

Expand Down

0 comments on commit 5194aa8

Please sign in to comment.