From 6408c140fc9618f43b64117149b4b78f6c48a30e Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 6 Mar 2026 17:12:36 +0800 Subject: [PATCH] Implement tunnel mode support in Tai service - Introduce a new "tunnel" scheme for the Tai client, allowing connections through Yao's reverse proxy. - Enhance the Tai registry to manage tunnel-connected nodes and their ports. - Add WebSocket and reverse proxy routes for tunnel connections in the OpenAPI server. - Implement tunnel-specific proxy and VNC handling to facilitate communication with containerized environments. - Update gRPC environment variable handling to support tunnel connections. These changes improve the Tai service's flexibility and connectivity options, enabling better integration with remote and containerized environments. --- cmd/start.go | 5 + grpc/grpc.go | 9 + openapi/openapi.go | 7 + sandbox/v2/grpc.go | 32 +- tai/proxy/connect.go | 10 + tai/proxy/proxy.go | 22 ++ tai/registry/registry.go | 401 ++++++++++++++++++++++ tai/registry/registry_test.go | 499 ++++++++++++++++++++++++++++ tai/registry/testing.go | 31 ++ tai/tai.go | 110 ++++++- tai/tunnel/proxy.go | 172 ++++++++++ tai/tunnel/server.go | 267 +++++++++++++++ tai/tunnel/server_test.go | 604 ++++++++++++++++++++++++++++++++++ tai/vnc/vnc.go | 24 ++ 14 files changed, 2188 insertions(+), 5 deletions(-) create mode 100644 tai/registry/registry.go create mode 100644 tai/registry/registry_test.go create mode 100644 tai/registry/testing.go create mode 100644 tai/tunnel/proxy.go create mode 100644 tai/tunnel/server.go create mode 100644 tai/tunnel/server_test.go diff --git a/cmd/start.go b/cmd/start.go index 18d2e367..689e3fe4 100644 --- a/cmd/start.go +++ b/cmd/start.go @@ -35,6 +35,7 @@ import ( "github.com/yaoapp/yao/service" "github.com/yaoapp/yao/setup" "github.com/yaoapp/yao/share" + tairegistry "github.com/yaoapp/yao/tai/registry" itask "github.com/yaoapp/yao/task" ) @@ -231,6 +232,10 @@ var startCmd = &cobra.Command{ ischedule.Start() defer ischedule.Stop() + // Initialize the global Tai registry for tunnel and direct connections + // (must happen before HTTP/gRPC start so handlers can access it) + tairegistry.Init(nil) + // Start HTTP Server srv, err := service.Start(config.Conf) defer func() { diff --git a/grpc/grpc.go b/grpc/grpc.go index 1f204034..b5cc6f1c 100644 --- a/grpc/grpc.go +++ b/grpc/grpc.go @@ -192,6 +192,15 @@ func Stop() { addrs = nil } +// GRPCServer returns the active gRPC server instance. +// Used by the Tai tunnel server to serve data channel connections +// on the existing gRPC server. +func GRPCServer() *grpc.Server { + mu.Lock() + defer mu.Unlock() + return server +} + // Addr returns all addresses the gRPC server is listening on. func Addr() []string { mu.Lock() diff --git a/openapi/openapi.go b/openapi/openapi.go index 9f5e6653..224546a8 100644 --- a/openapi/openapi.go +++ b/openapi/openapi.go @@ -28,6 +28,7 @@ import ( "github.com/yaoapp/yao/openapi/team" openapiTrace "github.com/yaoapp/yao/openapi/trace" "github.com/yaoapp/yao/openapi/user" + taitunnel "github.com/yaoapp/yao/tai/tunnel" ) // Server is the OpenAPI server @@ -175,6 +176,12 @@ func (openapi *OpenAPI) Attach(router *gin.Engine) { sandbox.SetPathPrefix(baseURL) sandbox.Attach(group.Group("/sandbox"), openapi.OAuth) + // Tai tunnel WebSocket and reverse proxy routes + group.GET("/ws/tai", taitunnel.HandleControl) + group.GET("/ws/tai/data/:channel_id", taitunnel.HandleData) + group.Any("/tai/:taiID/proxy/*path", taitunnel.HandleProxy) + group.GET("/tai/:taiID/vnc/*path", taitunnel.HandleVNC) + // Custom handlers (Defined by developer) } diff --git a/sandbox/v2/grpc.go b/sandbox/v2/grpc.go index 522f0123..e676d5a7 100644 --- a/sandbox/v2/grpc.go +++ b/sandbox/v2/grpc.go @@ -4,6 +4,7 @@ import ( "crypto/rand" "encoding/hex" "fmt" + "net/url" "strconv" "strings" ) @@ -35,6 +36,7 @@ func RevokeContainerTokens(refresh string) error { } // BuildGRPCEnv builds the gRPC environment variables for a sandbox container. +// Supports tai:// (direct), tunnel:// (NAT traversal), and local modes. func BuildGRPCEnv(pool *Pool, sandboxID, access, refresh string, grpcPort int) map[string]string { portStr := strconv.Itoa(grpcPort) env := map[string]string{ @@ -42,12 +44,34 @@ func BuildGRPCEnv(pool *Pool, sandboxID, access, refresh string, grpcPort int) m "YAO_TOKEN": access, "YAO_REFRESH_TOKEN": refresh, } - if pool != nil && strings.Contains(pool.Addr, "tai://") { - taiHost := strings.TrimPrefix(pool.Addr, "tai://") + + if pool == nil { + env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%s", portStr) + return env + } + + switch { + case strings.HasPrefix(pool.Addr, "tunnel://"): env["YAO_GRPC_TAI"] = "enable" - env["YAO_GRPC_ADDR"] = fmt.Sprintf("%s:9100", taiHost) + env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%d", grpcPort) env["YAO_GRPC_UPSTREAM"] = fmt.Sprintf("127.0.0.1:%s", portStr) - } else { + + case strings.HasPrefix(pool.Addr, "tai://"): + u, err := url.Parse(pool.Addr) + if err != nil { + env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%s", portStr) + return env + } + taiHost := u.Hostname() + taiPort := u.Port() + if taiPort == "" { + taiPort = "9100" + } + env["YAO_GRPC_TAI"] = "enable" + env["YAO_GRPC_ADDR"] = fmt.Sprintf("%s:%s", taiHost, taiPort) + env["YAO_GRPC_UPSTREAM"] = fmt.Sprintf("127.0.0.1:%s", portStr) + + default: env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%s", portStr) } return env diff --git a/tai/proxy/connect.go b/tai/proxy/connect.go index ecd9f043..b8663d0b 100644 --- a/tai/proxy/connect.go +++ b/tai/proxy/connect.go @@ -21,6 +21,16 @@ func (r *remoteProxy) Connect(ctx context.Context, containerID string, opts Conn return connect(ctx, baseURL, opts.Protocol, r.client) } +// --- Tunnel Connect --- + +func (t *tunnelProxy) Connect(ctx context.Context, containerID string, opts ConnectOptions) (*Connection, error) { + baseURL, err := t.URL(ctx, containerID, opts.Port, opts.Path) + if err != nil { + return nil, err + } + return connect(ctx, baseURL, opts.Protocol, http.DefaultClient) +} + // --- Local Connect --- func (l *localProxy) Connect(ctx context.Context, containerID string, opts ConnectOptions) (*Connection, error) { diff --git a/tai/proxy/proxy.go b/tai/proxy/proxy.go index c79ca5d4..ba556259 100644 --- a/tai/proxy/proxy.go +++ b/tai/proxy/proxy.go @@ -73,6 +73,28 @@ func (r *remoteProxy) Healthz(ctx context.Context) error { return nil } +// --- Tunnel implementation --- + +type tunnelProxy struct { + taiID string + yaoBase string // e.g. "http://yao-host:5099" +} + +// NewTunnel creates a Proxy that routes through Yao's reverse proxy for +// tunnel-connected Tai instances. URLs point to {yaoBase}/tai/{taiID}/proxy/*. +func NewTunnel(taiID, yaoBase string) Proxy { + return &tunnelProxy{taiID: taiID, yaoBase: strings.TrimRight(yaoBase, "/")} +} + +func (t *tunnelProxy) URL(_ context.Context, containerID string, port int, path string) (string, error) { + path = strings.TrimPrefix(path, "/") + return fmt.Sprintf("%s/tai/%s/proxy/%s:%d/%s", t.yaoBase, t.taiID, containerID, port, path), nil +} + +func (t *tunnelProxy) Healthz(_ context.Context) error { + return nil +} + // --- Local implementation --- type localProxy struct { diff --git a/tai/registry/registry.go b/tai/registry/registry.go new file mode 100644 index 00000000..7bb15711 --- /dev/null +++ b/tai/registry/registry.go @@ -0,0 +1,401 @@ +package registry + +import ( + "crypto/rand" + "encoding/hex" + "fmt" + "io" + "log/slog" + "net" + "sync" + "time" + + "github.com/gorilla/websocket" +) + +// TaiNode represents a registered Tai instance (direct or tunnel). +// Internal use only; external callers receive NodeSnapshot via Get()/List(). +type TaiNode struct { + TaiID string + MachineID string + Version string + Auth AuthInfo + Mode string // "direct" | "tunnel" + Addr string // direct mode: "tai-host"; tunnel mode: empty + YaoBase string // Yao server base URL reported by Tai (tunnel mode) + Ports map[string]int // {"grpc":9100, "http":8080, "vnc":6080, "docker":2375} + Capabilities map[string]bool + + ControlConn *websocket.Conn + connMu sync.Mutex // protects ControlConn writes + + Status string // "online" | "offline" | "connecting" + ConnectedAt time.Time + LastPing time.Time + PoolName string + + localListeners map[int]*tunnelListener +} + +// NodeSnapshot is a read-only copy of TaiNode fields safe to use outside locks. +type NodeSnapshot struct { + TaiID string + MachineID string + Version string + Auth AuthInfo + Mode string + Addr string + YaoBase string + Ports map[string]int + Capabilities map[string]bool + Status string + ConnectedAt time.Time + LastPing time.Time + PoolName string +} + +func (n *TaiNode) snapshot() NodeSnapshot { + ports := make(map[string]int, len(n.Ports)) + for k, v := range n.Ports { + ports[k] = v + } + caps := make(map[string]bool, len(n.Capabilities)) + for k, v := range n.Capabilities { + caps[k] = v + } + return NodeSnapshot{ + TaiID: n.TaiID, MachineID: n.MachineID, Version: n.Version, + Auth: n.Auth, Mode: n.Mode, Addr: n.Addr, YaoBase: n.YaoBase, + Ports: ports, Capabilities: caps, + Status: n.Status, ConnectedAt: n.ConnectedAt, LastPing: n.LastPing, + PoolName: n.PoolName, + } +} + +// AuthInfo holds Yao user authorization extracted from OAuth token. +type AuthInfo struct { + Subject string + UserID string + ClientID string + Scope string + TeamID string + TenantID string +} + +// pendingChannel represents a channel awaiting Tai's data WS connection. +type pendingChannel struct { + taiID string + result chan net.Conn + timer *time.Timer +} + +// tunnelListener wraps a TCP listener that bridges each accepted connection +// through the WS tunnel to a specific Tai port. +type tunnelListener struct { + listener net.Listener + taiID string + port int + cancel func() +} + +var ( + global *Registry + once sync.Once +) + +// Registry manages all Tai nodes (direct and tunnel). +type Registry struct { + mu sync.RWMutex + nodes map[string]*TaiNode + pending map[string]*pendingChannel + logger *slog.Logger +} + +// Init initializes the global registry singleton. +func Init(logger *slog.Logger) { + once.Do(func() { + if logger == nil { + logger = slog.Default() + } + global = &Registry{ + nodes: make(map[string]*TaiNode), + pending: make(map[string]*pendingChannel), + logger: logger, + } + }) +} + +// Global returns the global registry instance. +func Global() *Registry { + return global +} + +// Register adds or updates a Tai node in the registry. +func (r *Registry) Register(node *TaiNode) { + r.mu.Lock() + defer r.mu.Unlock() + + node.Status = "online" + node.ConnectedAt = time.Now() + node.LastPing = time.Now() + if node.localListeners == nil { + node.localListeners = make(map[int]*tunnelListener) + } + r.nodes[node.TaiID] = node + + r.logger.Info("tai node registered", + "tai_id", node.TaiID, "mode", node.Mode, "version", node.Version) +} + +// Unregister removes a Tai node and closes its local listeners and control connection. +func (r *Registry) Unregister(taiID string) { + r.mu.Lock() + node, ok := r.nodes[taiID] + if ok { + for _, tl := range node.localListeners { + tl.cancel() + tl.listener.Close() + } + node.connMu.Lock() + if node.ControlConn != nil { + node.ControlConn.Close() + node.ControlConn = nil + } + node.connMu.Unlock() + delete(r.nodes, taiID) + } + r.mu.Unlock() + + if ok { + r.logger.Info("tai node unregistered", "tai_id", taiID) + } +} + +// Get returns a snapshot of a Tai node by ID. Returns nil, false if not found. +func (r *Registry) Get(taiID string) (*NodeSnapshot, bool) { + r.mu.RLock() + defer r.mu.RUnlock() + n, ok := r.nodes[taiID] + if !ok { + return nil, false + } + snap := n.snapshot() + return &snap, true +} + +// List returns snapshots of all registered Tai nodes. +func (r *Registry) List() []NodeSnapshot { + r.mu.RLock() + defer r.mu.RUnlock() + result := make([]NodeSnapshot, 0, len(r.nodes)) + for _, n := range r.nodes { + result = append(result, n.snapshot()) + } + return result +} + +// WriteControlJSON sends a JSON message on the node's control channel +// with proper serialization. Returns error if node not found or not tunnel. +func (r *Registry) WriteControlJSON(taiID string, v interface{}) error { + r.mu.RLock() + node := r.nodes[taiID] + r.mu.RUnlock() + + if node == nil { + return fmt.Errorf("tai node %s not found", taiID) + } + + node.connMu.Lock() + defer node.connMu.Unlock() + if node.ControlConn == nil { + return fmt.Errorf("tai node %s has no active control channel", taiID) + } + return node.ControlConn.WriteJSON(v) +} + +// UpdatePing records a heartbeat timestamp. +func (r *Registry) UpdatePing(taiID string) { + r.mu.Lock() + defer r.mu.Unlock() + if n, ok := r.nodes[taiID]; ok { + n.LastPing = time.Now() + } +} + +// RequestChannel sends an "open" command to a tunnel-connected Tai via its +// control channel. Returns a channel_id that Tai will use to connect back. +// Blocks until the data channel is established or timeout. +func (r *Registry) RequestChannel(taiID string, targetPort int) (string, chan net.Conn, error) { + r.mu.RLock() + node := r.nodes[taiID] + r.mu.RUnlock() + + if node == nil { + return "", nil, fmt.Errorf("tai node %s not found", taiID) + } + if node.Mode != "tunnel" { + return "", nil, fmt.Errorf("tai node %s is not a tunnel node", taiID) + } + node.connMu.Lock() + hasConn := node.ControlConn != nil + node.connMu.Unlock() + if !hasConn { + return "", nil, fmt.Errorf("tai node %s has no active control channel", taiID) + } + + channelID, err := generateChannelID() + if err != nil { + return "", nil, fmt.Errorf("generate channel_id: %w", err) + } + + resultCh := make(chan net.Conn, 1) + timer := time.AfterFunc(30*time.Second, func() { + r.mu.Lock() + if pc, ok := r.pending[channelID]; ok { + close(pc.result) + delete(r.pending, channelID) + } + r.mu.Unlock() + }) + + r.mu.Lock() + r.pending[channelID] = &pendingChannel{taiID: taiID, result: resultCh, timer: timer} + r.mu.Unlock() + + msg := map[string]interface{}{ + "type": "open", + "channel_id": channelID, + "target_port": targetPort, + } + if err := r.WriteControlJSON(taiID, msg); err != nil { + r.mu.Lock() + delete(r.pending, channelID) + r.mu.Unlock() + timer.Stop() + return "", nil, fmt.Errorf("send open command: %w", err) + } + + return channelID, resultCh, nil +} + +// AcceptDataChannel resolves a pending channel when Tai connects its data WS. +// The taiID must match the node that requested the channel via RequestChannel. +func (r *Registry) AcceptDataChannel(channelID, taiID string, conn net.Conn) error { + r.mu.Lock() + pc, ok := r.pending[channelID] + if ok { + delete(r.pending, channelID) + } + r.mu.Unlock() + + if !ok { + return fmt.Errorf("no pending channel for %s", channelID) + } + if pc.taiID != taiID { + pc.timer.Stop() + close(pc.result) + return fmt.Errorf("channel %s: tai_id mismatch (expected %s, got %s)", channelID, pc.taiID, taiID) + } + pc.timer.Stop() + pc.result <- conn + return nil +} + +// OpenLocalListener creates a localhost TCP listener that tunnels every +// accepted connection to the specified port on the given Tai node. +// Returns the listener address (e.g. "127.0.0.1:54321"). +func (r *Registry) OpenLocalListener(taiID string, targetPort int) (net.Listener, error) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return nil, fmt.Errorf("listen: %w", err) + } + + ctx, cancel := newContext() + tl := &tunnelListener{listener: ln, taiID: taiID, port: targetPort, cancel: cancel} + + r.mu.Lock() + node := r.nodes[taiID] + if node == nil { + r.mu.Unlock() + cancel() + ln.Close() + return nil, fmt.Errorf("tai node %s not found", taiID) + } + node.localListeners[targetPort] = tl + r.mu.Unlock() + + go func() { + for { + conn, err := ln.Accept() + if err != nil { + select { + case <-ctx.Done(): + return + default: + r.logger.Debug("tunnel listener accept error", "err", err) + return + } + } + go r.bridgeTunnelConn(taiID, targetPort, conn) + } + }() + + r.logger.Info("tunnel local listener started", + "tai_id", taiID, "target_port", targetPort, "local_addr", ln.Addr().String()) + return ln, nil +} + +func (r *Registry) bridgeTunnelConn(taiID string, targetPort int, localConn net.Conn) { + channelID, resultCh, err := r.RequestChannel(taiID, targetPort) + if err != nil { + localConn.Close() + r.logger.Error("request channel failed", "tai_id", taiID, "port", targetPort, "err", err) + return + } + + remoteConn, ok := <-resultCh + if !ok || remoteConn == nil { + localConn.Close() + r.logger.Error("data channel timeout", "tai_id", taiID, "channel_id", channelID) + return + } + + bridgeTCP(localConn, remoteConn) +} + +// bridgeTCP copies bytes bidirectionally between two net.Conn, closing both when done. +func bridgeTCP(a, b net.Conn) { + var wg sync.WaitGroup + wg.Add(2) + + cp := func(dst, src net.Conn) { + defer wg.Done() + io.Copy(dst, src) + dst.Close() + } + + go cp(a, b) + go cp(b, a) + wg.Wait() +} + +func generateChannelID() (string, error) { + b := make([]byte, 32) + if _, err := rand.Read(b); err != nil { + return "", err + } + return hex.EncodeToString(b), nil +} + +type contextCancel struct { + done chan struct{} +} + +func newContext() (*contextCancel, func()) { + cc := &contextCancel{done: make(chan struct{})} + return cc, func() { close(cc.done) } +} + +func (c *contextCancel) Done() <-chan struct{} { + return c.done +} diff --git a/tai/registry/registry_test.go b/tai/registry/registry_test.go new file mode 100644 index 00000000..25cf9a39 --- /dev/null +++ b/tai/registry/registry_test.go @@ -0,0 +1,499 @@ +package registry + +import ( + "log/slog" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/gorilla/websocket" +) + +// newTestRegistry creates a standalone registry for testing (bypasses global singleton). +func newTestRegistry() *Registry { + return &Registry{ + nodes: make(map[string]*TaiNode), + pending: make(map[string]*pendingChannel), + logger: slog.Default(), + } +} + +func TestRegister_SetsFieldsAndOnline(t *testing.T) { + r := newTestRegistry() + node := &TaiNode{ + TaiID: "tai-001", + MachineID: "m-abc", + Version: "1.0.0", + Mode: "tunnel", + Ports: map[string]int{"grpc": 9100}, + } + r.Register(node) + + snap, ok := r.Get("tai-001") + if !ok { + t.Fatal("expected node to exist after Register") + } + if snap.Status != "online" { + t.Errorf("Status = %q, want online", snap.Status) + } + if snap.MachineID != "m-abc" { + t.Errorf("MachineID = %q, want m-abc", snap.MachineID) + } + if snap.ConnectedAt.IsZero() { + t.Error("ConnectedAt should be set") + } +} + +func TestRegister_Overwrite(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001", Version: "1.0"}) + r.Register(&TaiNode{TaiID: "tai-001", Version: "2.0"}) + + snap, ok := r.Get("tai-001") + if !ok { + t.Fatal("node should exist") + } + if snap.Version != "2.0" { + t.Errorf("Version = %q, want 2.0 after re-register", snap.Version) + } +} + +func TestUnregister_RemovesNode(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001"}) + r.Unregister("tai-001") + + if _, ok := r.Get("tai-001"); ok { + t.Error("expected node to be removed after Unregister") + } +} + +func TestUnregister_Nonexistent(t *testing.T) { + r := newTestRegistry() + r.Unregister("ghost") +} + +func TestGet_NotFound(t *testing.T) { + r := newTestRegistry() + if _, ok := r.Get("missing"); ok { + t.Error("expected false for missing node") + } +} + +func TestList_Empty(t *testing.T) { + r := newTestRegistry() + if got := r.List(); len(got) != 0 { + t.Errorf("List() = %d items, want 0", len(got)) + } +} + +func TestList_MultipleNodes(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "a"}) + r.Register(&TaiNode{TaiID: "b"}) + r.Register(&TaiNode{TaiID: "c"}) + + list := r.List() + if len(list) != 3 { + t.Errorf("List() = %d items, want 3", len(list)) + } + + ids := map[string]bool{} + for _, snap := range list { + ids[snap.TaiID] = true + } + for _, id := range []string{"a", "b", "c"} { + if !ids[id] { + t.Errorf("missing node %q in List()", id) + } + } +} + +func TestSnapshot_DeepCopy(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{ + TaiID: "tai-001", + Ports: map[string]int{"grpc": 9100, "http": 8080}, + }) + + snap, _ := r.Get("tai-001") + snap.Ports["grpc"] = 0 + + snap2, _ := r.Get("tai-001") + if snap2.Ports["grpc"] != 9100 { + t.Error("snapshot modification leaked into registry node") + } +} + +func TestUpdatePing(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001"}) + time.Sleep(10 * time.Millisecond) + + r.UpdatePing("tai-001") + snap, _ := r.Get("tai-001") + if snap.LastPing.Before(snap.ConnectedAt) { + t.Error("LastPing should be after ConnectedAt") + } +} + +func TestUpdatePing_NonexistentNode(t *testing.T) { + r := newTestRegistry() + r.UpdatePing("ghost") +} + +func TestWriteControlJSON_NoNode(t *testing.T) { + r := newTestRegistry() + err := r.WriteControlJSON("missing", map[string]string{"type": "test"}) + if err == nil { + t.Fatal("expected error for missing node") + } +} + +func TestWriteControlJSON_NilConn(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001"}) + err := r.WriteControlJSON("tai-001", map[string]string{"type": "test"}) + if err == nil { + t.Fatal("expected error for nil ControlConn") + } +} + +func TestRequestChannel_NotFound(t *testing.T) { + r := newTestRegistry() + _, _, err := r.RequestChannel("ghost", 9100) + if err == nil { + t.Fatal("expected error for missing node") + } +} + +func TestRequestChannel_DirectMode(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001", Mode: "direct"}) + _, _, err := r.RequestChannel("tai-001", 9100) + if err == nil { + t.Fatal("expected error for direct-mode node") + } +} + +func TestAcceptDataChannel_NotPending(t *testing.T) { + r := newTestRegistry() + pipe1, pipe2 := net.Pipe() + defer pipe1.Close() + defer pipe2.Close() + + err := r.AcceptDataChannel("unknown-channel", "tai-001", pipe1) + if err == nil { + t.Fatal("expected error for non-pending channel") + } +} + +func TestAcceptDataChannel_TaiIDMismatch(t *testing.T) { + r := newTestRegistry() + + resultCh := make(chan net.Conn, 1) + timer := time.AfterFunc(5*time.Second, func() {}) + r.mu.Lock() + r.pending["ch-001"] = &pendingChannel{taiID: "tai-owner", result: resultCh, timer: timer} + r.mu.Unlock() + + pipe1, pipe2 := net.Pipe() + defer pipe1.Close() + defer pipe2.Close() + + err := r.AcceptDataChannel("ch-001", "tai-intruder", pipe1) + if err == nil { + t.Fatal("expected error for tai_id mismatch") + } +} + +func TestAcceptDataChannel_Success(t *testing.T) { + r := newTestRegistry() + + resultCh := make(chan net.Conn, 1) + timer := time.AfterFunc(5*time.Second, func() {}) + r.mu.Lock() + r.pending["ch-002"] = &pendingChannel{taiID: "tai-001", result: resultCh, timer: timer} + r.mu.Unlock() + + pipe1, pipe2 := net.Pipe() + defer pipe2.Close() + + if err := r.AcceptDataChannel("ch-002", "tai-001", pipe1); err != nil { + t.Fatalf("AcceptDataChannel: %v", err) + } + + select { + case conn := <-resultCh: + if conn == nil { + t.Fatal("expected non-nil conn") + } + conn.Close() + case <-time.After(time.Second): + t.Fatal("timeout waiting for conn on resultCh") + } +} + +func TestGenerateChannelID_Unique(t *testing.T) { + seen := make(map[string]bool) + for i := 0; i < 100; i++ { + id, err := generateChannelID() + if err != nil { + t.Fatalf("generateChannelID: %v", err) + } + if len(id) != 64 { + t.Errorf("len = %d, want 64 hex chars", len(id)) + } + if seen[id] { + t.Fatalf("duplicate channel ID: %s", id) + } + seen[id] = true + } +} + +func TestBridgeTCP(t *testing.T) { + a1, a2 := net.Pipe() + b1, b2 := net.Pipe() + + go bridgeTCP(a2, b1) + + msg := []byte("hello tunnel") + go func() { + a1.Write(msg) + a1.Close() + }() + + buf := make([]byte, 64) + n, _ := b2.Read(buf) + if string(buf[:n]) != "hello tunnel" { + t.Errorf("got %q, want %q", buf[:n], "hello tunnel") + } + b2.Close() +} + +func TestConcurrentRegisterGet(t *testing.T) { + r := newTestRegistry() + var wg sync.WaitGroup + + for i := 0; i < 50; i++ { + wg.Add(2) + id := "tai-" + string(rune('A'+i%26)) + + go func() { + defer wg.Done() + r.Register(&TaiNode{TaiID: id, Mode: "tunnel"}) + }() + + go func() { + defer wg.Done() + r.Get(id) + r.List() + }() + } + + wg.Wait() +} + +func TestWriteControlJSON_Success(t *testing.T) { + done := make(chan map[string]string, 1) + + srv := newWSServer(func(conn *websocket.Conn) { + var msg map[string]string + conn.ReadJSON(&msg) + done <- msg + conn.Close() + }) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn}) + + payload := map[string]string{"type": "test", "data": "hello"} + if err := r.WriteControlJSON("tai-001", payload); err != nil { + t.Fatalf("WriteControlJSON: %v", err) + } + + select { + case got := <-done: + if got["type"] != "test" { + t.Errorf("type = %q, want test", got["type"]) + } + if got["data"] != "hello" { + t.Errorf("data = %q, want hello", got["data"]) + } + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for server to receive message") + } +} + +func TestRequestChannel_Success(t *testing.T) { + openCh := make(chan map[string]interface{}, 1) + + srv := newWSServer(func(conn *websocket.Conn) { + var msg map[string]interface{} + conn.ReadJSON(&msg) + openCh <- msg + time.Sleep(time.Second) + conn.Close() + }) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn}) + + channelID, resultCh, err := r.RequestChannel("tai-001", 9100) + if err != nil { + t.Fatalf("RequestChannel: %v", err) + } + if channelID == "" { + t.Fatal("channelID should not be empty") + } + if len(channelID) != 64 { + t.Errorf("channelID len = %d, want 64", len(channelID)) + } + if resultCh == nil { + t.Fatal("resultCh should not be nil") + } + + select { + case cmd := <-openCh: + if cmd["type"] != "open" { + t.Errorf("cmd type = %v, want open", cmd["type"]) + } + if cmd["channel_id"] != channelID { + t.Errorf("cmd channel_id = %v, want %s", cmd["channel_id"], channelID) + } + if tp, ok := cmd["target_port"].(float64); !ok || int(tp) != 9100 { + t.Errorf("cmd target_port = %v, want 9100", cmd["target_port"]) + } + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for open command") + } +} + +func TestRequestChannel_NoControlConn(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel"}) + + _, _, err := r.RequestChannel("tai-001", 9100) + if err == nil { + t.Fatal("expected error for nil ControlConn") + } +} + +func TestOpenLocalListener_Success(t *testing.T) { + r := newTestRegistry() + + controlCh := make(chan map[string]interface{}, 1) + srv := newWSServer(func(conn *websocket.Conn) { + for { + var msg map[string]interface{} + if err := conn.ReadJSON(&msg); err != nil { + return + } + controlCh <- msg + } + }) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + + r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn}) + + ln, err := r.OpenLocalListener("tai-001", 9100) + if err != nil { + t.Fatalf("OpenLocalListener: %v", err) + } + defer ln.Close() + + addr := ln.Addr().String() + if addr == "" { + t.Fatal("listener address should not be empty") + } + if !strings.HasPrefix(addr, "127.0.0.1:") { + t.Errorf("addr = %q, want 127.0.0.1:*", addr) + } + + conn, err := net.DialTimeout("tcp", addr, time.Second) + if err != nil { + t.Fatalf("connect to local listener: %v", err) + } + defer conn.Close() + + select { + case cmd := <-controlCh: + if cmd["type"] != "open" { + t.Errorf("open cmd type = %v, want open", cmd["type"]) + } + if _, ok := cmd["channel_id"].(string); !ok { + t.Error("open cmd missing channel_id") + } + if tp, ok := cmd["target_port"].(float64); !ok || int(tp) != 9100 { + t.Errorf("target_port = %v, want 9100", cmd["target_port"]) + } + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for open command from local listener") + } +} + +func TestOpenLocalListener_NodeNotFound(t *testing.T) { + r := newTestRegistry() + _, err := r.OpenLocalListener("ghost", 9100) + if err == nil { + t.Fatal("expected error for missing node") + } +} + +func newWSServer(handler func(*websocket.Conn)) *httptest.Server { + up := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }} + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := up.Upgrade(w, r, nil) + if err != nil { + return + } + handler(conn) + })) +} + +func TestNodeSnapshot_AuthInfo(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{ + TaiID: "tai-001", + Auth: AuthInfo{ + Subject: "user123", + ClientID: "tai-001", + Scope: "tai:tunnel", + }, + }) + + snap, ok := r.Get("tai-001") + if !ok { + t.Fatal("node not found") + } + if snap.Auth.Subject != "user123" { + t.Errorf("Auth.Subject = %q, want user123", snap.Auth.Subject) + } + if snap.Auth.Scope != "tai:tunnel" { + t.Errorf("Auth.Scope = %q, want tai:tunnel", snap.Auth.Scope) + } +} diff --git a/tai/registry/testing.go b/tai/registry/testing.go new file mode 100644 index 00000000..6da804b0 --- /dev/null +++ b/tai/registry/testing.go @@ -0,0 +1,31 @@ +package registry + +import ( + "log/slog" + "net" + "time" +) + +// NewForTest creates a standalone Registry for use in tests. +// Not intended for production use. +func NewForTest() *Registry { + return &Registry{ + nodes: make(map[string]*TaiNode), + pending: make(map[string]*pendingChannel), + logger: slog.Default(), + } +} + +// SetGlobalForTest replaces the global registry singleton for testing. +// Not intended for production use. +func SetGlobalForTest(r *Registry) { + global = r +} + +// SetPendingForTest injects a pending channel entry for testing. +// Not intended for production use. +func (r *Registry) SetPendingForTest(channelID, taiID string, result chan net.Conn, timer *time.Timer) { + r.mu.Lock() + defer r.mu.Unlock() + r.pending[channelID] = &pendingChannel{taiID: taiID, result: result, timer: timer} +} diff --git a/tai/tai.go b/tai/tai.go index 1fc52ec0..76acf990 100644 --- a/tai/tai.go +++ b/tai/tai.go @@ -3,6 +3,7 @@ package tai import ( "context" "fmt" + "net" "net/http" "net/url" "strconv" @@ -10,6 +11,7 @@ import ( "time" "github.com/yaoapp/yao/tai/proxy" + "github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/sandbox" sipb "github.com/yaoapp/yao/tai/serverinfo/pb" "github.com/yaoapp/yao/tai/vnc" @@ -124,7 +126,7 @@ func mergedPorts(p Ports) Ports { // Client provides unified access to all Tai SDK sub-packages. type Client struct { - scheme string // "tai" or "docker" + scheme string // "tai", "docker", or "tunnel" host string addr string ports Ports @@ -135,6 +137,9 @@ type Client struct { prx proxy.Proxy vc vnc.VNC grpcConn *grpc.ClientConn + + // tunnel mode: local listeners that bridge to Tai via WS + tunnelListeners []net.Listener } // New creates a Client based on the address protocol: @@ -172,6 +177,8 @@ func New(addr string, opts ...Option) (*Client, error) { return c.initLocal(cfg) case "tai": return c.initRemote(cfg) + case "tunnel": + return c.initTunnel(cfg) default: return nil, fmt.Errorf("unsupported scheme: %s", scheme) } @@ -256,9 +263,97 @@ func (c *Client) initRemote(cfg *config) (*Client, error) { hc := cfg.httpClient c.prx = proxy.NewRemote(c.host, c.ports.HTTP, hc) c.vc = vnc.NewRemote(c.host, c.ports.VNC, hc) + + if reg := registry.Global(); reg != nil { + reg.Register(®istry.TaiNode{ + TaiID: c.host, + Mode: "direct", + Addr: c.host, + Ports: map[string]int{ + "grpc": c.ports.GRPC, + "http": c.ports.HTTP, + "vnc": c.ports.VNC, + "docker": c.ports.Docker, + "k8s": c.ports.K8s, + }, + }) + } + return c, nil } +func (c *Client) initTunnel(cfg *config) (*Client, error) { + reg := registry.Global() + if reg == nil { + return nil, fmt.Errorf("tai registry not initialized") + } + + taiID := c.host // for tunnel:// scheme, host stores the taiID + node, ok := reg.Get(taiID) + if !ok || node.Status != "online" { + return nil, fmt.Errorf("tai node %s not online", taiID) + } + + c.ports = Ports{ + GRPC: nodePort(node.Ports, "grpc", 9100), + HTTP: nodePort(node.Ports, "http", 8080), + VNC: nodePort(node.Ports, "vnc", 6080), + Docker: nodePort(node.Ports, "docker", 2375), + } + + grpcLn, err := reg.OpenLocalListener(taiID, c.ports.GRPC) + if err != nil { + return nil, fmt.Errorf("open grpc tunnel listener: %w", err) + } + c.tunnelListeners = append(c.tunnelListeners, grpcLn) + + grpcAddr := grpcLn.Addr().String() + conn, err := grpc.NewClient("passthrough:///"+grpcAddr, + grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + grpcLn.Close() + return nil, fmt.Errorf("grpc dial tunnel %s: %w", grpcAddr, err) + } + c.grpcConn = conn + c.vol = volume.NewRemote(conn) + + dockerLn, err := reg.OpenLocalListener(taiID, c.ports.Docker) + if err != nil { + conn.Close() + grpcLn.Close() + return nil, fmt.Errorf("open docker tunnel listener: %w", err) + } + c.tunnelListeners = append(c.tunnelListeners, dockerLn) + + sbAddr := fmt.Sprintf("tcp://%s", dockerLn.Addr().String()) + sb, err := sandbox.NewDocker(sbAddr) + if err != nil { + c.closeTunnelListeners() + conn.Close() + return nil, err + } + c.sb = sb + c.img = sandbox.NewDockerImage(sandbox.DockerCli(sb)) + + c.prx = proxy.NewTunnel(taiID, node.YaoBase) + c.vc = vnc.NewTunnel(taiID, node.YaoBase) + return c, nil +} + +func (c *Client) closeTunnelListeners() { + for _, ln := range c.tunnelListeners { + ln.Close() + } + c.tunnelListeners = nil +} + +func nodePort(ports map[string]int, key string, fallback int) int { + if p, ok := ports[key]; ok && p > 0 { + return p + } + return fallback +} + // Close releases all resources. func (c *Client) Close() error { var errs []error @@ -277,6 +372,12 @@ func (c *Client) Close() error { errs = append(errs, err) } } + c.closeTunnelListeners() + if c.scheme == "tai" { + if reg := registry.Global(); reg != nil { + reg.Unregister(c.host) + } + } if len(errs) > 0 { return fmt.Errorf("close: %v", errs) } @@ -355,6 +456,13 @@ func parseAddr(addr string) (scheme, host, dockerAddr string, grpcPort int, err } return "tai", hostname, "", grpcPort, nil + case "tunnel": + taiID := u.Host + if taiID == "" { + return "", "", "", 0, fmt.Errorf("tunnel:// requires a tai ID") + } + return "tunnel", taiID, "", 0, nil + case "docker": return "docker", "", addr, 0, nil diff --git a/tai/tunnel/proxy.go b/tai/tunnel/proxy.go new file mode 100644 index 00000000..a02d92aa --- /dev/null +++ b/tai/tunnel/proxy.go @@ -0,0 +1,172 @@ +package tunnel + +import ( + "bufio" + "io" + "log/slog" + "net" + "net/http" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/yaoapp/yao/tai/registry" +) + +// HandleProxy handles HTTP reverse proxy requests for a tunnel-connected Tai: +// ANY /tai/:taiID/proxy/*path +// Opens a data channel to Tai's HTTP port, forwards the HTTP request, +// and streams the response back. +func HandleProxy(c *gin.Context) { + logger := slog.Default() + reg := registry.Global() + if reg == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "registry not initialized"}) + return + } + + taiID := c.Param("taiID") + node, ok := reg.Get(taiID) + if !ok || node.Status != "online" { + c.JSON(http.StatusBadGateway, gin.H{"error": "tai node not available"}) + return + } + + httpPort := node.Ports["http"] + if httpPort == 0 { + httpPort = 8080 + } + + channelID, resultCh, err := reg.RequestChannel(taiID, httpPort) + if err != nil { + logger.Error("request channel failed", "tai_id", taiID, "err", err) + c.JSON(http.StatusBadGateway, gin.H{"error": "tunnel channel failed"}) + return + } + + remoteConn, ok := <-resultCh + if !ok || remoteConn == nil { + logger.Error("data channel timeout", "tai_id", taiID, "channel_id", channelID) + c.JSON(http.StatusGatewayTimeout, gin.H{"error": "data channel timeout"}) + return + } + defer remoteConn.Close() + + path := c.Param("path") + outReq, err := http.NewRequestWithContext(c.Request.Context(), c.Request.Method, "http://tai-tunnel"+path, c.Request.Body) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "build request failed"}) + return + } + outReq.Header = c.Request.Header.Clone() + outReq.Host = c.Request.Host + + if err := outReq.Write(remoteConn); err != nil { + logger.Error("write request to tunnel", "err", err) + c.JSON(http.StatusBadGateway, gin.H{"error": "write to tunnel failed"}) + return + } + + resp, err := http.ReadResponse(bufio.NewReader(remoteConn), outReq) + if err != nil { + logger.Error("read response from tunnel", "err", err) + c.JSON(http.StatusBadGateway, gin.H{"error": "read from tunnel failed"}) + return + } + defer resp.Body.Close() + + for k, vv := range resp.Header { + for _, v := range vv { + c.Writer.Header().Add(k, v) + } + } + c.Writer.WriteHeader(resp.StatusCode) + io.Copy(c.Writer, resp.Body) +} + +// HandleVNC handles VNC WebSocket proxying for a tunnel-connected Tai: +// GET /tai/:taiID/vnc/*path +// Upgrades the client connection to WebSocket, opens a data channel to +// Tai's VNC port, and bridges the two WebSocket connections. +func HandleVNC(c *gin.Context) { + logger := slog.Default() + reg := registry.Global() + if reg == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "registry not initialized"}) + return + } + + taiID := c.Param("taiID") + node, ok := reg.Get(taiID) + if !ok || node.Status != "online" { + c.JSON(http.StatusBadGateway, gin.H{"error": "tai node not available"}) + return + } + + vncPort := node.Ports["vnc"] + if vncPort == 0 { + vncPort = 6080 + } + + channelID, resultCh, err := reg.RequestChannel(taiID, vncPort) + if err != nil { + logger.Error("request vnc channel failed", "tai_id", taiID, "err", err) + c.JSON(http.StatusBadGateway, gin.H{"error": "tunnel channel failed"}) + return + } + + clientConn, err := upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + logger.Error("ws upgrade client failed", "err", err) + return + } + + taiConn, ok := <-resultCh + if !ok || taiConn == nil { + logger.Error("vnc data channel timeout", "tai_id", taiID, "channel_id", channelID) + clientConn.Close() + return + } + + bridgeWSToConn(clientConn, taiConn) +} + +// bridgeWSToConn bridges a client WebSocket to a net.Conn (tunnel data channel). +func bridgeWSToConn(clientWS *websocket.Conn, taiConn net.Conn) { + done := make(chan struct{}, 2) + + // client WS -> tai conn + go func() { + defer func() { done <- struct{}{} }() + for { + _, data, err := clientWS.ReadMessage() + if err != nil { + return + } + if _, err := taiConn.Write(data); err != nil { + return + } + } + }() + + // tai conn -> client WS + go func() { + defer func() { done <- struct{}{} }() + buf := make([]byte, 32*1024) + for { + n, err := taiConn.Read(buf) + if n > 0 { + if wErr := clientWS.WriteMessage(websocket.BinaryMessage, buf[:n]); wErr != nil { + return + } + } + if err != nil { + return + } + } + }() + + <-done + clientWS.Close() + taiConn.Close() + <-done +} diff --git a/tai/tunnel/server.go b/tai/tunnel/server.go new file mode 100644 index 00000000..1910358c --- /dev/null +++ b/tai/tunnel/server.go @@ -0,0 +1,267 @@ +package tunnel + +import ( + "fmt" + "io" + "log/slog" + "net" + "net/http" + "strings" + "sync" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + oauth "github.com/yaoapp/yao/openapi/oauth" + "github.com/yaoapp/yao/tai/registry" +) + +var upgrader = websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, +} + +// HandleControl handles the Tai control channel WebSocket: GET /ws/tai. +// Authenticates via Bearer token, reads register + ping messages, +// and maintains the Tai node in the global registry. +func HandleControl(c *gin.Context) { + logger := slog.Default() + reg := registry.Global() + if reg == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "registry not initialized"}) + return + } + + bearer := extractBearer(c.Request) + if bearer == "" { + c.JSON(http.StatusUnauthorized, gin.H{"error": "missing authorization"}) + return + } + + authInfo, err := authenticateBearerFunc(bearer) + if err != nil { + logger.Warn("tunnel auth failed", "err", err) + c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"}) + return + } + + conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + logger.Error("ws upgrade failed", "err", err) + return + } + + // Read the register message + var regMsg registerMessage + if err := conn.ReadJSON(®Msg); err != nil { + logger.Error("read register message", "err", err) + conn.Close() + return + } + if regMsg.Type != "register" { + logger.Error("expected register message", "got", regMsg.Type) + conn.Close() + return + } + if regMsg.TaiID == "" { + logger.Error("register message missing tai_id") + conn.Close() + return + } + + node := ®istry.TaiNode{ + TaiID: regMsg.TaiID, + MachineID: regMsg.MachineID, + Version: regMsg.Version, + Auth: authInfo, + Mode: "tunnel", + YaoBase: regMsg.Server, + Ports: regMsg.Ports, + Capabilities: regMsg.Capabilities, + ControlConn: conn, + } + reg.Register(node) + defer func() { + reg.Unregister(regMsg.TaiID) + logger.Info("tai tunnel disconnected", "tai_id", regMsg.TaiID) + }() + + if err := reg.WriteControlJSON(regMsg.TaiID, map[string]string{"type": "registered", "tai_id": regMsg.TaiID}); err != nil { + logger.Error("write registered response", "err", err) + return + } + + logger.Info("tai tunnel connected", "tai_id", regMsg.TaiID, "version", regMsg.Version) + + for { + var msg controlMsg + if err := conn.ReadJSON(&msg); err != nil { + if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) { + logger.Debug("control channel read error", "err", err) + } + return + } + + switch msg.Type { + case "ping": + reg.UpdatePing(regMsg.TaiID) + if err := reg.WriteControlJSON(regMsg.TaiID, map[string]string{"type": "pong"}); err != nil { + logger.Debug("pong write failed", "err", err) + return + } + default: + logger.Debug("unknown control message", "type", msg.Type) + } + } +} + +// HandleData handles a Tai data channel WebSocket: GET /ws/tai/data/:channel_id. +// Authenticates via Bearer token, verifies the caller matches the pending +// channel's owner, then wraps the WS as a net.Conn for bidirectional bridging. +func HandleData(c *gin.Context) { + logger := slog.Default() + reg := registry.Global() + if reg == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "registry not initialized"}) + return + } + + bearer := extractBearer(c.Request) + if bearer == "" { + c.JSON(http.StatusUnauthorized, gin.H{"error": "missing authorization"}) + return + } + authInfo, err := authenticateBearerFunc(bearer) + if err != nil { + logger.Warn("data channel auth failed", "err", err) + c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"}) + return + } + + channelID := c.Param("channel_id") + if channelID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "missing channel_id"}) + return + } + + conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + logger.Error("ws data upgrade failed", "err", err) + return + } + + wsConn := newWSConn(conn) + if err := reg.AcceptDataChannel(channelID, authInfo.ClientID, wsConn); err != nil { + logger.Debug("accept data channel failed", "channel_id", channelID, "err", err) + conn.Close() + return + } +} + +// registerMessage is the JSON structure for Tai's register message. +type registerMessage struct { + Type string `json:"type"` + TaiID string `json:"tai_id"` + MachineID string `json:"machine_id"` + Version string `json:"version"` + Server string `json:"server"` + Ports map[string]int `json:"ports"` + Capabilities map[string]bool `json:"capabilities"` +} + +// controlMsg is a generic control channel message. +type controlMsg struct { + Type string `json:"type"` +} + +func extractBearer(r *http.Request) string { + auth := r.Header.Get("Authorization") + if len(auth) > 7 && strings.EqualFold(auth[:7], "bearer ") { + return auth[7:] + } + return "" +} + +var authenticateBearerFunc = authenticateBearerDefault + +func authenticateBearerDefault(token string) (registry.AuthInfo, error) { + svc := oauth.OAuth + if svc == nil { + return registry.AuthInfo{}, fmt.Errorf("oauth service not initialized") + } + + result, err := svc.AuthenticateToken(oauth.AuthInput{ + AccessToken: token, + }) + if err != nil { + return registry.AuthInfo{}, err + } + + info := registry.AuthInfo{} + if result.Info != nil { + info.Subject = result.Info.Subject + info.UserID = result.Info.UserID + info.ClientID = result.Info.ClientID + info.Scope = result.Info.Scope + info.TeamID = result.Info.TeamID + info.TenantID = result.Info.TenantID + } + return info, nil +} + +// wsConn wraps a gorilla/websocket.Conn to implement net.Conn for raw byte bridging. +type wsConn struct { + ws *websocket.Conn + reader io.Reader + mu sync.Mutex +} + +func newWSConn(ws *websocket.Conn) *wsConn { + return &wsConn{ws: ws} +} + +func (c *wsConn) Read(p []byte) (int, error) { + for { + if c.reader != nil { + n, err := c.reader.Read(p) + if n > 0 { + return n, nil + } + c.reader = nil + if err != nil && err != io.EOF { + return 0, err + } + } + _, reader, err := c.ws.NextReader() + if err != nil { + return 0, err + } + c.reader = reader + } +} + +func (c *wsConn) Write(p []byte) (int, error) { + c.mu.Lock() + defer c.mu.Unlock() + err := c.ws.WriteMessage(websocket.BinaryMessage, p) + if err != nil { + return 0, err + } + return len(p), nil +} + +func (c *wsConn) Close() error { + return c.ws.Close() +} + +func (c *wsConn) LocalAddr() net.Addr { return c.ws.LocalAddr() } +func (c *wsConn) RemoteAddr() net.Addr { return c.ws.RemoteAddr() } + +func (c *wsConn) SetDeadline(t time.Time) error { + if err := c.ws.SetReadDeadline(t); err != nil { + return err + } + return c.ws.SetWriteDeadline(t) +} + +func (c *wsConn) SetReadDeadline(t time.Time) error { return c.ws.SetReadDeadline(t) } +func (c *wsConn) SetWriteDeadline(t time.Time) error { return c.ws.SetWriteDeadline(t) } diff --git a/tai/tunnel/server_test.go b/tai/tunnel/server_test.go new file mode 100644 index 00000000..186e13b5 --- /dev/null +++ b/tai/tunnel/server_test.go @@ -0,0 +1,604 @@ +package tunnel + +import ( + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/yaoapp/yao/tai/registry" +) + +func init() { + gin.SetMode(gin.TestMode) +} + +func setupTestRegistry() *registry.Registry { + r := registry.NewForTest() + registry.SetGlobalForTest(r) + return r +} + +func mockAuth(info registry.AuthInfo, authErr error) func() { + old := authenticateBearerFunc + authenticateBearerFunc = func(token string) (registry.AuthInfo, error) { + return info, authErr + } + return func() { authenticateBearerFunc = old } +} + +// --- extractBearer --- + +func TestExtractBearer(t *testing.T) { + tests := []struct { + name string + header string + want string + }{ + {"valid", "Bearer abc123", "abc123"}, + {"lowercase", "bearer xyz", "xyz"}, + {"empty", "", ""}, + {"no_scheme", "abc123", ""}, + {"only_bearer", "Bearer ", ""}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + r := &http.Request{Header: http.Header{}} + if tt.header != "" { + r.Header.Set("Authorization", tt.header) + } + got := extractBearer(r) + if got != tt.want { + t.Errorf("extractBearer(%q) = %q, want %q", tt.header, got, tt.want) + } + }) + } +} + +// --- wsConn --- + +func TestWSConn_EchoRoundTrip(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + + wc := newWSConn(conn) + buf := make([]byte, 256) + n, err := wc.Read(buf) + if err != nil { + return + } + wc.Write(buf[:n]) + })) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + + if resp.StatusCode != http.StatusSwitchingProtocols { + t.Errorf("handshake status = %d, want 101", resp.StatusCode) + } + + msg := []byte("hello tunnel") + if err := conn.WriteMessage(websocket.BinaryMessage, msg); err != nil { + t.Fatalf("write: %v", err) + } + + mt, reply, err := conn.ReadMessage() + if err != nil { + t.Fatalf("read: %v", err) + } + if mt != websocket.BinaryMessage { + t.Errorf("type = %d, want BinaryMessage(%d)", mt, websocket.BinaryMessage) + } + if string(reply) != "hello tunnel" { + t.Errorf("reply = %q, want %q", reply, "hello tunnel") + } +} + +func TestWSConn_MultipleMessages(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + + wc := newWSConn(conn) + for i := 0; i < 3; i++ { + buf := make([]byte, 256) + n, err := wc.Read(buf) + if err != nil { + return + } + wc.Write(buf[:n]) + } + })) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + + for i, msg := range []string{"one", "two", "three"} { + conn.WriteMessage(websocket.BinaryMessage, []byte(msg)) + _, reply, err := conn.ReadMessage() + if err != nil { + t.Fatalf("round %d read: %v", i, err) + } + if string(reply) != msg { + t.Errorf("round %d: got %q, want %q", i, reply, msg) + } + } +} + +func TestWSConn_ImplementsNetConn(t *testing.T) { + var _ net.Conn = (*wsConn)(nil) +} + +func TestWSConn_LocalRemoteAddr(t *testing.T) { + addrCh := make(chan [2]net.Addr, 1) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + wc := newWSConn(conn) + addrCh <- [2]net.Addr{wc.LocalAddr(), wc.RemoteAddr()} + wc.Close() + })) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + + select { + case addrs := <-addrCh: + if addrs[0] == nil { + t.Error("LocalAddr should not be nil") + } + if addrs[1] == nil { + t.Error("RemoteAddr should not be nil") + } + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for addresses") + } +} + +// --- HandleControl --- + +func newGinRouter() *gin.Engine { + r := gin.New() + r.GET("/ws/tai", HandleControl) + r.GET("/ws/tai/data/:channel_id", HandleData) + return r +} + +func TestHandleControl_NoRegistry(t *testing.T) { + registry.SetGlobalForTest(nil) + defer setupTestRegistry() + + restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil) + defer restore() + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai" + _, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{ + "Authorization": []string{"Bearer test-token"}, + }) + if err == nil { + t.Fatal("expected dial to fail when registry is nil") + } + if resp != nil && resp.StatusCode != http.StatusServiceUnavailable { + t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusServiceUnavailable) + } +} + +func TestHandleControl_NoAuth(t *testing.T) { + setupTestRegistry() + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai" + _, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err == nil { + t.Fatal("expected dial to fail without auth") + } + if resp != nil && resp.StatusCode != http.StatusUnauthorized { + t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusUnauthorized) + } +} + +func TestHandleControl_AuthFailed(t *testing.T) { + setupTestRegistry() + restore := mockAuth(registry.AuthInfo{}, fmt.Errorf("bad token")) + defer restore() + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai" + _, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{ + "Authorization": []string{"Bearer bad-token"}, + }) + if err == nil { + t.Fatal("expected dial to fail with bad auth") + } + if resp != nil && resp.StatusCode != http.StatusUnauthorized { + t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusUnauthorized) + } +} + +func TestHandleControl_RegisterAndPing(t *testing.T) { + reg := setupTestRegistry() + restore := mockAuth(registry.AuthInfo{ + ClientID: "tai-001", + Subject: "user-test", + Scope: "tai:tunnel", + }, nil) + defer restore() + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai" + conn, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{ + "Authorization": []string{"Bearer valid-token"}, + }) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + + if resp.StatusCode != http.StatusSwitchingProtocols { + t.Errorf("handshake = %d, want 101", resp.StatusCode) + } + + regMsg := registerMessage{ + Type: "register", + TaiID: "tai-001", + MachineID: "m-test", + Version: "2.0", + Ports: map[string]int{"grpc": 9100}, + } + if err := conn.WriteJSON(regMsg); err != nil { + t.Fatalf("write register: %v", err) + } + + var registered map[string]string + if err := conn.ReadJSON(®istered); err != nil { + t.Fatalf("read registered: %v", err) + } + if registered["type"] != "registered" { + t.Errorf("response type = %q, want registered", registered["type"]) + } + if registered["tai_id"] != "tai-001" { + t.Errorf("response tai_id = %q, want tai-001", registered["tai_id"]) + } + + snap, ok := reg.Get("tai-001") + if !ok { + t.Fatal("node not found in registry after register") + } + if snap.Status != "online" { + t.Errorf("Status = %q, want online", snap.Status) + } + if snap.MachineID != "m-test" { + t.Errorf("MachineID = %q, want m-test", snap.MachineID) + } + if snap.Version != "2.0" { + t.Errorf("Version = %q, want 2.0", snap.Version) + } + if snap.Mode != "tunnel" { + t.Errorf("Mode = %q, want tunnel", snap.Mode) + } + if snap.Auth.ClientID != "tai-001" { + t.Errorf("Auth.ClientID = %q, want tai-001", snap.Auth.ClientID) + } + if snap.Auth.Subject != "user-test" { + t.Errorf("Auth.Subject = %q, want user-test", snap.Auth.Subject) + } + if snap.Ports["grpc"] != 9100 { + t.Errorf("Ports[grpc] = %d, want 9100", snap.Ports["grpc"]) + } + + time.Sleep(10 * time.Millisecond) + if err := conn.WriteJSON(map[string]string{"type": "ping"}); err != nil { + t.Fatalf("write ping: %v", err) + } + + var pong map[string]string + if err := conn.ReadJSON(&pong); err != nil { + t.Fatalf("read pong: %v", err) + } + if pong["type"] != "pong" { + t.Errorf("pong type = %q, want pong", pong["type"]) + } + + snap2, _ := reg.Get("tai-001") + if !snap2.LastPing.After(snap.LastPing) { + t.Error("LastPing should be updated after ping") + } + + conn.WriteMessage(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) + time.Sleep(100 * time.Millisecond) + + if _, ok := reg.Get("tai-001"); ok { + t.Error("node should be unregistered after connection close") + } +} + +func TestHandleControl_BadRegisterType(t *testing.T) { + setupTestRegistry() + restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil) + defer restore() + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai" + conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{ + "Authorization": []string{"Bearer valid-token"}, + }) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + + conn.WriteJSON(map[string]string{"type": "not-register"}) + _, _, readErr := conn.ReadMessage() + if readErr == nil { + t.Error("expected connection to close for bad register type") + } +} + +func TestHandleControl_MissingTaiID(t *testing.T) { + setupTestRegistry() + restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil) + defer restore() + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai" + conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{ + "Authorization": []string{"Bearer valid-token"}, + }) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + + conn.WriteJSON(map[string]string{"type": "register"}) + _, _, readErr := conn.ReadMessage() + if readErr == nil { + t.Error("expected connection to close for missing tai_id") + } +} + +// --- HandleData --- + +func TestHandleData_NoAuth(t *testing.T) { + setupTestRegistry() + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/ch-001" + _, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err == nil { + t.Fatal("expected dial to fail without auth") + } + if resp != nil && resp.StatusCode != http.StatusUnauthorized { + t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusUnauthorized) + } +} + +func TestHandleData_AcceptSuccess(t *testing.T) { + reg := setupTestRegistry() + restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil) + defer restore() + + resultCh := make(chan net.Conn, 1) + timer := time.AfterFunc(5*time.Second, func() {}) + reg.SetPendingForTest("ch-test-123", "tai-001", resultCh, timer) + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/ch-test-123" + conn, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{ + "Authorization": []string{"Bearer valid-token"}, + }) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + + if resp.StatusCode != http.StatusSwitchingProtocols { + t.Errorf("status = %d, want 101", resp.StatusCode) + } + + select { + case c := <-resultCh: + if c == nil { + t.Fatal("expected non-nil conn from resultCh") + } + c.Close() + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for conn on resultCh") + } +} + +func TestHandleData_ChannelNotPending(t *testing.T) { + setupTestRegistry() + restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil) + defer restore() + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/nonexistent" + conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{ + "Authorization": []string{"Bearer valid-token"}, + }) + if err != nil { + return + } + defer conn.Close() + + _, _, readErr := conn.ReadMessage() + if readErr == nil { + t.Error("expected connection to close for non-pending channel") + } +} + +func TestHandleData_TaiIDMismatch(t *testing.T) { + reg := setupTestRegistry() + restore := mockAuth(registry.AuthInfo{ClientID: "tai-intruder"}, nil) + defer restore() + + resultCh := make(chan net.Conn, 1) + timer := time.AfterFunc(5*time.Second, func() {}) + reg.SetPendingForTest("ch-mismatch", "tai-owner", resultCh, timer) + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/ch-mismatch" + conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{ + "Authorization": []string{"Bearer valid-token"}, + }) + if err != nil { + return + } + defer conn.Close() + + _, _, readErr := conn.ReadMessage() + if readErr == nil { + t.Error("expected connection to close for tai_id mismatch") + } +} + +// --- Full open-channel flow --- + +func TestHandleControl_OpenChannelAndBridge(t *testing.T) { + reg := setupTestRegistry() + restore := mockAuth(registry.AuthInfo{ + ClientID: "tai-001", + Subject: "user-test", + }, nil) + defer restore() + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai" + ctrlConn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{ + "Authorization": []string{"Bearer valid-token"}, + }) + if err != nil { + t.Fatalf("dial control: %v", err) + } + defer ctrlConn.Close() + + ctrlConn.WriteJSON(registerMessage{ + Type: "register", + TaiID: "tai-001", + Ports: map[string]int{"grpc": 9100}, + }) + var registered map[string]string + if err := ctrlConn.ReadJSON(®istered); err != nil { + t.Fatalf("read registered: %v", err) + } + if registered["type"] != "registered" { + t.Fatalf("expected registered, got %v", registered) + } + + var wg sync.WaitGroup + wg.Add(1) + var requestErr error + var channelConn net.Conn + go func() { + defer wg.Done() + _, resultCh, err := reg.RequestChannel("tai-001", 9100) + if err != nil { + requestErr = err + return + } + channelConn = <-resultCh + }() + + time.Sleep(50 * time.Millisecond) + + var openCmd map[string]interface{} + if err := ctrlConn.ReadJSON(&openCmd); err != nil { + t.Fatalf("read open cmd: %v", err) + } + if openCmd["type"] != "open" { + t.Errorf("open type = %v, want open", openCmd["type"]) + } + channelID, ok := openCmd["channel_id"].(string) + if !ok || channelID == "" { + t.Fatalf("missing channel_id: %v", openCmd) + } + if tp, ok := openCmd["target_port"].(float64); !ok || int(tp) != 9100 { + t.Errorf("target_port = %v, want 9100", openCmd["target_port"]) + } + + dataURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/" + channelID + dataConn, _, err := websocket.DefaultDialer.Dial(dataURL, http.Header{ + "Authorization": []string{"Bearer valid-token"}, + }) + if err != nil { + t.Fatalf("dial data: %v", err) + } + defer dataConn.Close() + + wg.Wait() + if requestErr != nil { + t.Fatalf("RequestChannel: %v", requestErr) + } + if channelConn == nil { + t.Fatal("expected non-nil conn from RequestChannel") + } + defer channelConn.Close() + + payload := []byte("grpc-payload-test") + dataConn.WriteMessage(websocket.BinaryMessage, payload) + + buf := make([]byte, 256) + n, err := channelConn.Read(buf) + if err != nil && err != io.EOF { + t.Fatalf("read bridged: %v", err) + } + if string(buf[:n]) != "grpc-payload-test" { + t.Errorf("bridged data = %q, want %q", buf[:n], "grpc-payload-test") + } +} diff --git a/tai/vnc/vnc.go b/tai/vnc/vnc.go index ec03348d..ba36b7bb 100644 --- a/tai/vnc/vnc.go +++ b/tai/vnc/vnc.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "net/http" + "strings" "github.com/yaoapp/yao/tai/sandbox" ) @@ -51,6 +52,29 @@ func (r *remoteVNC) Ping(ctx context.Context, containerID string) error { return nil } +// --- Tunnel implementation --- + +type tunnelVNC struct { + taiID string + yaoBase string // e.g. "http://yao-host:5099" +} + +// NewTunnel creates a VNC that routes through Yao's reverse proxy +// for tunnel-connected Tai instances. +func NewTunnel(taiID, yaoBase string) VNC { + return &tunnelVNC{taiID: taiID, yaoBase: strings.TrimRight(yaoBase, "/")} +} + +func (t *tunnelVNC) URL(_ context.Context, containerID string) (string, error) { + base := strings.Replace(t.yaoBase, "http://", "ws://", 1) + base = strings.Replace(base, "https://", "wss://", 1) + return fmt.Sprintf("%s/tai/%s/vnc/%s/ws", base, t.taiID, containerID), nil +} + +func (t *tunnelVNC) Ping(_ context.Context, _ string) error { + return nil +} + // --- Local implementation --- type localVNC struct {