Use copy for info response and add intake test (#83)

* Use copy to forward apm server response body to agent

* Remove unncessary check if body is nil and add intake test

* Close body after reading

* Write headers before body to response

Author: estolfo

apm-lambda-extension/extension/http_server_test.go

@@ -22,6 +22,7 @@ import (
 	"net"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"gotest.tools/assert"
@@ -37,6 +38,7 @@ func TestInfoProxy(t *testing.T) {
 			assert.Equal(t, 1, len(r.Header[key]))
 			assert.Equal(t, headers[key], r.Header[key][0])
 		}
+		w.Header().Add("test", "header")
 		w.Write([]byte(`{"foo": "bar"}`))
 	}))
 	defer apmServer.Close()
@@ -75,6 +77,91 @@ func TestInfoProxy(t *testing.T) {
 	} else {
 		body, _ := ioutil.ReadAll(resp.Body)
 		assert.Equal(t, string(body), wantResp)
+		assert.Equal(t, "header", resp.Header.Get("test"))
 		resp.Body.Close()
 	}
 }
+
+func TestInfoProxyErrorStatusCode(t *testing.T) {
+	// Create apm server and handler
+	apmServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(401)
+	}))
+	defer apmServer.Close()
+
+	// Create extension config and start the server
+	dataChannel := make(chan AgentData, 100)
+	config := extensionConfig{
+		apmServerUrl:               apmServer.URL,
+		apmServerSecretToken:       "foo",
+		apmServerApiKey:            "bar",
+		dataReceiverServerPort:     ":1234",
+		dataReceiverTimeoutSeconds: 15,
+	}
+
+	StartHttpServer(dataChannel, &config)
+	defer agentDataServer.Close()
+
+	hosts, _ := net.LookupHost("localhost")
+	url := "http://" + hosts[0] + ":1234"
+
+	// Create a request to send to the extension
+	req, err := http.NewRequest("GET", url, nil)
+	if err != nil {
+		t.Logf("Could not create request")
+	}
+
+	// Send the request to the extension
+	client := &http.Client{}
+	resp, err := client.Do(req)
+	if err != nil {
+		t.Logf("Error fetching %s, [%v]", agentDataServer.Addr, err)
+		t.Fail()
+	} else {
+		assert.Equal(t, 401, 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.
+	agentRequestBody := `{"metadata": {"service": {"name": "1234_service-12a3","node": {"configured_name": "node-123"},"version": "5.1.3","environment": "staging","language": {"name": "ecmascript","version": "8"},"runtime": {"name": "node","version": "8.0.0"},"framework": {"name": "Express","version": "1.2.3"},"agent": {"name": "elastic-node","version": "3.14.0"}},"user": {"id": "123user", "username": "bar", "email": "bar@user.com"}, "labels": {"tag0": null, "tag1": "one", "tag2": 2}, "process": {"pid": 1234,"ppid": 6789,"title": "node","argv": ["node","server.js"]},"system": {"hostname": "prod1.example.com","architecture": "x64","platform": "darwin", "container": {"id": "container-id"}, "kubernetes": {"namespace": "namespace1", "pod": {"uid": "pod-uid", "name": "pod-name"}, "node": {"name": "node-name"}}},"cloud":{"account":{"id":"account_id","name":"account_name"},"availability_zone":"cloud_availability_zone","instance":{"id":"instance_id","name":"instance_name"},"machine":{"type":"machine_type"},"project":{"id":"project_id","name":"project_name"},"provider":"cloud_provider","region":"cloud_region","service":{"name":"lambda"}}}}
+{"transaction": { "id": "945254c567a5417e", "trace_id": "0123456789abcdef0123456789abcdef", "parent_id": "abcdefabcdef01234567", "type": "request", "duration": 32.592981,  "span_count": { "started": 43 }}}
+{"transaction": { "id": "00xxxxFFaaaa1234", "trace_id": "0123456789abcdef0123456789abcdef", "name": "amqp receive", "parent_id": "abcdefabcdef01234567", "type": "messaging", "duration": 3, "span_count": { "started": 1 }, "context": {"message": {"queue": { "name": "new_users"}, "age":{ "ms": 1577958057123}, "headers": {"user_id": "1ax3", "involved_services": ["user", "auth"]}, "body": "user created", "routing_key": "user-created-transaction"}},"session":{"id":"sunday","sequence":123}}}
+`
+
+	// Create extension config
+	dataChannel := make(chan AgentData, 100)
+	config := extensionConfig{
+		apmServerSecretToken:       "foo",
+		apmServerApiKey:            "bar",
+		dataReceiverServerPort:     ":1234",
+		dataReceiverTimeoutSeconds: 15,
+	}
+
+	// Start extension server
+	StartHttpServer(dataChannel, &config)
+	defer agentDataServer.Close()
+
+	// Create a request to send to the extension
+	hosts, _ := net.LookupHost("localhost")
+	url := "http://" + hosts[0] + ":1234/intake/v2/events"
+	req, err := http.NewRequest("POST", url, strings.NewReader(agentRequestBody))
+	if err != nil {
+		t.Logf("Could not create request")
+	}
+	// Add headers to the request
+	for name, value := range headers {
+		req.Header.Add(name, value)
+	}
+
+	// Send the request to the extension
+	client := &http.Client{}
+	resp, err := client.Do(req)
+	if err != nil {
+		t.Logf("Error fetching %s, [%v]", agentDataServer.Addr, err)
+		t.Fail()
+	} else {
+		assert.Equal(t, 202, resp.StatusCode)
+	}
+}

apm-lambda-extension/extension/route_handlers.go

@@ -18,6 +18,7 @@
 package extension
 
 import (
+	"io"
 	"io/ioutil"
 	"log"
 	"net/http"
@@ -46,30 +47,33 @@ func handleInfoRequest(apmServerUrl string) func(w http.ResponseWriter, r *http.
 			return
 		}
 
-		resp, err := client.Do(req)
+		// Send request to apm server
+		serverResp, err := client.Do(req)
 		if err != nil {
 			log.Printf("error forwarding info request (`/`) to APM Server: %v", err)
 			return
 		}
-		defer resp.Body.Close()
-		body, err := ioutil.ReadAll(resp.Body)
-		if err != nil {
-			log.Printf("could not read info request response to APM Server: %v", err)
-			return
-		}
 
-		// send status code
-		w.WriteHeader(resp.StatusCode)
+		// If WriteHeader is not called explicitly, the first call to Write
+		// will trigger an implicit WriteHeader(http.StatusOK).
+		if serverResp.StatusCode != 200 {
+			w.WriteHeader(serverResp.StatusCode)
+		}
 
 		// send every header received
-		for name, values := range resp.Header {
+		for name, values := range serverResp.Header {
 			// Loop over all values for the name.
 			for _, value := range values {
 				w.Header().Add(name, value)
 			}
 		}
-		// send body
-		w.Write([]byte(body))
+
+		// copy body to request sent back to the agent
+		_, err = io.Copy(w, serverResp.Body)
+		if err != nil {
+			log.Printf("could not read info request response to APM Server: %v", err)
+			return
+		}
 	}
 }
 
@@ -80,12 +84,8 @@ func handleIntakeV2Events(agentDataChan chan AgentData) func(w http.ResponseWrit
 		w.WriteHeader(http.StatusAccepted)
 		w.Write([]byte("ok"))
 
-		if r.Body == nil {
-			log.Println("No body in agent request")
-			return
-		}
-
 		rawBytes, err := ioutil.ReadAll(r.Body)
+		defer r.Body.Close()
 		if err != nil {
 			log.Println("Could not read bytes from agent request body")
 			return