go-omni/httpx/bridge_test.go

package httpx_test

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

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

func TestLegacyResponseTrackingAndMiddleware(t *testing.T) {
	var logs bytes.Buffer
	closeLog, err := logx.Setup(logx.Config{Writer: &logs, Format: "json"})
	if err != nil {
		t.Fatal(err)
	}
	defer closeLog()
	app := omni.New(omni.WithLogger(true))
	app.GET("/legacy", httpx.WrapHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.NotFound(w, r) })))
	app.GET("/late", func(c *httpx.Context) error {
		c.Raw().WriteHeader(404)
		_, _ = c.Raw().Write([]byte("original"))
		return errors.New("late error")
	})
	app.GET("/panic", func(c *httpx.Context) error { _, _ = c.Raw().Write([]byte("partial")); panic("late panic") })
	for _, tc := range []struct {
		path   string
		status int
		body   string
	}{{"/legacy", 404, "404 page not found\n"}, {"/late", 404, "original"}, {"/panic", 200, "partial"}} {
		w := httptest.NewRecorder()
		app.ServeHTTP(w, httptest.NewRequest("GET", tc.path, nil))
		if w.Code != tc.status || w.Body.String() != tc.body {
			t.Fatalf("%s: %d %q", tc.path, w.Code, w.Body.String())
		}
	}
	count := 0
	for _, line := range strings.Split(logs.String(), "\n") {
		var e map[string]any
		if json.Unmarshal([]byte(line), &e) != nil || e["msg"] != "http.request" {
			continue
		}
		count++
		if e["path"] != "/panic" && e["status"] != float64(404) {
			t.Fatalf("wrong log: %s", line)
		}
	}
	if count != 3 {
		t.Fatalf("access logs: %d", count)
	}

	type key struct{}
	app = omni.New()
	app.Use(httpx.WrapFuncMiddleware(func(next http.HandlerFunc) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("X-Before", "yes")
			next(w, r.WithContext(context.WithValue(r.Context(), key{}, "identity")))
		}
	}))
	app.GET("/context", func(c *httpx.Context) error {
		if c.Context().Value(key{}) != "identity" || c.RequestID() == "" {
			t.Error("lost context")
		}
		return omni.NotFound("MISSING", "missing")
	})
	w := httptest.NewRecorder()
	app.ServeHTTP(w, httptest.NewRequest("GET", "/context", nil))
	if w.Code != 404 || w.Header().Get("X-Before") != "yes" || strings.Count(w.Body.String(), "MISSING") != 1 {
		t.Fatalf("middleware response: %d %s", w.Code, w.Body)
	}
}

type plainWriter struct {
	header http.Header
	status int
}

func (w *plainWriter) Header() http.Header         { return w.header }
func (w *plainWriter) WriteHeader(status int)      { w.status = status }
func (w *plainWriter) Write(p []byte) (int, error) { return len(p), nil }
func TestWriterCapabilitiesAndInformationalStatus(t *testing.T) {
	w := &plainWriter{header: http.Header{}}
	c := httpx.NewContext(httptest.NewRequest("GET", "/", nil), w, nil)
	if _, ok := c.Raw().(http.Flusher); ok {
		t.Fatal("invented Flusher")
	}
	if _, ok := c.Raw().(http.Hijacker); ok {
		t.Fatal("invented Hijacker")
	}
	if err := http.NewResponseController(c.Raw()).Flush(); !errors.Is(err, http.ErrNotSupported) {
		t.Fatalf("Flush: %v", err)
	}
	c.Raw().WriteHeader(103)
	if c.Responded() {
		t.Fatal("103 committed final response")
	}
	c.Raw().WriteHeader(201)
	c.Raw().WriteHeader(500)
	if !c.Responded() || c.Status() != 201 || w.status != 201 {
		t.Fatal("first final status not preserved")
	}
	recorder := httptest.NewRecorder()
	c = httpx.NewContext(httptest.NewRequest("GET", "/", nil), recorder, nil)
	if err := http.NewResponseController(c.Raw()).Flush(); err != nil {
		t.Fatal(err)
	}
	if !c.Responded() || c.Status() != 200 || !recorder.Flushed {
		t.Fatal("flush not tracked")
	}
}

