go-omni/middleware/middleware_test.go

package middleware_test

import (
	"bytes"
	"encoding/json"
	"errors"
	"io"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
	"git.yuntuo.hk/k2safe/go-omni/httpx"
	"git.yuntuo.hk/k2safe/go-omni/logx"
	"git.yuntuo.hk/k2safe/go-omni/middleware"
)

func TestAccessLogReportsCommittedStatus(t *testing.T) {
	for _, tc := range []struct {
		name    string
		handler httpx.Handler
		status  int
	}{
		{"unwritten error", func(c *httpx.Context) error { return apierrors.NotFound("MISSING", "missing") }, 404},
		{"encoding error", func(c *httpx.Context) error { return c.OK(make(chan int)) }, 500},
		{"committed error", func(c *httpx.Context) error { _ = c.Created("saved"); return errors.New("later error") }, 201},
		{"committed panic", func(c *httpx.Context) error { _ = c.Accepted("saved"); panic("later panic") }, 202},
	} {
		t.Run(tc.name, func(t *testing.T) {
			var output bytes.Buffer
			logger, err := logx.New(logx.Config{Format: "json", Writer: &output})
			if err != nil {
				t.Fatal(err)
			}
			logx.SetDefault(logger)
			defer logx.SetDefault(nil)
			req := httptest.NewRequest("GET", "/users/42", nil)
			req.Header.Set("X-Request-ID", "logRequest")
			req = req.WithContext(httpx.WithRoutePattern(req.Context(), "/users/:id"))
			w := httptest.NewRecorder()
			ctx := httpx.NewContext(req, w, nil)
			chain := middleware.RequestID("X-Request-ID")(middleware.AccessLog()(middleware.Recovery()(tc.handler)))
			if err := chain(ctx); err != nil {
				ctx.WriteError(err)
			}
			var entry struct {
				Message   string `json:"msg"`
				Status    int
				Route     string
				RequestID string `json:"requestId"`
			}
			decoder := json.NewDecoder(&output)
			for {
				if err := decoder.Decode(&entry); err != nil {
					t.Fatal(err)
				}
				if entry.Message == "http.request" {
					break
				}
			}
			if entry.Status != tc.status || w.Code != tc.status || entry.Route != "/users/:id" || entry.RequestID != "logRequest" {
				t.Fatalf("log=%s response=%d; expected status %d and request metadata", &output, w.Code, tc.status)
			}
		})
	}
}

func TestAccessLogUsesVerifiedClientIP(t *testing.T) {
	loggerOutput := bytes.Buffer{}
	logger, err := logx.New(logx.Config{Format: "json", Writer: &loggerOutput})
	if err != nil {
		t.Fatal(err)
	}
	logx.SetDefault(logger)
	defer logx.SetDefault(nil)
	resolver, err := httpx.NewClientIPResolver("127.0.0.1")
	if err != nil {
		t.Fatal(err)
	}
	request := httptest.NewRequest("GET", "/health", nil)
	request.RemoteAddr = "127.0.0.1:443"
	request.Header.Set("X-Forwarded-For", "198.51.100.7")
	request = request.WithContext(httpx.WithClientIPResolver(request.Context(), resolver))
	ctx := httpx.NewContext(request, httptest.NewRecorder(), nil)
	if err := middleware.AccessLog()(func(c *httpx.Context) error { return c.OK("ok") })(ctx); err != nil {
		t.Fatal(err)
	}
	var entry struct {
		ClientIP string `json:"clientIp"`
	}
	if err := json.Unmarshal(loggerOutput.Bytes(), &entry); err != nil {
		t.Fatal(err)
	}
	if entry.ClientIP != "198.51.100.7" {
		t.Fatalf("expected verified client IP, got %q", entry.ClientIP)
	}
}

func TestSlowAndPanicLogs(t *testing.T) {
	var out bytes.Buffer
	logger, _ := logx.New(logx.Config{Writer: &out})
	logx.SetDefault(logger)
	defer logx.SetDefault(nil)
	ctx := httpx.NewContext(httptest.NewRequest("GET", "/panic?secret=hidden", nil), httptest.NewRecorder(), nil)
	chain := middleware.RequestID("X-Request-ID")(middleware.AccessLog(middleware.WithSlowRequestThreshold(time.Nanosecond))(middleware.Recovery()(func(c *httpx.Context) error { panic("boom") })))
	err := chain(ctx)
	if err == nil {
		t.Fatal("panic not recovered")
	}
	decoder := json.NewDecoder(&out)
	count := 0
	for {
		var entry map[string]any
		err := decoder.Decode(&entry)
		if err == io.EOF {
			break
		}
		if err != nil {
			t.Fatal(err)
		}
		count++
		if entry["requestId"] == "" || entry["level"] != "ERROR" {
			t.Fatal(entry)
		}
		if entry["msg"] == "http.panic" && !strings.Contains(entry["stack"].(string), "TestSlowAndPanicLogs") {
			t.Fatal(entry)
		}
		if entry["msg"] == "http.request" && (entry["slow"] != true || entry["status"] != float64(500)) {
			t.Fatal(entry)
		}
	}
	if count != 2 {
		t.Fatal(count)
	}
	out.Reset()
	ctx = httpx.NewContext(httptest.NewRequest("GET", "/slow", nil), httptest.NewRecorder(), nil)
	err = middleware.AccessLog(middleware.WithSlowRequestThreshold(time.Nanosecond))(func(c *httpx.Context) error { return c.OK("done") })(ctx)
	var entry map[string]any
	if err != nil || json.Unmarshal(out.Bytes(), &entry) != nil || entry["level"] != "WARN" || entry["slow"] != true {
		t.Fatal(out.String(), err)
	}
}