From 9c9701ed7e052a147aed3d9ce95b6659b9f6f7e8 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 14 Mar 2026 21:51:32 +0800 Subject: [PATCH] feat(tunnel): refactor VNC and proxy handling with structured routing - Updated VNC and proxy handling in the tunnel to utilize a structured routing approach, enhancing clarity and maintainability. - Replaced direct port checks with a new `forwardRoute` struct to encapsulate routing information, including channel type, container ID, and port. - Modified request handling to streamline the forwarding process and improve error handling for unknown routes. - Enhanced tests to validate the new routing logic and ensure consistent behavior across VNC and proxy requests. Made-with: Cursor --- openapi/computer/computer.go | 4 +- openapi/sandbox/manage.go | 2 +- tai/api/register.go | 1 + tai/tunnel/forward.go | 148 ++++++++++++------ tai/tunnel/forward_test.go | 240 +++++++++++++---------------- tai/tunnel/grpc_handler.go | 82 ++++++++-- tai/tunnel/grpc_handler_test.go | 16 +- tai/tunnel/taipb/tunnel.pb.go | 126 +++++++++------ tai/tunnel/taipb/tunnel_grpc.pb.go | 4 +- tai/types/types.go | 1 + 10 files changed, 379 insertions(+), 245 deletions(-) diff --git a/openapi/computer/computer.go b/openapi/computer/computer.go index 93016184..09311000 100644 --- a/openapi/computer/computer.go +++ b/openapi/computer/computer.go @@ -203,7 +203,7 @@ func nodeToHostOption(s taitypes.NodeMeta) computerOption { Status: status, Mode: s.Mode, Addr: addr, - VNC: s.Ports.VNC > 0, + VNC: s.Capabilities.VNC, System: computerSystemInfo{ OS: s.System.OS, Arch: s.System.Arch, @@ -239,7 +239,7 @@ func nodeToNodeOption(s taitypes.NodeMeta) computerOption { Status: status, Mode: s.Mode, Addr: addr, - VNC: s.Ports.VNC > 0, + VNC: s.Capabilities.VNC, System: computerSystemInfo{ OS: s.System.OS, Arch: s.System.Arch, diff --git a/openapi/sandbox/manage.go b/openapi/sandbox/manage.go index 316d1c58..3df56482 100644 --- a/openapi/sandbox/manage.go +++ b/openapi/sandbox/manage.go @@ -197,7 +197,7 @@ func hostToResponse(s taitypes.NodeMeta) sandboxResponse { Policy: "persistent", Mode: s.Mode, Addr: addr, - VNC: s.Ports.VNC > 0, + VNC: s.Capabilities.VNC, CreatedAt: s.ConnectedAt, LastActive: s.LastPing, System: sandboxSystemInfo{ diff --git a/tai/api/register.go b/tai/api/register.go index 27378169..fa6a6d59 100644 --- a/tai/api/register.go +++ b/tai/api/register.go @@ -294,6 +294,7 @@ func capsFromMap(m map[string]bool) types.Capabilities { Docker: m["docker"], K8s: m["k8s"], HostExec: m["host_exec"], + VNC: m["vnc"], } } diff --git a/tai/tunnel/forward.go b/tai/tunnel/forward.go index 81ee5896..4bd46302 100644 --- a/tai/tunnel/forward.go +++ b/tai/tunnel/forward.go @@ -5,6 +5,7 @@ import ( "fmt" "io" "net/http" + "strconv" "strings" "github.com/gin-gonic/gin" @@ -12,11 +13,24 @@ import ( "github.com/yaoapp/yao/tai/types" ) +const defaultVNCPort = 5900 + +// forwardRoute holds the structured routing information extracted from the +// incoming request URL. It is passed to RequestForward so that Yao can +// populate the TunnelControl proto fields and Tai can route directly without +// parsing the first packet. +type forwardRoute struct { + channelType string // "proxy" | "vnc" + containerID string // target container or "__host__" + containerPort int // container-internal port (vnc default 5900) + subpath string // rewritten request path for the container +} + // HandleForward handles HTTP/VNC/any TCP-level forwarding through the gRPC tunnel. // Route: ANY /tai/:taiID/proxy/*path and GET /tai/:taiID/vnc/*path // // It hijacks the browser's raw TCP connection, asks Tai to open a Forward stream -// to the resolved target port, rewrites the request path, and then performs +// with explicit routing information, rewrites the request path, and then performs // bidirectional byte-level bridging. No protocol parsing beyond HTTP hijack. func (h *TunnelHandler) HandleForward(c *gin.Context) { logger := h.logger @@ -30,16 +44,18 @@ func (h *TunnelHandler) HandleForward(c *gin.Context) { return } - targetPort := resolveTargetPort(c, node) - if targetPort == 0 { - c.JSON(http.StatusBadRequest, gin.H{"error": "cannot resolve target port"}) + route, err := resolveRoute(c, node) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) return } - rewrittenReq := rewriteRequest(c.Request, taiID) - logger.Debug("[forward] "+node.Mode+" → tai:"+fmt.Sprintf("%d", targetPort), + rewrittenReq := rewriteRequest(c.Request, taiID, route) + logger.Debug("[forward] "+node.Mode+" → tai", "tai_id", taiID, - "addr", node.Addr, + "type", route.channelType, + "container", route.containerID, + "container_port", route.containerPort, "path", rewrittenReq.URL.Path, ) @@ -55,10 +71,10 @@ func (h *TunnelHandler) HandleForward(c *gin.Context) { } defer browserConn.Close() - fwd, err := h.RequestForward(taiID, targetPort) + fwd, err := h.RequestForward(taiID, route) if err != nil { logger.Error("[forward] stream failed", - "tai_id", taiID, "port", targetPort, "err", err) + "tai_id", taiID, "type", route.channelType, "err", err) browserConn.Write([]byte("HTTP/1.1 502 Bad Gateway\r\n\r\n")) return } @@ -94,47 +110,91 @@ func HandleForwardLazy(c *gin.Context) { h.HandleForward(c) } -// resolveTargetPort determines the Tai-side port from the route pattern. -func resolveTargetPort(c *gin.Context, node *types.NodeMeta) int { +// resolveRoute extracts structured routing info from the request URL path. +// +// For proxy requests (/tai/:taiID/proxy/{containerID}:{port}/{subpath}): +// +// channelType = "proxy", containerPort from URL, subpath = remaining path. +// +// For VNC requests (/tai/:taiID/vnc/{containerID}/ws): +// +// channelType = "vnc", containerPort = 5900, subpath = /vnc/{containerID}/ws. +func resolveRoute(c *gin.Context, node *types.NodeMeta) (*forwardRoute, error) { path := c.Request.URL.Path - - if strings.Contains(path, "/vnc/") { - if node.Ports.VNC != 0 { - return node.Ports.VNC - } - return 16080 - } - if strings.Contains(path, "/proxy/") { - if node.Ports.HTTP != 0 { - return node.Ports.HTTP - } - return 8099 - } - return 0 -} - -// rewriteRequest clones the request and strips the Yao-side route prefix, -// leaving only what the Tai-side handler expects. -// -// The Tai httpproxy expects /{containerID}:{port}/..., so the /proxy prefix -// is stripped. The Tai VNC router expects /vnc/{containerID}/ws, so the /vnc -// prefix is kept. -// -// /v1/tai/abc/proxy/cid:8080/foo → /cid:8080/foo -// /v1/tai/abc/vnc/cid/ws → /vnc/cid/ws -func rewriteRequest(orig *http.Request, taiID string) *http.Request { - r := orig.Clone(orig.Context()) + taiID := c.Param("taiID") marker := "/tai/" + taiID - if idx := strings.Index(r.URL.Path, marker); idx >= 0 { - rest := r.URL.Path[idx+len(marker):] - rest = strings.TrimPrefix(rest, "/proxy") - if rest == "" { - rest = "/" + idx := strings.Index(path, marker) + if idx < 0 { + return nil, fmt.Errorf("cannot locate /tai/%s in path", taiID) + } + rest := path[idx+len(marker):] + + if strings.HasPrefix(rest, "/vnc/") { + // /vnc/{containerID}/ws → containerID, port=5900 + tail := strings.TrimPrefix(rest, "/vnc/") + containerID := tail + if slashIdx := strings.IndexByte(tail, '/'); slashIdx >= 0 { + containerID = tail[:slashIdx] } - r.URL.Path = rest + if containerID == "" { + return nil, fmt.Errorf("missing container ID in VNC path: %s", path) + } + return &forwardRoute{ + channelType: "vnc", + containerID: containerID, + containerPort: defaultVNCPort, + subpath: rest, // keep /vnc/{containerID}/ws + }, nil } + if strings.HasPrefix(rest, "/proxy/") { + // /proxy/{containerID}:{port}/{subpath} + proxyPath := strings.TrimPrefix(rest, "/proxy") + // proxyPath = /{containerID}:{port}/{subpath} + proxyPath = strings.TrimPrefix(proxyPath, "/") + if proxyPath == "" { + return nil, fmt.Errorf("empty proxy path") + } + + slash := strings.IndexByte(proxyPath, '/') + var head, subpath string + if slash == -1 { + head = proxyPath + subpath = "/" + } else { + head = proxyPath[:slash] + subpath = proxyPath[slash:] + } + + colon := strings.LastIndexByte(head, ':') + if colon < 0 { + return nil, fmt.Errorf("missing port in proxy path: %s", path) + } + containerID := head[:colon] + portStr := head[colon+1:] + port, err := strconv.Atoi(portStr) + if err != nil { + return nil, fmt.Errorf("invalid port %q in proxy path: %w", portStr, err) + } + return &forwardRoute{ + channelType: "proxy", + containerID: containerID, + containerPort: port, + subpath: subpath, + }, nil + } + + return nil, fmt.Errorf("unknown route pattern: %s", rest) +} + +// rewriteRequest clones the request and sets the path to the route's subpath. +// +// For proxy: the path becomes the subpath (e.g. /foo/bar). +// For VNC: the path keeps /vnc/{containerID}/ws as-is. +func rewriteRequest(orig *http.Request, taiID string, route *forwardRoute) *http.Request { + r := orig.Clone(orig.Context()) + r.URL.Path = route.subpath r.RequestURI = r.URL.RequestURI() return r } diff --git a/tai/tunnel/forward_test.go b/tai/tunnel/forward_test.go index 703e5198..834c33ff 100644 --- a/tai/tunnel/forward_test.go +++ b/tai/tunnel/forward_test.go @@ -16,155 +16,131 @@ func init() { gin.SetMode(gin.TestMode) } -func TestResolveTargetPort_VNC(t *testing.T) { +func TestResolveRoute_Proxy(t *testing.T) { tests := []struct { - name string - path string - vncPort int - wantPort int + name string + path string + wantType string + wantContainer string + wantPort int + wantSubpath string }{ - {"default_vnc", "/tai/abc/vnc/websockify", 0, 16080}, - {"custom_vnc", "/tai/abc/vnc/websockify", 5900, 5900}, + { + "basic_proxy", + "/tai/abc/proxy/cid123:8080/foo/bar", + "proxy", "cid123", 8080, "/foo/bar", + }, + { + "proxy_root", + "/tai/abc/proxy/cid:3000", + "proxy", "cid", 3000, "/", + }, + { + "proxy_host", + "/v1/tai/abc/proxy/__host__:9090/api", + "proxy", "__host__", 9090, "/api", + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = &http.Request{URL: &url.URL{Path: tt.path}} - node := &types.NodeMeta{Ports: types.Ports{VNC: tt.vncPort}} - got := resolveTargetPort(c, node) - if got != tt.wantPort { - t.Errorf("resolveTargetPort = %d, want %d", got, tt.wantPort) + c.Params = gin.Params{{Key: "taiID", Value: "abc"}} + node := &types.NodeMeta{} + + r, err := resolveRoute(c, node) + if err != nil { + t.Fatalf("resolveRoute error: %v", err) + } + if r.channelType != tt.wantType { + t.Errorf("channelType = %q, want %q", r.channelType, tt.wantType) + } + if r.containerID != tt.wantContainer { + t.Errorf("containerID = %q, want %q", r.containerID, tt.wantContainer) + } + if r.containerPort != tt.wantPort { + t.Errorf("containerPort = %d, want %d", r.containerPort, tt.wantPort) + } + if r.subpath != tt.wantSubpath { + t.Errorf("subpath = %q, want %q", r.subpath, tt.wantSubpath) } }) } } -func TestResolveTargetPort_Proxy(t *testing.T) { +func TestResolveRoute_VNC(t *testing.T) { tests := []struct { - name string - path string - httpPort int - wantPort int + name string + path string + wantContainer string + wantPort int }{ - {"default_proxy", "/tai/abc/proxy/api/v1/foo", 0, 8099}, - {"custom_proxy", "/tai/abc/proxy/api/v1/foo", 9090, 9090}, + {"vnc_basic", "/tai/abc/vnc/container1/ws", "container1", defaultVNCPort}, + {"vnc_host", "/v1/tai/abc/vnc/__host__/ws", "__host__", defaultVNCPort}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = &http.Request{URL: &url.URL{Path: tt.path}} - node := &types.NodeMeta{Ports: types.Ports{HTTP: tt.httpPort}} - got := resolveTargetPort(c, node) - if got != tt.wantPort { - t.Errorf("resolveTargetPort = %d, want %d", got, tt.wantPort) + c.Params = gin.Params{{Key: "taiID", Value: "abc"}} + node := &types.NodeMeta{} + + r, err := resolveRoute(c, node) + if err != nil { + t.Fatalf("resolveRoute error: %v", err) + } + if r.channelType != "vnc" { + t.Errorf("channelType = %q, want vnc", r.channelType) + } + if r.containerID != tt.wantContainer { + t.Errorf("containerID = %q, want %q", r.containerID, tt.wantContainer) + } + if r.containerPort != tt.wantPort { + t.Errorf("containerPort = %d, want %d", r.containerPort, tt.wantPort) } }) } } -func TestResolveTargetPort_Unknown(t *testing.T) { +func TestResolveRoute_Unknown(t *testing.T) { c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = &http.Request{URL: &url.URL{Path: "/tai/abc/unknown/something"}} + c.Params = gin.Params{{Key: "taiID", Value: "abc"}} node := &types.NodeMeta{} - got := resolveTargetPort(c, node) - if got != 0 { - t.Errorf("resolveTargetPort = %d, want 0", got) + + _, err := resolveRoute(c, node) + if err == nil { + t.Error("expected error for unknown route") } } -func TestRewriteRequest(t *testing.T) { - tests := []struct { - name string - origPath string - taiID string - wantPath string - wantURI string - }{ - { - "proxy_path", - "/tai/abc123/proxy/api/v1/data", - "abc123", - "/api/v1/data", - "/api/v1/data", - }, - { - "vnc_path", - "/tai/node-1/vnc/websockify", - "node-1", - "/vnc/websockify", - "/vnc/websockify", - }, - { - "with_query", - "/tai/node-1/proxy/api?foo=bar", - "node-1", - "/api", - "/api?foo=bar", - }, - { - "exact_prefix", - "/tai/node-1", - "node-1", - "/", - "/", - }, - { - "with_base_url", - "/v1/tai/node-1/proxy/api/v1/data", - "node-1", - "/api/v1/data", - "/api/v1/data", - }, - { - "with_base_url_vnc", - "/v1/tai/abc123/vnc/__host__/ws", - "abc123", - "/vnc/__host__/ws", - "/vnc/__host__/ws", - }, - { - "proxy_container_port", - "/v1/tai/abc/proxy/cid123:8080/foo", - "abc", - "/cid123:8080/foo", - "/cid123:8080/foo", - }, - { - "no_match", - "/other/path", - "node-1", - "/other/path", - "/other/path", - }, +func TestRewriteRequest_Proxy(t *testing.T) { + u, _ := url.Parse("http://localhost/v1/tai/abc/proxy/cid:8080/foo") + orig := &http.Request{ + Method: "GET", + URL: u, + RequestURI: u.RequestURI(), + Host: "localhost", + Header: http.Header{}, + } + route := &forwardRoute{ + channelType: "proxy", + containerID: "cid", + containerPort: 8080, + subpath: "/foo", } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - u, _ := url.Parse("http://localhost" + tt.origPath) - orig := &http.Request{ - Method: "GET", - URL: u, - RequestURI: u.RequestURI(), - Host: "localhost", - Header: http.Header{}, - } - got := rewriteRequest(orig, tt.taiID) - - if got.URL.Path != tt.wantPath { - t.Errorf("path = %q, want %q", got.URL.Path, tt.wantPath) - } - if got.RequestURI != tt.wantURI { - t.Errorf("requestURI = %q, want %q", got.RequestURI, tt.wantURI) - } - if got == orig { - t.Error("rewriteRequest should return a clone, not the original") - } - }) + got := rewriteRequest(orig, "abc", route) + if got.URL.Path != "/foo" { + t.Errorf("path = %q, want /foo", got.URL.Path) + } + if got == orig { + t.Error("rewriteRequest should return a clone") } } -func TestRewriteRequest_PreservesHeaders(t *testing.T) { - u, _ := url.Parse("http://localhost/tai/node-1/vnc/websockify") +func TestRewriteRequest_VNC(t *testing.T) { + u, _ := url.Parse("http://localhost/tai/node-1/vnc/cid/ws") orig := &http.Request{ Method: "GET", URL: u, @@ -175,14 +151,20 @@ func TestRewriteRequest_PreservesHeaders(t *testing.T) { "Upgrade": {"websocket"}, }, } + route := &forwardRoute{ + channelType: "vnc", + containerID: "cid", + containerPort: 5900, + subpath: "/vnc/cid/ws", + } - got := rewriteRequest(orig, "node-1") + got := rewriteRequest(orig, "node-1", route) + if got.URL.Path != "/vnc/cid/ws" { + t.Errorf("path = %q, want /vnc/cid/ws", got.URL.Path) + } if got.Header.Get("Connection") != "Upgrade" { t.Error("expected Connection header preserved") } - if got.Header.Get("Upgrade") != "websocket" { - t.Error("expected Upgrade header preserved") - } } func TestHandleForwardLazy_NilHandler(t *testing.T) { @@ -217,29 +199,24 @@ func TestHandleForward_NodeNotFound(t *testing.T) { } } -func TestHandleForward_NodeOffline(t *testing.T) { +func TestHandleForward_UnknownRoute(t *testing.T) { reg := registry.NewForTest() h := NewTunnelHandler(reg) reg.Register(®istry.TaiNode{ - TaiID: "offline-node", + TaiID: "online-node", Mode: "tunnel", - Ports: types.Ports{HTTP: 8099}, }) - // Manually set status to offline via a Get() — the node is online by default - // after Register, but we need an offline one. We'll use Unregister + re-register - // pattern. Actually, let's just test with a node that doesn't exist: - // the NodeNotFound test above covers that case. Instead, test zero port. w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) - c.Request = httptest.NewRequest("GET", "/tai/offline-node/unknown/foo", nil) - c.Params = gin.Params{{Key: "taiID", Value: "offline-node"}} + c.Request = httptest.NewRequest("GET", "/tai/online-node/unknown/foo", nil) + c.Params = gin.Params{{Key: "taiID", Value: "online-node"}} h.HandleForward(c) if w.Code != http.StatusBadRequest { - t.Errorf("expected 400 for unresolvable port, got %d", w.Code) + t.Errorf("expected 400 for unresolvable route, got %d", w.Code) } } @@ -268,7 +245,6 @@ func TestHandleForward_ViaRealHTTP(t *testing.T) { reg.Register(®istry.TaiNode{ TaiID: "http-node", Mode: "tunnel", - Ports: types.Ports{HTTP: 8099}, }) router := gin.New() @@ -277,16 +253,12 @@ func TestHandleForward_ViaRealHTTP(t *testing.T) { srv := httptest.NewServer(router) defer srv.Close() - resp, err := http.Get(srv.URL + "/tai/http-node/proxy/api") + resp, err := http.Get(srv.URL + "/tai/http-node/proxy/cid:8080/api") if err != nil { t.Fatal(err) } defer resp.Body.Close() - // RequestForward will fail (no register stream) → hijacked conn gets "502" - // or the response will be a 502 written before hijack. - // Since hijack happens, the actual HTTP status may not be set normally. - // We just verify no panic and the request completes. if resp.StatusCode == 200 { t.Error("expected non-200 response for failed forward") } diff --git a/tai/tunnel/grpc_handler.go b/tai/tunnel/grpc_handler.go index bdf39deb..d7fb41c7 100644 --- a/tai/tunnel/grpc_handler.go +++ b/tai/tunnel/grpc_handler.go @@ -191,7 +191,9 @@ func (h *TunnelHandler) Forward(stream taipb.TaiTunnel_ForwardServer) error { // RequestForward sends an "open" command to Tai via the Register stream and // waits for Tai to call back with a Forward stream. Returns the Forward stream. -func (h *TunnelHandler) RequestForward(taiID string, targetPort int) (taipb.TaiTunnel_ForwardServer, error) { +// +// route may be nil for raw TCP tunnels (gRPC, Docker API, K8s API). +func (h *TunnelHandler) RequestForward(taiID string, route *forwardRoute) (taipb.TaiTunnel_ForwardServer, error) { stream := h.reg.GetRegisterStream(taiID) if stream == nil { return nil, fmt.Errorf("tai %s: no active register stream", taiID) @@ -217,23 +219,30 @@ func (h *TunnelHandler) RequestForward(taiID string, targetPort int) (taipb.TaiT return nil, fmt.Errorf("tai %s: register stream type mismatch", taiID) } + ctrl := &taipb.TunnelControl{ + Type: "open", + ChannelId: channelID, + } + if route != nil { + ctrl.ChannelType = route.channelType + ctrl.ContainerId = route.containerID + ctrl.ContainerPort = int32(route.containerPort) + } + short := registry.ShortChannelID(channelID) h.logger.Debug("[forward] sending open command", - "tai_id", taiID, "port", targetPort, "channel_id", short) + "tai_id", taiID, "channel_type", ctrl.ChannelType, + "container", ctrl.ContainerId, "channel_id", short) mu.Lock() - sendErr := regStream.Send(&taipb.TunnelControl{ - Type: "open", - ChannelId: channelID, - TargetPort: int32(targetPort), - }) + sendErr := regStream.Send(ctrl) mu.Unlock() if sendErr != nil { return nil, fmt.Errorf("send open: %w", sendErr) } h.logger.Debug("[forward] open sent, waiting for callback", - "tai_id", taiID, "port", targetPort, "channel_id", short) + "tai_id", taiID, "channel_id", short) select { case fwd := <-waitCh: @@ -261,8 +270,9 @@ func (h *TunnelHandler) connectTunnelNode(taiID string) { // bridgeConn bridges a local TCP connection to a Tai port via gRPC Forward stream. // Called by registry.OpenLocalListener for each accepted TCP connection. +// Uses raw TCP forwarding (TargetPort only, no container routing). func (h *TunnelHandler) bridgeConn(taiID string, targetPort int, localConn net.Conn) { - fwd, err := h.RequestForward(taiID, targetPort) + fwd, err := h.requestForwardRaw(taiID, targetPort) if err != nil { localConn.Close() h.logger.Error("request forward failed", @@ -274,6 +284,59 @@ func (h *TunnelHandler) bridgeConn(taiID string, targetPort int, localConn net.C bridgeTCP(localConn, streamConn) } +// requestForwardRaw sends an "open" command with only TargetPort (no container +// routing). Used by bridgeConn for raw TCP tunnels (gRPC, Docker API, K8s API). +func (h *TunnelHandler) requestForwardRaw(taiID string, targetPort int) (taipb.TaiTunnel_ForwardServer, error) { + stream := h.reg.GetRegisterStream(taiID) + if stream == nil { + return nil, fmt.Errorf("tai %s: no active register stream", taiID) + } + + muVal, ok := h.sendMu.Load(taiID) + if !ok { + return nil, fmt.Errorf("tai %s: no send mutex (stream closing?)", taiID) + } + mu := muVal.(*sync.Mutex) + + channelID, err := registry.GenerateChannelID() + if err != nil { + return nil, fmt.Errorf("generate channel_id: %w", err) + } + + waitCh := make(chan taipb.TaiTunnel_ForwardServer, 1) + h.pending.Store(channelID, waitCh) + defer h.pending.Delete(channelID) + + regStream, ok := stream.(taipb.TaiTunnel_RegisterServer) + if !ok { + return nil, fmt.Errorf("tai %s: register stream type mismatch", taiID) + } + + short := registry.ShortChannelID(channelID) + h.logger.Debug("[forward] sending open command (raw)", + "tai_id", taiID, "port", targetPort, "channel_id", short) + + mu.Lock() + sendErr := regStream.Send(&taipb.TunnelControl{ + Type: "open", + ChannelId: channelID, + TargetPort: int32(targetPort), + }) + mu.Unlock() + if sendErr != nil { + return nil, fmt.Errorf("send open: %w", sendErr) + } + + select { + case fwd := <-waitCh: + return fwd, nil + case <-time.After(10 * time.Second): + return nil, fmt.Errorf("tai %s: forward timeout (10s) channel=%s", taiID, short) + case <-regStream.Context().Done(): + return nil, fmt.Errorf("tai %s: register stream closed while waiting for forward", taiID) + } +} + // forwardConn wraps a Forward stream as a net.Conn-like reader/writer. type forwardConn struct { stream taipb.TaiTunnel_ForwardServer @@ -364,6 +427,7 @@ func capsFromProto(c *taipb.Capabilities) types.Capabilities { Docker: c.Docker, K8s: c.K8S, HostExec: c.HostExec, + VNC: c.Vnc, } } diff --git a/tai/tunnel/grpc_handler_test.go b/tai/tunnel/grpc_handler_test.go index 278dd300..b5d97558 100644 --- a/tai/tunnel/grpc_handler_test.go +++ b/tai/tunnel/grpc_handler_test.go @@ -245,7 +245,7 @@ func TestRequestForward_NoRegisterStream(t *testing.T) { reg.Register(®istry.TaiNode{TaiID: "no-stream", Mode: "tunnel"}) - _, err := h.RequestForward("no-stream", 8099) + _, err := h.requestForwardRaw("no-stream", 8099) if err == nil { t.Fatal("expected error when no register stream") } @@ -258,7 +258,7 @@ func TestRequestForward_TypeMismatch(t *testing.T) { reg.Register(®istry.TaiNode{TaiID: "bad-type", Mode: "tunnel"}) reg.SetRegisterStream("bad-type", "not-a-stream") - _, err := h.RequestForward("bad-type", 8099) + _, err := h.requestForwardRaw("bad-type", 8099) if err == nil { t.Fatal("expected error for type mismatch") } @@ -348,7 +348,7 @@ drainLoop: requestDone.Add(1) go func() { defer requestDone.Done() - requestResult, requestErr = h.RequestForward(taiID, 8099) + requestResult, requestErr = h.requestForwardRaw(taiID, 8099) }() // Receive the "open" command @@ -748,7 +748,7 @@ func TestRequestForward_Timeout(t *testing.T) { // by never sending Forward). We'll use a short context cancel to avoid waiting. done := make(chan error, 1) go func() { - _, err := h.RequestForward(taiID, 8099) + _, err := h.requestForwardRaw(taiID, 8099) done <- err }() @@ -828,7 +828,7 @@ drained: for i := 0; i < N; i++ { port := 8099 + i go func(port int) { - _, err := h.RequestForward(taiID, port) + _, err := h.requestForwardRaw(taiID, port) results <- err }(port) } @@ -926,7 +926,7 @@ drained2: // Start RequestForward fwdResult := make(chan error, 1) go func() { - _, err := h.RequestForward(taiID, 8099) + _, err := h.requestForwardRaw(taiID, 8099) fwdResult <- err }() @@ -1101,7 +1101,7 @@ proxyDrained: }() // Now do an actual RequestForward + simulate browser side - fwd, err := h.RequestForward(taiID, 8099) + fwd, err := h.requestForwardRaw(taiID, 8099) if err != nil { t.Fatal("RequestForward:", err) } @@ -1267,7 +1267,7 @@ vncDrained: }() // Send WS upgrade request through tunnel - fwd, err := h.RequestForward(taiID, 16080) + fwd, err := h.requestForwardRaw(taiID, 16080) if err != nil { t.Fatal("RequestForward:", err) } diff --git a/tai/tunnel/taipb/tunnel.pb.go b/tai/tunnel/taipb/tunnel.pb.go index 33fe5105..ddc6cd8d 100644 --- a/tai/tunnel/taipb/tunnel.pb.go +++ b/tai/tunnel/taipb/tunnel.pb.go @@ -2,7 +2,7 @@ // versions: // protoc-gen-go v1.36.11 // protoc v4.25.0 -// source: tunnel.proto +// source: tunnel/proto/tunnel.proto package taipb @@ -33,8 +33,11 @@ type TunnelControl struct { Caps *Capabilities `protobuf:"bytes,7,opt,name=caps,proto3" json:"caps,omitempty"` System *SystemInfo `protobuf:"bytes,8,opt,name=system,proto3" json:"system,omitempty"` // Carried on "open" (Yao → Tai) - ChannelId string `protobuf:"bytes,10,opt,name=channel_id,json=channelId,proto3" json:"channel_id,omitempty"` - TargetPort int32 `protobuf:"varint,11,opt,name=target_port,json=targetPort,proto3" json:"target_port,omitempty"` + ChannelId string `protobuf:"bytes,10,opt,name=channel_id,json=channelId,proto3" json:"channel_id,omitempty"` + TargetPort int32 `protobuf:"varint,11,opt,name=target_port,json=targetPort,proto3" json:"target_port,omitempty"` + ChannelType string `protobuf:"bytes,12,opt,name=channel_type,json=channelType,proto3" json:"channel_type,omitempty"` // "proxy" | "vnc" | "" (legacy/raw TCP) + ContainerId string `protobuf:"bytes,13,opt,name=container_id,json=containerId,proto3" json:"container_id,omitempty"` // target container or "__host__" + ContainerPort int32 `protobuf:"varint,14,opt,name=container_port,json=containerPort,proto3" json:"container_port,omitempty"` // container-internal port (vnc default 5900) // Carried on "registered" (Yao → Tai) TaiId string `protobuf:"bytes,20,opt,name=tai_id,json=taiId,proto3" json:"tai_id,omitempty"` unknownFields protoimpl.UnknownFields @@ -43,7 +46,7 @@ type TunnelControl struct { func (x *TunnelControl) Reset() { *x = TunnelControl{} - mi := &file_tunnel_proto_msgTypes[0] + mi := &file_tunnel_proto_tunnel_proto_msgTypes[0] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -55,7 +58,7 @@ func (x *TunnelControl) String() string { func (*TunnelControl) ProtoMessage() {} func (x *TunnelControl) ProtoReflect() protoreflect.Message { - mi := &file_tunnel_proto_msgTypes[0] + mi := &file_tunnel_proto_tunnel_proto_msgTypes[0] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -68,7 +71,7 @@ func (x *TunnelControl) ProtoReflect() protoreflect.Message { // Deprecated: Use TunnelControl.ProtoReflect.Descriptor instead. func (*TunnelControl) Descriptor() ([]byte, []int) { - return file_tunnel_proto_rawDescGZIP(), []int{0} + return file_tunnel_proto_tunnel_proto_rawDescGZIP(), []int{0} } func (x *TunnelControl) GetType() string { @@ -141,6 +144,27 @@ func (x *TunnelControl) GetTargetPort() int32 { return 0 } +func (x *TunnelControl) GetChannelType() string { + if x != nil { + return x.ChannelType + } + return "" +} + +func (x *TunnelControl) GetContainerId() string { + if x != nil { + return x.ContainerId + } + return "" +} + +func (x *TunnelControl) GetContainerPort() int32 { + if x != nil { + return x.ContainerPort + } + return 0 +} + func (x *TunnelControl) GetTaiId() string { if x != nil { return x.TaiId @@ -157,7 +181,7 @@ type ForwardData struct { func (x *ForwardData) Reset() { *x = ForwardData{} - mi := &file_tunnel_proto_msgTypes[1] + mi := &file_tunnel_proto_tunnel_proto_msgTypes[1] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -169,7 +193,7 @@ func (x *ForwardData) String() string { func (*ForwardData) ProtoMessage() {} func (x *ForwardData) ProtoReflect() protoreflect.Message { - mi := &file_tunnel_proto_msgTypes[1] + mi := &file_tunnel_proto_tunnel_proto_msgTypes[1] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -182,7 +206,7 @@ func (x *ForwardData) ProtoReflect() protoreflect.Message { // Deprecated: Use ForwardData.ProtoReflect.Descriptor instead. func (*ForwardData) Descriptor() ([]byte, []int) { - return file_tunnel_proto_rawDescGZIP(), []int{1} + return file_tunnel_proto_tunnel_proto_rawDescGZIP(), []int{1} } func (x *ForwardData) GetData() []byte { @@ -205,7 +229,7 @@ type Ports struct { func (x *Ports) Reset() { *x = Ports{} - mi := &file_tunnel_proto_msgTypes[2] + mi := &file_tunnel_proto_tunnel_proto_msgTypes[2] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -217,7 +241,7 @@ func (x *Ports) String() string { func (*Ports) ProtoMessage() {} func (x *Ports) ProtoReflect() protoreflect.Message { - mi := &file_tunnel_proto_msgTypes[2] + mi := &file_tunnel_proto_tunnel_proto_msgTypes[2] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -230,7 +254,7 @@ func (x *Ports) ProtoReflect() protoreflect.Message { // Deprecated: Use Ports.ProtoReflect.Descriptor instead. func (*Ports) Descriptor() ([]byte, []int) { - return file_tunnel_proto_rawDescGZIP(), []int{2} + return file_tunnel_proto_tunnel_proto_rawDescGZIP(), []int{2} } func (x *Ports) GetGrpc() int32 { @@ -273,13 +297,14 @@ type Capabilities struct { Docker bool `protobuf:"varint,1,opt,name=docker,proto3" json:"docker,omitempty"` K8S bool `protobuf:"varint,2,opt,name=k8s,proto3" json:"k8s,omitempty"` HostExec bool `protobuf:"varint,3,opt,name=host_exec,json=hostExec,proto3" json:"host_exec,omitempty"` + Vnc bool `protobuf:"varint,4,opt,name=vnc,proto3" json:"vnc,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } func (x *Capabilities) Reset() { *x = Capabilities{} - mi := &file_tunnel_proto_msgTypes[3] + mi := &file_tunnel_proto_tunnel_proto_msgTypes[3] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -291,7 +316,7 @@ func (x *Capabilities) String() string { func (*Capabilities) ProtoMessage() {} func (x *Capabilities) ProtoReflect() protoreflect.Message { - mi := &file_tunnel_proto_msgTypes[3] + mi := &file_tunnel_proto_tunnel_proto_msgTypes[3] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -304,7 +329,7 @@ func (x *Capabilities) ProtoReflect() protoreflect.Message { // Deprecated: Use Capabilities.ProtoReflect.Descriptor instead. func (*Capabilities) Descriptor() ([]byte, []int) { - return file_tunnel_proto_rawDescGZIP(), []int{3} + return file_tunnel_proto_tunnel_proto_rawDescGZIP(), []int{3} } func (x *Capabilities) GetDocker() bool { @@ -328,6 +353,13 @@ func (x *Capabilities) GetHostExec() bool { return false } +func (x *Capabilities) GetVnc() bool { + if x != nil { + return x.Vnc + } + return false +} + type SystemInfo struct { state protoimpl.MessageState `protogen:"open.v1"` Os string `protobuf:"bytes,1,opt,name=os,proto3" json:"os,omitempty"` @@ -340,7 +372,7 @@ type SystemInfo struct { func (x *SystemInfo) Reset() { *x = SystemInfo{} - mi := &file_tunnel_proto_msgTypes[4] + mi := &file_tunnel_proto_tunnel_proto_msgTypes[4] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -352,7 +384,7 @@ func (x *SystemInfo) String() string { func (*SystemInfo) ProtoMessage() {} func (x *SystemInfo) ProtoReflect() protoreflect.Message { - mi := &file_tunnel_proto_msgTypes[4] + mi := &file_tunnel_proto_tunnel_proto_msgTypes[4] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -365,7 +397,7 @@ func (x *SystemInfo) ProtoReflect() protoreflect.Message { // Deprecated: Use SystemInfo.ProtoReflect.Descriptor instead. func (*SystemInfo) Descriptor() ([]byte, []int) { - return file_tunnel_proto_rawDescGZIP(), []int{4} + return file_tunnel_proto_tunnel_proto_rawDescGZIP(), []int{4} } func (x *SystemInfo) GetOs() string { @@ -396,12 +428,12 @@ func (x *SystemInfo) GetShell() string { return "" } -var File_tunnel_proto protoreflect.FileDescriptor +var File_tunnel_proto_tunnel_proto protoreflect.FileDescriptor -const file_tunnel_proto_rawDesc = "" + +const file_tunnel_proto_tunnel_proto_rawDesc = "" + "\n" + - "\ftunnel.proto\x12\n" + - "tai.tunnel\"\xf6\x02\n" + + "\x19tunnel/proto/tunnel.proto\x12\n" + + "tai.tunnel\"\xe3\x03\n" + "\rTunnelControl\x12\x12\n" + "\x04type\x18\x01 \x01(\tR\x04type\x12\x17\n" + "\anode_id\x18\x02 \x01(\tR\x06nodeId\x12\x1d\n" + @@ -416,7 +448,10 @@ const file_tunnel_proto_rawDesc = "" + "channel_id\x18\n" + " \x01(\tR\tchannelId\x12\x1f\n" + "\vtarget_port\x18\v \x01(\x05R\n" + - "targetPort\x12\x15\n" + + "targetPort\x12!\n" + + "\fchannel_type\x18\f \x01(\tR\vchannelType\x12!\n" + + "\fcontainer_id\x18\r \x01(\tR\vcontainerId\x12%\n" + + "\x0econtainer_port\x18\x0e \x01(\x05R\rcontainerPort\x12\x15\n" + "\x06tai_id\x18\x14 \x01(\tR\x05taiId\"!\n" + "\vForwardData\x12\x12\n" + "\x04data\x18\x01 \x01(\fR\x04data\"k\n" + @@ -425,11 +460,12 @@ const file_tunnel_proto_rawDesc = "" + "\x04http\x18\x02 \x01(\x05R\x04http\x12\x10\n" + "\x03vnc\x18\x03 \x01(\x05R\x03vnc\x12\x16\n" + "\x06docker\x18\x04 \x01(\x05R\x06docker\x12\x10\n" + - "\x03k8s\x18\x05 \x01(\x05R\x03k8s\"U\n" + + "\x03k8s\x18\x05 \x01(\x05R\x03k8s\"g\n" + "\fCapabilities\x12\x16\n" + "\x06docker\x18\x01 \x01(\bR\x06docker\x12\x10\n" + "\x03k8s\x18\x02 \x01(\bR\x03k8s\x12\x1b\n" + - "\thost_exec\x18\x03 \x01(\bR\bhostExec\"b\n" + + "\thost_exec\x18\x03 \x01(\bR\bhostExec\x12\x10\n" + + "\x03vnc\x18\x04 \x01(\bR\x03vnc\"b\n" + "\n" + "SystemInfo\x12\x0e\n" + "\x02os\x18\x01 \x01(\tR\x02os\x12\x12\n" + @@ -438,29 +474,29 @@ const file_tunnel_proto_rawDesc = "" + "\x05shell\x18\x04 \x01(\tR\x05shell2\x92\x01\n" + "\tTaiTunnel\x12D\n" + "\bRegister\x12\x19.tai.tunnel.TunnelControl\x1a\x19.tai.tunnel.TunnelControl(\x010\x01\x12?\n" + - "\aForward\x12\x17.tai.tunnel.ForwardData\x1a\x17.tai.tunnel.ForwardData(\x010\x01B(Z&github.com/yaoapp/yao/tai/tunnel/taipbb\x06proto3" + "\aForward\x12\x17.tai.tunnel.ForwardData\x1a\x17.tai.tunnel.ForwardData(\x010\x01B$Z\"github.com/yaoapp/tai/tunnel/taipbb\x06proto3" var ( - file_tunnel_proto_rawDescOnce sync.Once - file_tunnel_proto_rawDescData []byte + file_tunnel_proto_tunnel_proto_rawDescOnce sync.Once + file_tunnel_proto_tunnel_proto_rawDescData []byte ) -func file_tunnel_proto_rawDescGZIP() []byte { - file_tunnel_proto_rawDescOnce.Do(func() { - file_tunnel_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_tunnel_proto_rawDesc), len(file_tunnel_proto_rawDesc))) +func file_tunnel_proto_tunnel_proto_rawDescGZIP() []byte { + file_tunnel_proto_tunnel_proto_rawDescOnce.Do(func() { + file_tunnel_proto_tunnel_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_tunnel_proto_tunnel_proto_rawDesc), len(file_tunnel_proto_tunnel_proto_rawDesc))) }) - return file_tunnel_proto_rawDescData + return file_tunnel_proto_tunnel_proto_rawDescData } -var file_tunnel_proto_msgTypes = make([]protoimpl.MessageInfo, 5) -var file_tunnel_proto_goTypes = []any{ +var file_tunnel_proto_tunnel_proto_msgTypes = make([]protoimpl.MessageInfo, 5) +var file_tunnel_proto_tunnel_proto_goTypes = []any{ (*TunnelControl)(nil), // 0: tai.tunnel.TunnelControl (*ForwardData)(nil), // 1: tai.tunnel.ForwardData (*Ports)(nil), // 2: tai.tunnel.Ports (*Capabilities)(nil), // 3: tai.tunnel.Capabilities (*SystemInfo)(nil), // 4: tai.tunnel.SystemInfo } -var file_tunnel_proto_depIdxs = []int32{ +var file_tunnel_proto_tunnel_proto_depIdxs = []int32{ 2, // 0: tai.tunnel.TunnelControl.ports:type_name -> tai.tunnel.Ports 3, // 1: tai.tunnel.TunnelControl.caps:type_name -> tai.tunnel.Capabilities 4, // 2: tai.tunnel.TunnelControl.system:type_name -> tai.tunnel.SystemInfo @@ -475,26 +511,26 @@ var file_tunnel_proto_depIdxs = []int32{ 0, // [0:3] is the sub-list for field type_name } -func init() { file_tunnel_proto_init() } -func file_tunnel_proto_init() { - if File_tunnel_proto != nil { +func init() { file_tunnel_proto_tunnel_proto_init() } +func file_tunnel_proto_tunnel_proto_init() { + if File_tunnel_proto_tunnel_proto != nil { return } type x struct{} out := protoimpl.TypeBuilder{ File: protoimpl.DescBuilder{ GoPackagePath: reflect.TypeOf(x{}).PkgPath(), - RawDescriptor: unsafe.Slice(unsafe.StringData(file_tunnel_proto_rawDesc), len(file_tunnel_proto_rawDesc)), + RawDescriptor: unsafe.Slice(unsafe.StringData(file_tunnel_proto_tunnel_proto_rawDesc), len(file_tunnel_proto_tunnel_proto_rawDesc)), NumEnums: 0, NumMessages: 5, NumExtensions: 0, NumServices: 1, }, - GoTypes: file_tunnel_proto_goTypes, - DependencyIndexes: file_tunnel_proto_depIdxs, - MessageInfos: file_tunnel_proto_msgTypes, + GoTypes: file_tunnel_proto_tunnel_proto_goTypes, + DependencyIndexes: file_tunnel_proto_tunnel_proto_depIdxs, + MessageInfos: file_tunnel_proto_tunnel_proto_msgTypes, }.Build() - File_tunnel_proto = out.File - file_tunnel_proto_goTypes = nil - file_tunnel_proto_depIdxs = nil + File_tunnel_proto_tunnel_proto = out.File + file_tunnel_proto_tunnel_proto_goTypes = nil + file_tunnel_proto_tunnel_proto_depIdxs = nil } diff --git a/tai/tunnel/taipb/tunnel_grpc.pb.go b/tai/tunnel/taipb/tunnel_grpc.pb.go index 973075e6..eb0241b2 100644 --- a/tai/tunnel/taipb/tunnel_grpc.pb.go +++ b/tai/tunnel/taipb/tunnel_grpc.pb.go @@ -2,7 +2,7 @@ // versions: // - protoc-gen-go-grpc v1.6.1 // - protoc v4.25.0 -// source: tunnel.proto +// source: tunnel/proto/tunnel.proto package taipb @@ -147,5 +147,5 @@ var TaiTunnel_ServiceDesc = grpc.ServiceDesc{ ClientStreams: true, }, }, - Metadata: "tunnel.proto", + Metadata: "tunnel/proto/tunnel.proto", } diff --git a/tai/types/types.go b/tai/types/types.go index a6c80be7..d47364c3 100644 --- a/tai/types/types.go +++ b/tai/types/types.go @@ -24,6 +24,7 @@ type Capabilities struct { Docker bool `json:"docker"` K8s bool `json:"k8s"` HostExec bool `json:"host_exec"` + VNC bool `json:"vnc"` } // SystemInfo describes the host machine running Tai.