func TestLegacyStreamingAcrossTransports(t *testing.T) {
	for _, transport := range []string{omni.TransportNetHTTP, omni.TransportNative} {
		t.Run(transport, func(t *testing.T) {
			release := make(chan struct{})
			defer close(release)
			app := omni.New(omni.WithTransport(transport))
			app.GET("/stream", httpx.WrapHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "text/plain")
				_, _ = io.WriteString(w, "first\n")
				if err := http.NewResponseController(w).Flush(); err != nil {
					t.Error(err)
					return
				}
				select {
				case <-release:
				case <-r.Context().Done():
				}
			})))
			app.GET("/file", httpx.WrapHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Disposition", `attachment; filename="data.bin"`)
				_, _ = io.Copy(w, strings.NewReader(strings.Repeat("x", 2<<20)))
			})))
			app.HandleWithBodyLimit("POST", "/bounded", 32, nil, func(c *httpx.Context) error {
				_, err := io.Copy(io.Discard, c.Request().Body)
				if err != nil {
					return err
				}
				return c.NoContent()
			})
			listener, err := net.Listen("tcp", "127.0.0.1:0")
			if err != nil {
				t.Fatal(err)
			}
			ctx, cancel := context.WithCancel(context.Background())
			done := make(chan error, 1)
			go func() { done <- app.ServeContext(ctx, listener) }()
			defer func() {
				cancel()
				select {
				case <-done:
				case <-time.After(15 * time.Second):
					t.Error("shutdown timed out")
				}
			}()
			client := &http.Client{Timeout: 3 * time.Second}
			defer client.CloseIdleConnections()
			for _, body := range []io.Reader{strings.NewReader(strings.Repeat("x", 33)), io.NopCloser(strings.NewReader(strings.Repeat("x", 33)))} {
				res, err := client.Post("http://"+listener.Addr().String()+"/bounded", "application/octet-stream", body)
				if err != nil {
					t.Fatal(err)
				}
				_, _ = io.Copy(io.Discard, res.Body)
				res.Body.Close()
				if res.StatusCode != 413 {
					t.Fatalf("real transport body limit: %d", res.StatusCode)
				}
			}
			res, err := client.Get("http://" + listener.Addr().String() + "/stream")
			if err != nil {
				t.Fatal(err)
			}
			first := make([]byte, 6)
			_, err = io.ReadFull(res.Body, first)
			res.Body.Close()
			if err != nil || string(first) != "first\n" {
				t.Fatalf("stream was buffered: %s %v", first, err)
			}
			res, err = client.Get("http://" + listener.Addr().String() + "/file")
			if err != nil {
				t.Fatal(err)
			}
			size, err := io.Copy(io.Discard, res.Body)
			res.Body.Close()
			if err != nil || size != 2<<20 || res.Header.Get("Content-Disposition") == "" {
				t.Fatalf("file: %d %v", size, err)
			}
		})
	}
}

func TestStandardMiddlewareStateAndShortCircuit(t *testing.T) {
	constructions, calls := 0, 0
	app := omni.New()
	defer app.Close()
	app.Use(httpx.WrapMiddleware(func(next http.Handler) http.Handler {
		constructions++
		count := 0
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			count++
			if count == 2 {
				http.Error(w, "denied", 401)
				return
			}
			next.ServeHTTP(w, r)
		})
	}))
	app.GET("/", func(c *httpx.Context) error { calls++; return c.NoContent() })
	for _, status := range []int{204, 401, 204} {
		w := httptest.NewRecorder()
		app.ServeHTTP(w, httptest.NewRequest("GET", "/", nil))
		if w.Code != status {
			t.Fatalf("got %d want %d", w.Code, status)
		}
	}
	if constructions != 1 || calls != 2 {
		t.Fatalf("middleware lifecycle: %d constructions, %d calls", constructions, calls)
	}
}

