telemt-api/internal/proxy/mihomo_test.go

package proxy

import (
	"context"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"
)

func TestMihomoForwardRewritesPathAndAuth(t *testing.T) {
	target, err := url.Parse("http://127.0.0.1:9090")
	if err != nil {
		t.Fatal(err)
	}
	cap := &captureTransport{}
	h := NewMihomoForward(target, "/api/mtg/mihomo", "Bearer testsecret", cap, nil)

	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "http://gw/api/mtg/mihomo/proxies", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Authorization", "Bearer client-should-not-forward")
	req.Header.Set("Host", "public-gateway.example:8888")
	h.ServeHTTP(httptest.NewRecorder(), req)

	if cap.got == nil {
		t.Fatal("no outgoing request captured")
	}
	if got, want := cap.got.Host, "127.0.0.1:9090"; got != want {
		t.Fatalf("Host: got %q want %q (upstream must not see client Host)", got, want)
	}
	if got := cap.got.Header.Get("Authorization"); got != "Bearer testsecret" {
		t.Fatalf("Authorization: got %q want Bearer testsecret", got)
	}
	want, err := url.Parse("http://127.0.0.1:9090/proxies")
	if err != nil {
		t.Fatal(err)
	}
	assertSameURL(t, cap.got.URL, want)
}

func TestUseMihomoWebSocketTunnel(t *testing.T) {
	t.Run("gorilla IsWebSocketUpgrade", func(t *testing.T) {
		r := httptest.NewRequest(http.MethodGet, "/api/x/mihomo/traffic", nil)
		r.Header.Set("Connection", "Upgrade")
		r.Header.Set("Upgrade", "websocket")
		r.Header.Set("Sec-WebSocket-Version", "13")
		r.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
		if !useMihomoWebSocketTunnel("traffic", r) {
			t.Fatal("expected tunnel for full WS handshake")
		}
	})

	t.Run("traffic with key only (no Connection token)", func(t *testing.T) {
		r := httptest.NewRequest(http.MethodGet, "/api/x/mihomo/traffic", nil)
		r.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
		if !useMihomoWebSocketTunnel("traffic", r) {
			t.Fatal("expected tunnel for traffic+Sec-WebSocket-Key")
		}
	})

	t.Run("memory with key only", func(t *testing.T) {
		r := httptest.NewRequest(http.MethodGet, "/api/x/mihomo/memory", nil)
		r.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
		if !useMihomoWebSocketTunnel("memory", r) {
			t.Fatal("expected tunnel for memory+Sec-WebSocket-Key")
		}
	})

	t.Run("traffic streaming GET without key", func(t *testing.T) {
		r := httptest.NewRequest(http.MethodGet, "/api/x/mihomo/traffic", nil)
		if useMihomoWebSocketTunnel("traffic", r) {
			t.Fatal("plain GET /traffic must use HTTP forwarder")
		}
	})

	t.Run("proxies path never tunnel", func(t *testing.T) {
		r := httptest.NewRequest(http.MethodGet, "/api/x/mihomo/proxies", nil)
		r.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
		if useMihomoWebSocketTunnel("proxies", r) {
			t.Fatal("proxies must not use WS tunnel")
		}
	})

	t.Run("connections with key only", func(t *testing.T) {
		r := httptest.NewRequest(http.MethodGet, "/api/x/mihomo/connections", nil)
		r.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
		if !useMihomoWebSocketTunnel("connections", r) {
			t.Fatal("expected tunnel for connections+Sec-WebSocket-Key")
		}
	})

	t.Run("connections REST GET without key", func(t *testing.T) {
		r := httptest.NewRequest(http.MethodGet, "/api/x/mihomo/connections", nil)
		if useMihomoWebSocketTunnel("connections", r) {
			t.Fatal("plain GET /connections must use HTTP forwarder")
		}
	})
}

func TestApplyMihomoWSDialAuth(t *testing.T) {
	t.Run("connections strips Bearer for token query", func(t *testing.T) {
		u, _ := url.Parse("http://127.0.0.1:9090/connections")
		hdr := applyMihomoWSDialAuth("connections", u, "Bearer mysecret")
		if hdr.Get("Authorization") != "" {
			t.Fatal("connections WS must not set Authorization")
		}
		if u.RawQuery != "token=mysecret" {
			t.Fatalf("query: %q", u.RawQuery)
		}
	})

	t.Run("connections raw token", func(t *testing.T) {
		u, _ := url.Parse("http://127.0.0.1:9090/connections")
		_ = applyMihomoWSDialAuth("connections", u, "rawonly")
		if u.RawQuery != "token=rawonly" {
			t.Fatalf("query: %q", u.RawQuery)
		}
	})

	t.Run("traffic uses Authorization", func(t *testing.T) {
		u, _ := url.Parse("http://127.0.0.1:9090/traffic")
		hdr := applyMihomoWSDialAuth("traffic", u, "Bearer t")
		if hdr.Get("Authorization") != "Bearer t" {
			t.Fatalf("Authorization: %q", hdr.Get("Authorization"))
		}
	})
}

func TestMihomoWebSocketURL(t *testing.T) {
	u, err := url.Parse("http://172.20.0.2:9090/traffic?q=1")
	if err != nil {
		t.Fatal(err)
	}
	s, err := mihomoWebSocketURL(u)
	if err != nil {
		t.Fatal(err)
	}
	if want := "ws://172.20.0.2:9090/traffic?q=1"; s != want {
		t.Fatalf("got %q want %q", s, want)
	}
	u2, _ := url.Parse("https://example.com/mem")
	s2, err := mihomoWebSocketURL(u2)
	if err != nil {
		t.Fatal(err)
	}
	if want := "wss://example.com/mem"; s2 != want {
		t.Fatalf("got %q want %q", s2, want)
	}
}