middleware_test.go

package logjam

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"sort"
	"strings"
	"testing"
	"time"

	"log"

	"github.com/golang/snappy"
	"github.com/gorilla/mux"
	"github.com/pebbe/zmq4"
	. "github.com/smartystreets/goconvey/convey"
)

func TestObfuscateIP(t *testing.T) {
	Convey("Obfuscate IP", t, func() {
		ips := map[string]string{
			"0000:0000:0000:0000:0000:FFFF:C0A8:1": "192.168.0.XXX",
			"192.168.0.1":                          "192.168.0.XXX",
			"::FFFF:192.168.0.1":                   "192.168.0.XXX",
			"::FFFF:C0A8:1":                        "192.168.0.XXX",
			"fe80::da50:e6ff:fedb:c252":            "fe80::da50:e6ff:fedb:XXXX",
			"::fedb:c252":                          "::fedb:XXXX",
			"invalid":                              "invalid",
		}

		// this is mostly so test order is always deterministic
		keys := []string{}
		for k := range ips {
			keys = append(keys, k)
		}
		sort.Strings(keys)

		for _, key := range keys {
			Convey(key, func() {
				So(obfuscateIP(key), ShouldEqual, ips[key])
			})
		}
	})
}

type recoveryHandler struct {
	panicked *bool
	handler  http.Handler
}

func (h recoveryHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	defer func() { *h.panicked = recover() != nil }()
	h.handler.ServeHTTP(w, req)
}