func TestStandardMiddlewareTransformsErrorAndPreservesCancellation(t *testing.T) {
	app := omni.New()
	defer app.Close()
	app.Use(httpx.WrapMiddleware(func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// A legacy middleware may buffer its own small response. The bridge must
			// report its eventual physical status, and encode errors before it unwinds.
			temp := httptest.NewRecorder()
			next.ServeHTTP(temp, r)
			if temp.Code != 404 {
				t.Errorf("inner error not encoded: %d", temp.Code)
			}
			w.WriteHeader(409)
			_, _ = io.WriteString(w, strings.ToUpper(temp.Body.String()))
		})
	}))
	app.GET("/", func(c *httpx.Context) error {
		if !errors.Is(c.Context().Err(), context.Canceled) {
			t.Error("cancellation lost")
		}
		return omni.NotFound("missing", "missing")
	})
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	w := httptest.NewRecorder()
	app.ServeHTTP(w, httptest.NewRequest("GET", "/", nil).WithContext(ctx))
	if w.Code != 409 || !strings.Contains(w.Body.String(), "MISSING") {
		t.Fatalf("%d %s", w.Code, w.Body)
	}
}

func TestRouteBodyLimitBeforeMiddleware(t *testing.T) {
	for _, length := range []int64{-1, 33} {
		t.Run(fmt.Sprint(length), func(t *testing.T) {
			app := omni.New()
			defer app.Close()
			read := 0
			app.Use(func(next httpx.Handler) httpx.Handler {
				return func(c *httpx.Context) error {
					body, err := io.ReadAll(c.Request().Body)
					read = len(body)
					if err != nil {
						return err
					}
					return next(c)
				}
			})
			app.Group("/api").HandleWithBodyLimit("POST", "/upload", 32, nil, func(c *httpx.Context) error { t.Error("called oversized handler"); return nil })
			r := httptest.NewRequest("POST", "/api/upload", strings.NewReader(strings.Repeat("x", 33)))
			r.ContentLength = length
			w := httptest.NewRecorder()
			app.ServeHTTP(w, r)
			if w.Code != 413 || read > 32 {
				t.Fatalf("limit: status=%d read=%d", w.Code, read)
			}
		})
	}
}

func TestMultipartCleanupOnPanic(t *testing.T) {
	temp := t.TempDir()
	t.Setenv("TMPDIR", temp)
	var body bytes.Buffer
	form := multipart.NewWriter(&body)
	part, err := form.CreateFormFile("file", "x.bin")
	if err != nil {
		t.Fatal(err)
	}
	_, _ = part.Write(bytes.Repeat([]byte("x"), 4096))
	_ = form.Close()
	app := omni.New()
	defer app.Close()
	app.HandleWithBodyLimit("POST", "/", 8192, nil, httpx.WrapHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := r.ParseMultipartForm(1); err != nil {
			t.Error(err)
			return
		}
		panic("after disk spill")
	})))
	r := httptest.NewRequest("POST", "/", &body)
	r.Header.Set("Content-Type", form.FormDataContentType())
	w := httptest.NewRecorder()
	app.ServeHTTP(w, r)
	entries, err := os.ReadDir(temp)
	if w.Code != 500 || err != nil || len(entries) != 0 {
		t.Fatalf("cleanup: %d %v %v", w.Code, entries, err)
	}
}

func TestStandardMiddlewareCommitsBeforeCallingNext(t *testing.T) {
	app := omni.New()
	defer app.Close()
	app.Use(httpx.WrapMiddleware(func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(202)
			_, _ = io.WriteString(w, "accepted")
			next.ServeHTTP(w, r)
		})
	}))
	app.GET("/", func(c *httpx.Context) error {
		if !c.Responded() || c.Status() != 202 {
			t.Error("lost middleware commit")
		}
		return errors.New("late failure")
	})
	w := httptest.NewRecorder()
	app.ServeHTTP(w, httptest.NewRequest("GET", "/", nil))
	if w.Code != 202 || w.Body.String() != "accepted" {
		t.Fatalf("late append: %d %q", w.Code, w.Body.String())
	}
}

func TestBridgeStandardPathValues(t *testing.T) {
	app := omni.New()
	defer app.Close()
	app.Use(httpx.WrapMiddleware(func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.PathValue("itemId") != "abc" {
				t.Error("middleware path value lost")
			}
			next.ServeHTTP(w, r)
		})
	}))
	app.GET("/items/:itemId", httpx.WrapHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, r.PathValue("itemId")) })))
	w := httptest.NewRecorder()
	app.ServeHTTP(w, httptest.NewRequest("GET", "/items/abc", nil))
	if w.Body.String() != "abc" {
		t.Fatal("handler path value lost")
	}
}