From 834920d6f97eabf8cfced8cec61b9b82b33dbe62 Mon Sep 17 00:00:00 2001
From: kruskal <99559985+kruskall@users.noreply.github.com>
Date: Mon, 14 Nov 2022 21:30:31 +0100
Subject: [PATCH 1/3] fix: do not return status 200 on server info request when
 unreachable

When proxying info requests to apm server we're loggin an error if
the server is unreachable.
The ErrorHandler should also write a status code back to avoid
returning a Status 200.
Add ModifyResponse to log a debug message with the body on unexpected
responses and let the ErrorHandler proxy the status code back to the
agent.
A StatusBadGateway is returned for generic errors.
---
 apmproxy/receiver.go      | 36 ++++++++++++++++++++++++++++++++++++
 apmproxy/receiver_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+)

diff --git a/apmproxy/receiver.go b/apmproxy/receiver.go
index ed321b06..20858fb8 100644
--- a/apmproxy/receiver.go
+++ b/apmproxy/receiver.go
@@ -73,6 +73,12 @@ func (c *Client) Shutdown() error {
 	return c.receiver.Shutdown(ctx)
 }
 
+type unexpectedStatusError int
+
+func (e unexpectedStatusError) Error() string {
+	return fmt.Sprintf("unexpected status code: %d", int(e))
+}
+
 // URL: http://server/
 func (c *Client) handleInfoRequest() (func(w http.ResponseWriter, r *http.Request), error) {
 	// Init reverse proxy
@@ -87,9 +93,39 @@ func (c *Client) handleInfoRequest() (func(w http.ResponseWriter, r *http.Reques
 	customTransport.ResponseHeaderTimeout = c.client.Timeout
 	reverseProxy.Transport = customTransport
 
+	reverseProxy.ModifyResponse = func(rsp *http.Response) error {
+		if rsp.StatusCode != http.StatusOK {
+			b, err := io.ReadAll(rsp.Body)
+			if err != nil {
+				c.logger.Warnf("failed to read version response body: %v", err)
+				b = []byte{}
+			}
+
+			c.logger.Debugf("failed to query version from the APM server: response body: %s", string(b))
+
+			// Return an error to the ErrorHandler that will proxy the status code back to the agent.
+			return fmt.Errorf("failed to query version from the APM server: response status: %w", unexpectedStatusError(rsp.StatusCode))
+		}
+
+		return nil
+	}
+
 	reverseProxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
 		c.UpdateStatus(r.Context(), Failing)
 		c.logger.Errorf("Error querying version from the APM server: %v", err)
+
+		uErr := new(unexpectedStatusError)
+		if errors.As(err, uErr) {
+			status := int(*uErr)
+
+			// Proxy the status code to the agent.
+			w.WriteHeader(status)
+			return
+		}
+
+		// Server is unreachable, return StatusBadGateway (default behaviour) to avoid
+		// returning a Status OK.
+		w.WriteHeader(http.StatusBadGateway)
 	}
 
 	return func(w http.ResponseWriter, r *http.Request) {
diff --git a/apmproxy/receiver_test.go b/apmproxy/receiver_test.go
index 155ae79a..bbbbcc2d 100644
--- a/apmproxy/receiver_test.go
+++ b/apmproxy/receiver_test.go
@@ -125,6 +125,44 @@ func TestInfoProxyErrorStatusCode(t *testing.T) {
 	assert.Equal(t, http.StatusUnauthorized, resp.StatusCode)
 }
 
+func TestInfoProxyUnreachable(t *testing.T) {
+	// Create apm server and handler
+	apmServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	// Shutdown
+	apmServer.Close()
+
+	// Create extension config and start the server
+	apmClient, err := apmproxy.NewClient(
+		apmproxy.WithURL(apmServer.URL),
+		apmproxy.WithSecretToken("foo"),
+		apmproxy.WithAPIKey("bar"),
+		apmproxy.WithReceiverAddress(":1234"),
+		apmproxy.WithReceiverTimeout(15*time.Second),
+		apmproxy.WithLogger(zaptest.NewLogger(t).Sugar()),
+	)
+	require.NoError(t, err)
+
+	require.NoError(t, apmClient.StartReceiver())
+	defer func() {
+		require.NoError(t, apmClient.Shutdown())
+	}()
+
+	hosts, _ := net.LookupHost("localhost")
+	url := "http://" + hosts[0] + ":1234"
+
+	// Create a request to send to the extension
+	req, err := http.NewRequest(http.MethodGet, url, nil)
+	require.NoError(t, err)
+
+	// Send the request to the extension
+	client := &http.Client{}
+	resp, err := client.Do(req)
+	require.NoError(t, err)
+
+	// Make sure we don't get a 200 OK
+	assert.Equal(t, http.StatusBadGateway, resp.StatusCode)
+}
+
 func Test_handleInfoRequest(t *testing.T) {
 	headers := map[string]string{"Authorization": "test-value"}
 	// Copied from https://github.com/elastic/apm-server/blob/master/testdata/intake-v2/transactions.ndjson.

From c58ae5f6f8964bd08164fd0e5ccfb4fd811dcde7 Mon Sep 17 00:00:00 2001
From: kruskal <99559985+kruskall@users.noreply.github.com>
Date: Sun, 4 Dec 2022 14:05:07 +0100
Subject: [PATCH 2/3] fix: don't modify an error response from the server

Proxy the response as is without changing the status
code.
---
 apmproxy/receiver.go | 17 -----------------
 1 file changed, 17 deletions(-)

diff --git a/apmproxy/receiver.go b/apmproxy/receiver.go
index 20858fb8..ee294350 100644
--- a/apmproxy/receiver.go
+++ b/apmproxy/receiver.go
@@ -93,23 +93,6 @@ func (c *Client) handleInfoRequest() (func(w http.ResponseWriter, r *http.Reques
 	customTransport.ResponseHeaderTimeout = c.client.Timeout
 	reverseProxy.Transport = customTransport
 
-	reverseProxy.ModifyResponse = func(rsp *http.Response) error {
-		if rsp.StatusCode != http.StatusOK {
-			b, err := io.ReadAll(rsp.Body)
-			if err != nil {
-				c.logger.Warnf("failed to read version response body: %v", err)
-				b = []byte{}
-			}
-
-			c.logger.Debugf("failed to query version from the APM server: response body: %s", string(b))
-
-			// Return an error to the ErrorHandler that will proxy the status code back to the agent.
-			return fmt.Errorf("failed to query version from the APM server: response status: %w", unexpectedStatusError(rsp.StatusCode))
-		}
-
-		return nil
-	}
-
 	reverseProxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
 		c.UpdateStatus(r.Context(), Failing)
 		c.logger.Errorf("Error querying version from the APM server: %v", err)

From fa376cd2d725fd5416d4c7b087311692fc7bf06e Mon Sep 17 00:00:00 2001
From: kruskall <99559985+kruskall@users.noreply.github.com>
Date: Mon, 5 Dec 2022 02:08:56 +0100
Subject: [PATCH 3/3] refactor: remove unused code

Co-authored-by: Andrew Wilkins <axwalk@gmail.com>
---
 apmproxy/receiver.go | 14 --------------
 1 file changed, 14 deletions(-)

diff --git a/apmproxy/receiver.go b/apmproxy/receiver.go
index 316ff61a..f17e1580 100644
--- a/apmproxy/receiver.go
+++ b/apmproxy/receiver.go
@@ -73,12 +73,6 @@ func (c *Client) Shutdown() error {
 	return c.receiver.Shutdown(ctx)
 }
 
-type unexpectedStatusError int
-
-func (e unexpectedStatusError) Error() string {
-	return fmt.Sprintf("unexpected status code: %d", int(e))
-}
-
 // URL: http://server/
 func (c *Client) handleInfoRequest() (func(w http.ResponseWriter, r *http.Request), error) {
 	// Init reverse proxy
@@ -98,14 +92,6 @@ func (c *Client) handleInfoRequest() (func(w http.ResponseWriter, r *http.Reques
 		// is frozen while processing the request and context is canceled due to timeout.
 		c.logger.Errorf("Error querying version from the APM server: %v", err)
 
-		uErr := new(unexpectedStatusError)
-		if errors.As(err, uErr) {
-			status := int(*uErr)
-
-			// Proxy the status code to the agent.
-			w.WriteHeader(status)
-			return
-		}
 
 		// Server is unreachable, return StatusBadGateway (default behaviour) to avoid
 		// returning a Status OK.