feat(agentcompat): expose dashboard capability routes

Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
2026-09-19 17:50:12 +00:00 · 2026-07-20 04:30:50 +00:00
parent 0c6eb416a5
commit 2640b86d3a
42 changed files with 4083 additions and 271 deletions
@@ -0,0 +1,93 @@
+package rpc
+
+import (
+	"io"
+	"net/http"
+	"net/url"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/goccy/go-json"
+	"github.com/nezhahq/nezha/model"
+	"github.com/nezhahq/nezha/proto"
+)
+
+func TestServeNATClosesHijackedResourcesWhenStreamDisappearsBeforeUserTransfer(t *testing.T) {
+	fixture := newServeNATFixture(t)
+	conn := newServeNATConn()
+	body := &serveNATBody{}
+	writer := &serveNATResponseWriter{conn: conn}
+	request := &http.Request{Method: http.MethodPost, URL: &url.URL{Scheme: "http", Host: "example.test", Path: "/nat"}, Header: make(http.Header), Body: body, ContentLength: 0}
+	request.Header.Set("X-NAT-Test", "failure")
+	fixture.taskStream.onSend = func(task *proto.Task) error {
+		var nat model.TaskNAT
+		if err := json.Unmarshal([]byte(task.Data), &nat); err != nil {
+			return err
+		}
+		return fixture.handler.CloseStream(nat.StreamID)
+	}
+
+	ServeNAT(writer, request, &model.NAT{ServerID: fixture.server.ID, Host: "target.example"})
+
+	if got := body.closeCount.Load(); got == 0 {
+		t.Fatal("request body was not closed after failed transfer")
+	}
+	if got := conn.closeCount.Load(); got != 1 {
+		t.Fatalf("hijacked connection close count = %d, want 1", got)
+	}
+	select {
+	case <-conn.readDone:
+	case <-time.After(time.Second):
+		t.Fatal("hijacked connection read remained blocked after failed transfer")
+	}
+	if writer.status != 0 || writer.writes != 0 {
+		t.Fatalf("hijacked failure wrote HTTP response: status=%d writes=%d", writer.status, writer.writes)
+	}
+}
+
+func TestServeNATTransfersRequestAndRegistryOwnsSuccessfulCleanup(t *testing.T) {
+	fixture := newServeNATFixture(t)
+	conn := newServeNATConn()
+	body := &serveNATBody{}
+	agent := &serveNATAgent{readErr: io.EOF, writeDone: make(chan struct{})}
+	writer := &serveNATResponseWriter{conn: conn}
+	request := &http.Request{Method: http.MethodPost, URL: &url.URL{Scheme: "http", Host: "example.test", Path: "/nat"}, Header: make(http.Header), Body: body, ContentLength: 0}
+	request.Header.Set("X-NAT-Test", "success")
+	fixture.taskStream.onSend = func(task *proto.Task) error {
+		var nat model.TaskNAT
+		if err := json.Unmarshal([]byte(task.Data), &nat); err != nil {
+			return err
+		}
+		return fixture.handler.AgentConnected(nat.StreamID, agent)
+	}
+
+	serveDone := make(chan struct{})
+	go func() {
+		ServeNAT(writer, request, &model.NAT{ServerID: fixture.server.ID, Host: "target.example"})
+		close(serveDone)
+	}()
+
+	select {
+	case <-agent.writeDone:
+	case <-time.After(time.Second):
+		t.Fatal("agent did not receive the transferred NAT request")
+	}
+	select {
+	case <-serveDone:
+	case <-time.After(time.Second):
+		t.Fatal("ServeNAT did not finish after the successful stream was closed")
+	}
+	if got := body.closeCount.Load(); got == 0 {
+		t.Fatal("registry-owned request body was not closed")
+	}
+	if got := conn.closeCount.Load(); got != 1 {
+		t.Fatalf("registry-owned hijacked connection close count = %d, want 1", got)
+	}
+	if got := agent.closeCount.Load(); got != 1 {
+		t.Fatalf("registry-owned agent close count = %d, want 1", got)
+	}
+	if got := strings.ToLower(string(agent.writtenBytes())); !strings.Contains(got, "post /nat http/1.1") || !strings.Contains(got, "x-nat-test: success") {
+		t.Fatalf("agent received incomplete NAT request: %q", got)
+	}
+}