func TestMiddleware(t *testing.T) {
	os.Setenv("HOSTNAME", "test-machine")
	os.Setenv("DATACENTER", "dc")
	os.Setenv("CLUSTER", "a")
	os.Setenv("NAMESPACE", "logjam")
	setRequestEnv()

	router := mux.NewRouter()
	logger := Logger{Logger: log.New(ioutil.Discard, "", 0)}

	router.Path("/rest/app/vendor/v1/users/123").HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		ctx := req.Context()
		logger.Debug(ctx, "First Line")
		logger.Info(ctx, "Second Line")
		logger.Warn(ctx, "Third Line")
		logger.Error(ctx, "Fourth Line")
		logger.Exception(ctx, "X1", "found")
		logger.Exceptionf(ctx, "X2", "occurred %d time", 1)

		r := GetRequest(ctx)
		r.AddCount("rest_calls", 1)
		r.AddDuration("rest_time", 100*time.Millisecond)
		r.SetField("sender_id", "foobar")
		r.MeasureDuration("view_time", func() {
			time.Sleep(100 * time.Millisecond)
			w.WriteHeader(200)
			w.Write([]byte(`some body`))
		})
	})

	router.Path("/panic-before-writing-header").HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		time.Sleep(100 * time.Millisecond)
		panic("panic")
	})

	router.Path("/panic-after-writing-header").HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		time.Sleep(100 * time.Millisecond)
		w.WriteHeader(200)
		panic("panic")
	})

	socket, err := zmq4.NewSocket(zmq4.ROUTER)
	if err != nil {
		panic("could not create socket")
	}
	err = socket.Bind("inproc://middleware-test")
	if err != nil {
		panic("could not bind socket")
	}
	defer socket.Close()

	agentOptions := Options{
		Endpoints:    "inproc://middleware-test",
		AppName:      "appName",
		EnvName:      "envName",
		Logger:       logger,
		ObfuscateIPs: true,
	}
	agent := NewAgent(&agentOptions)
	defer agent.Shutdown()

	Convey("full request/response cycle - happy path", t, func() {
		m := agent.NewMiddleware(MiddlewareOptions{})
		server := httptest.NewServer(m(router))
		defer server.Close()

		callerID := "27ce93ab-05e7-48b8-a80c-6e076c32b75a"
		traceID := "2ac5d40fd8f54c3d9def295f1adac47d"
		actionName := "Rest::App::Vendor::V1::Users::Id#get"

		req, err := http.NewRequest("GET", server.URL+"/rest/app/vendor/v1/users/123", nil)
		req.Header.Set("X-Logjam-Caller-Id", callerID)
		req.Header.Set("X-Logjam-Trace-Id", traceID)
		req.Header.Set("Authorization", "4ec04124-bd41-49e2-9e30-5b189f5ca5f2")
		query := req.URL.Query()
		query.Set("single", "value")
		query.Set("multi", "value1")
		query.Add("multi", "value2")
		req.URL.RawQuery = query.Encode()
		So(err, ShouldBeNil)

		now := time.Now()
		res, err := server.Client().Do(req)

		So(err, ShouldBeNil)
		So(res.StatusCode, ShouldEqual, 200)
		requestID := res.Header.Get("X-Logjam-Request-Id")
		So(requestID, ShouldStartWith, "appName-envName-")
		requestParts := strings.Split(requestID, "-")
		uuid := requestParts[len(requestParts)-1]
		So(res.Header.Get("X-Logjam-Action"), ShouldEqual, actionName)
		So(res.Header.Get("X-Logjam-Caller-Id"), ShouldEqual, callerID)
		So(res.Header.Get("Http-Authorization"), ShouldEqual, "")

		msg, err := socket.RecvMessage(0)
		So(err, ShouldBeNil)
		So(msg, ShouldHaveLength, 5)

		So(msg[1], ShouldEqual, agent.AppName+"-"+agent.EnvName)
		So(msg[2], ShouldEqual, "logs."+agentOptions.AppName+"."+agentOptions.EnvName)

		payload, err := snappy.Decode(nil, []byte(msg[3]))
		So(err, ShouldBeNil)

		output := map[string]interface{}{}
		json.Unmarshal([]byte(payload), &output)

		So(output["action"], ShouldEqual, actionName)
		So(output["host"], ShouldEqual, "test-machine")
		So(output["ip"], ShouldEqual, "127.0.0.XXX")
		So(output["process_id"].(float64), ShouldNotEqual, 0)
		So(output["request_id"], ShouldEqual, uuid)
		So(output["trace_id"], ShouldEqual, traceID)
		So(output["started_at"], shouldHaveTimeFormat, timeFormat)
		startedAt, err := time.ParseInLocation(timeFormat, output["started_at"].(string), now.Location())
		So(err, ShouldBeNil)
		So(uint64(startedAt.UnixNano()/1000000), ShouldEqual, output["started_ms"])
		So(output["started_ms"], ShouldAlmostEqual, uint64(now.UnixNano()/1000000), 100)
		So(output["total_time"], ShouldBeGreaterThan, 100)
		So(output["total_time"], ShouldAlmostEqual, 100, 10)
		So(output["rest_calls"], ShouldEqual, 1)
		totalTime := output["total_time"].(float64)
		viewTime := output["view_time"].(float64)
		restTime := output["rest_time"].(float64)
		So(totalTime, ShouldBeGreaterThanOrEqualTo, viewTime+restTime)
		So(output["datacenter"], ShouldEqual, "dc")
		So(output["cluster"], ShouldEqual, "a")
		So(output["namespace"], ShouldEqual, "logjam")
		So(output["sender_id"], ShouldEqual, "foobar")

		exceptions := output["exceptions"]
		So(exceptions, ShouldHaveLength, 2)
		So(exceptions, ShouldContain, "X1")
		So(exceptions, ShouldContain, "X2")

		requestInfo := output["request_info"].(map[string]interface{})
		So(requestInfo["method"], ShouldEqual, "GET")

		So(requestInfo["url"], ShouldContainSubstring, "/rest/app/vendor/v1/users/123")
		So(requestInfo["url"], ShouldContainSubstring, "multi=value1")
		So(requestInfo["url"], ShouldContainSubstring, "multi=value2")
		So(requestInfo["url"], ShouldContainSubstring, "single=value")

		So(requestInfo["headers"], ShouldResemble, map[string]interface{}{
			"Accept-Encoding":    "gzip",
			"User-Agent":         "Go-http-client/1.1",
			"X-Logjam-Caller-Id": callerID,
			"X-Logjam-Trace-Id":  traceID,
		})

		So(requestInfo["query_parameters"], ShouldResemble, map[string]interface{}{
			"multi":  []interface{}{"value1", "value2"},
			"single": "value"})

		So(requestInfo["body_parameters"], ShouldBeNil)

		lines := output["lines"].([]interface{})
		So(lines, ShouldHaveLength, 6)

		line := lines[0].([]interface{})
		So(line, ShouldHaveLength, 3)
		So(line[0], ShouldEqual, DEBUG) // severity
		So(line[1], shouldHaveTimeFormat, timeFormat)
		So(line[2], ShouldEqual, "First Line")

		line = lines[1].([]interface{})
		So(line, ShouldHaveLength, 3)
		So(line[0], ShouldEqual, INFO) // severity
		So(line[1], shouldHaveTimeFormat, timeFormat)
		So(line[2], ShouldEqual, "Second Line")

		line = lines[2].([]interface{})
		So(line, ShouldHaveLength, 3)
		So(line[0], ShouldEqual, WARN) // severity
		So(line[1], shouldHaveTimeFormat, timeFormat)
		So(line[2], ShouldEqual, "Third Line")

		line = lines[3].([]interface{})
		So(line, ShouldHaveLength, 3)
		So(line[0], ShouldEqual, ERROR) // severity
		So(line[1], shouldHaveTimeFormat, timeFormat)
		So(line[2], ShouldEqual, "Fourth Line")

		line = lines[4].([]interface{})
		So(line, ShouldHaveLength, 3)
		So(line[0], ShouldEqual, ERROR) // severity
		So(line[1], shouldHaveTimeFormat, timeFormat)
		So(line[2], ShouldEqual, "X1: found")

		line = lines[5].([]interface{})
		So(line, ShouldHaveLength, 3)
		So(line[0], ShouldEqual, ERROR) // severity
		So(line[1], shouldHaveTimeFormat, timeFormat)
		So(line[2], ShouldEqual, "X2: occurred 1 time")

	})

	Convey("full request/response cycle - handling panics", t, func() {
		tests := []struct {
			Path       string
			ActionName string
			Code       int
			Panic      string
		}{
			{
				Path:       "/panic-before-writing-header",
				ActionName: "PanicBeforeWritingHeader#get",
				Code:       500,
				Panic:      "supressed",
			},
			{
				Path:       "/panic-after-writing-header",
				ActionName: "PanicAfterWritingHeader#get",
				Code:       200,
				Panic:      "supressed",
			},
			{
				Path:       "/panic-before-writing-header",
				ActionName: "PanicBeforeWritingHeader#get",
				Code:       500,
				Panic:      "bubbling up",
			},
		}

		for _, test := range tests {
			Convey(test.Path+" - "+test.Panic, func() {
				bubblePanics := test.Panic == "bubbling up"
				panicked := false
				r := agent.NewHandler(router, MiddlewareOptions{BubblePanics: bubblePanics})
				server := httptest.NewServer(recoveryHandler{handler: r, panicked: &panicked})
				defer server.Close()

				req, err := http.NewRequest("GET", server.URL+test.Path, nil)
				So(err, ShouldBeNil)

				now := time.Now()
				res, err := server.Client().Do(req)
				So(err, ShouldBeNil)
				So(panicked, ShouldEqual, bubblePanics)

				So(res.StatusCode, ShouldEqual, test.Code)
				requestID := res.Header.Get("X-Logjam-Request-Id")
				So(requestID, ShouldStartWith, "appName-envName-")
				requestParts := strings.Split(requestID, "-")
				uuid := requestParts[len(requestParts)-1]
				So(res.Header.Get("X-Logjam-Action"), ShouldEqual, test.ActionName)

				msg, err := socket.RecvMessage(0)
				So(err, ShouldBeNil)
				So(msg, ShouldHaveLength, 5)

				So(msg[1], ShouldEqual, agent.AppName+"-"+agent.EnvName)
				So(msg[2], ShouldEqual, "logs."+agentOptions.AppName+"."+agentOptions.EnvName)

				payload, err := snappy.Decode(nil, []byte(msg[3]))
				So(err, ShouldBeNil)

				output := map[string]interface{}{}
				json.Unmarshal([]byte(payload), &output)

				So(output["action"], ShouldEqual, test.ActionName)
				So(output["host"], ShouldEqual, "test-machine")
				So(output["ip"], ShouldEqual, "127.0.0.XXX")
				So(output["process_id"].(float64), ShouldNotEqual, 0)
				So(output["request_id"], ShouldEqual, uuid)
				So(output["started_at"], shouldHaveTimeFormat, timeFormat)
				startedAt, err := time.ParseInLocation(timeFormat, output["started_at"].(string), now.Location())
				So(err, ShouldBeNil)
				So(uint64(startedAt.UnixNano()/1000000), ShouldEqual, output["started_ms"])
				So(output["started_ms"], ShouldAlmostEqual, uint64(now.UnixNano()/1000000), 100)
				So(output["total_time"], ShouldBeGreaterThan, 100)
				So(output["total_time"], ShouldAlmostEqual, 100, 10)
				So(output["datacenter"], ShouldEqual, "dc")
				So(output["cluster"], ShouldEqual, "a")
				So(output["namespace"], ShouldEqual, "logjam")

				requestInfo := output["request_info"].(map[string]interface{})
				So(requestInfo["method"], ShouldEqual, "GET")

				So(requestInfo["url"], ShouldContainSubstring, test.Path)

				lines := output["lines"].([]interface{})
				So(lines, ShouldHaveLength, 1)

				line := lines[0].([]interface{})
				So(line, ShouldHaveLength, 3)
				So(line[0], ShouldEqual, FATAL) // severity
				So(line[1], shouldHaveTimeFormat, timeFormat)
				So(line[2], ShouldStartWith, `"panic":`)
			})
		}
	})
}

func TestSetCallHeaders(t *testing.T) {
	Convey("SetLogjamHeaders", t, func() {
		agentOptions := Options{
			Endpoints:    "127.0.0.1,localhost",
			AppName:      "appName",
			EnvName:      "envName",
			Logger:       log.New(ioutil.Discard, "", 0),
			ObfuscateIPs: true,
		}
		agent := NewAgent(&agentOptions)
		defer agent.Shutdown()

		incoming := httptest.NewRequest("GET", "/", nil)
		logjamRequest := agent.NewRequest("foobar")
		wrapped := logjamRequest.AugmentRequest(incoming)
		outgoing := httptest.NewRequest("GET", "/", nil)
		SetCallHeaders(wrapped.Context(), outgoing)
		So(outgoing.Header.Get("X-Logjam-Action"), ShouldEqual, "foobar")
		So(outgoing.Header.Get("X-Logjam-Caller-Id"), ShouldEqual, logjamRequest.id)
	})
}

func shouldHaveTimeFormat(actual interface{}, expected ...interface{}) string {
	_, err := time.Parse(expected[0].(string), actual.(string))
	if err != nil {
		return err.Error()
	}
	return ""
}