go-omni/examples/reference/transport_test.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	omni "git.yuntuo.hk/k2safe/go-omni"
	"git.yuntuo.hk/k2safe/go-omni/auth"
	"git.yuntuo.hk/k2safe/go-omni/examples/reference/repository"
	"git.yuntuo.hk/k2safe/go-omni/internal/testutil"
)

// networkHandler runs the same business assertions through a real TCP client.
// The application handler is only called by its configured server transport.
type networkHandler struct {
	t                            *testing.T
	client                       *http.Client
	address, scheme, token, alpn string
	proto                        int
	requests                     *int
}

func (h networkHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h.t.Helper()
	req := r.Clone(r.Context())
	req.RequestURI = ""
	req.URL.Scheme, req.URL.Host, req.Host = h.scheme, h.address, h.address
	if h.token != "" && req.Header.Get("Authorization") == "" {
		req.Header.Set("Authorization", "Bearer "+h.token)
	}
	response, err := h.client.Do(req)
	if err != nil {
		h.t.Fatal(err)
	}
	defer response.Body.Close()
	if strings.HasPrefix(req.URL.Path, "/api/v1/") {
		*h.requests++
	}
	if response.ProtoMajor != h.proto {
		h.t.Fatalf("protocol %s, want HTTP/%d", response.Proto, h.proto)
	}
	if h.scheme == "https" && (response.TLS == nil || len(response.TLS.VerifiedChains) == 0 || response.TLS.NegotiatedProtocol != h.alpn) {
		h.t.Fatalf("TLS/ALPN verification failed: %+v", response.TLS)
	}
	for key, values := range response.Header {
		w.Header()[key] = append([]string(nil), values...)
	}
	w.WriteHeader(response.StatusCode)
	if _, err := io.Copy(w, response.Body); err != nil {
		h.t.Fatal(err)
	}
}

type lockedLog struct {
	sync.Mutex
	bytes.Buffer
}

func (b *lockedLog) Write(p []byte) (int, error) {
	b.Lock()
	defer b.Unlock()
	return b.Buffer.Write(p)
}

func (b *lockedLog) snapshot() string {
	b.Lock()
	defer b.Unlock()
	return b.Buffer.String()
}

func TestBusinessConfiguredTransports(t *testing.T) {
	for _, test := range []struct {
		name, transport string
		tls             bool
		proto           int
		alpn            string
	}{
		{"nethttp", omni.TransportNetHTTP, false, 1, ""},
		{"native", omni.TransportNative, false, 1, ""},
		{"nethttpTLS", omni.TransportNetHTTP, true, 2, "h2"},
		{"nativeTLS", omni.TransportNative, true, 1, "http/1.1"},
		{"http2TLS", omni.TransportHTTP2, true, 2, "h2"},
	} {
		t.Run(test.name, func(t *testing.T) {
			cert, key, tlsClient := testutil.CertificateFiles(t)
			t.Setenv("REFERENCE_TEST_TRANSPORT", test.transport)
			t.Setenv("REFERENCE_TEST_SHUTDOWN_TIMEOUT", "3s")
			t.Setenv("REFERENCE_TEST_TLS_CERT_FILE", "")
			t.Setenv("REFERENCE_TEST_TLS_KEY_FILE", "")
			if test.tls {
				t.Setenv("REFERENCE_TEST_TLS_CERT_FILE", cert)
				t.Setenv("REFERENCE_TEST_TLS_KEY_FILE", key)
			}
			path := filepath.Join(t.TempDir(), "config.yaml")
			if err := os.WriteFile(path, []byte("transport: ${REFERENCE_TEST_TRANSPORT}\nshutdownTimeout: 3s\ntls:\n  certFile: \"${REFERENCE_TEST_TLS_CERT_FILE:-}\"\n  keyFile: \"${REFERENCE_TEST_TLS_KEY_FILE:-}\""), 0600); err != nil {
				t.Fatal(err)
			}
			settings, err := omni.LoadConfig(path)
			if err != nil {
				t.Fatal(err)
			}
			var logs lockedLog
			logger := testLogger(&logs)
			application := newProtectedApp(repository.NewMemoryUsers(), logger, auth.StaticAuthenticator{
				"admin": {UserID: "admin", Roles: []string{"admin"}},
				"user":  {UserID: "user", Roles: []string{"user"}},
			}, omni.WithConfig(settings))
			application.OnShutdown(logger.Close)
			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 <- application.ServeContext(ctx, listener) }()
			t.Cleanup(func() {
				cancel()
				select {
				case err := <-done:
					if err != nil {
						t.Error(err)
					}
				case <-time.After(5 * time.Second):
					t.Error("business server did not stop")
					listener.Close()
				}
			})
			client := &http.Client{Timeout: 5 * time.Second, Transport: &http.Transport{TLSClientConfig: tlsClient, ForceAttemptHTTP2: true}}
			defer client.CloseIdleConnections()
			scheme := "http"
			if test.tls {
				scheme = "https"
			}
			requests := 0
			network := networkHandler{t: t, client: client, address: listener.Addr().String(), scheme: scheme, proto: test.proto, alpn: test.alpn, requests: &requests}
			request(t, network, "POST", "/api/v1/users", `{"name":"Alice","email":"alice@example.com"}`, 401, "UNAUTHENTICATED")
			requestWithAuthorization(t, network, "DELETE", "/api/v1/users/1", "", "Bearer user", 403, "FORBIDDEN")
			network.token = "admin"
			checkUserCRUD(t, network)
			request(t, network, "POST", "/api/v1/users", `{"name":`, 400, "INVALID_BODY")
			count := 0
			for _, line := range strings.Split(strings.TrimSpace(logs.snapshot()), "\n") {
				var record map[string]any
				if err := json.Unmarshal([]byte(line), &record); err != nil {
					t.Fatal(err)
				}
				if record["requestId"] == "reference-test" {
					count++
				}
			}
			if count < requests {
				t.Fatalf("missing correlated business logs: %d", count)
			}
		})
	}
}

func TestReferenceExplicitConfig(t *testing.T) {
	previous := *configPath
	*configPath = "config.yaml"
	defer func() { *configPath = previous }()
	t.Setenv("GO_OMNI_TRANSPORT", "native")
	t.Setenv("GO_OMNI_HTTP_ADDR", "")
	t.Setenv("GO_OMNI_TLS_CERT_FILE", "")
	t.Setenv("GO_OMNI_TLS_KEY_FILE", "")
	c, err := referenceConfig()
	if err != nil || c.Address != "127.0.0.1:8080" || c.Transport != "native" {
		t.Fatal(c, err)
	}
	t.Setenv("GO_OMNI_ADDR", "127.0.0.1:1234")
	c, err = referenceConfig()
	if err != nil || c.Address != "127.0.0.1:8080" {
		t.Fatal("implicit alias", err)
	}
	t.Setenv("GO_OMNI_HTTP_ADDR", "127.0.0.1:8082")
	c, err = referenceConfig()
	if err != nil || c.Address != "127.0.0.1:8082" {
		t.Fatal(c, err)
	}
}