diff --git a/openapi/openapi.go b/openapi/openapi.go index 4c20ccc4..49f2c78f 100644 --- a/openapi/openapi.go +++ b/openapi/openapi.go @@ -27,12 +27,12 @@ import ( "github.com/yaoapp/yao/openapi/otp" "github.com/yaoapp/yao/openapi/response" "github.com/yaoapp/yao/openapi/sandbox" + openapiTai "github.com/yaoapp/yao/openapi/tai" "github.com/yaoapp/yao/openapi/team" openapiTrace "github.com/yaoapp/yao/openapi/trace" "github.com/yaoapp/yao/openapi/user" openapiWorkspace "github.com/yaoapp/yao/openapi/workspace" taiapi "github.com/yaoapp/yao/tai/api" - taitunnel "github.com/yaoapp/yao/tai/tunnel" ) // Server is the OpenAPI server @@ -191,9 +191,8 @@ func (openapi *OpenAPI) Attach(router *gin.Engine) { // Tai nodes handlers nodes.Attach(group.Group("/nodes"), openapi.OAuth) - // Tai tunnel: gRPC Forward-based HTTP/VNC transparent proxy - group.Any("/tai/:taiID/proxy/*path", taitunnel.HandleForwardLazy) - group.Any("/tai/:taiID/vnc/*path", taitunnel.HandleForwardLazy) + // Tai forward handlers (proxy + VNC, dispatches tunnel vs local) + openapiTai.Attach(group) // Tai direct registration API (uses /tai-nodes/ prefix to avoid routing conflict with /tai/:taiID/) group.POST("/tai-nodes/register", taiapi.HandleRegister) diff --git a/openapi/tai/proxy.go b/openapi/tai/proxy.go new file mode 100644 index 00000000..dabf9bef --- /dev/null +++ b/openapi/tai/proxy.go @@ -0,0 +1,60 @@ +package tai + +import ( + "net/http" + "strings" + + "github.com/gin-gonic/gin" + yaoTai "github.com/yaoapp/yao/tai" +) + +// handleLocalProxy resolves the container's HTTP address via Docker socket +// and reverse-proxies the request. +func handleLocalProxy(c *gin.Context, taiID string) { + res, ok := yaoTai.GetResources(taiID) + if !ok || res.Proxy == nil { + c.JSON(http.StatusBadGateway, gin.H{"error": "proxy not available for node " + taiID}) + return + } + + // path format: /{containerID}:{port}/{rest...} + raw := strings.TrimPrefix(c.Param("path"), "/") + colonIdx := strings.Index(raw, ":") + if colonIdx < 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid proxy path, expected /{containerID}:{port}/{path}"}) + return + } + + containerID := raw[:colonIdx] + rest := raw[colonIdx+1:] + slashIdx := strings.Index(rest, "/") + var portStr, subPath string + if slashIdx >= 0 { + portStr = rest[:slashIdx] + subPath = rest[slashIdx:] + } else { + portStr = rest + subPath = "/" + } + + var port int + for _, ch := range portStr { + if ch < '0' || ch > '9' { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid port in proxy path"}) + return + } + port = port*10 + int(ch-'0') + } + if port == 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "missing port in proxy path"}) + return + } + + targetURL, err := res.Proxy.URL(c.Request.Context(), containerID, port, subPath) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": "resolve proxy target: " + err.Error()}) + return + } + + reverseProxy(c, targetURL) +} diff --git a/openapi/tai/tai.go b/openapi/tai/tai.go new file mode 100644 index 00000000..6ddaeba5 --- /dev/null +++ b/openapi/tai/tai.go @@ -0,0 +1,39 @@ +package tai + +import ( + "github.com/gin-gonic/gin" + yaoTai "github.com/yaoapp/yao/tai" + taitunnel "github.com/yaoapp/yao/tai/tunnel" +) + +// Attach registers Tai forward routes on the given group. +// +// - ANY /tai/:taiID/proxy/*path — HTTP forward (tunnel or local) +// - GET /tai/:taiID/vnc/*path — VNC WebSocket forward (tunnel or local) +func Attach(group *gin.RouterGroup) { + group.Any("/tai/:taiID/proxy/*path", handleProxy) + group.GET("/tai/:taiID/vnc/*path", handleVNC) +} + +func handleProxy(c *gin.Context) { + taiID := c.Param("taiID") + if isLocalNode(taiID) { + handleLocalProxy(c, taiID) + return + } + taitunnel.HandleForwardLazy(c) +} + +func handleVNC(c *gin.Context) { + taiID := c.Param("taiID") + if isLocalNode(taiID) { + handleLocalVNC(c, taiID) + return + } + taitunnel.HandleForwardLazy(c) +} + +func isLocalNode(taiID string) bool { + meta, ok := yaoTai.GetNodeMeta(taiID) + return ok && meta.Mode == "local" +} diff --git a/openapi/tai/util.go b/openapi/tai/util.go new file mode 100644 index 00000000..30ae9a5d --- /dev/null +++ b/openapi/tai/util.go @@ -0,0 +1,95 @@ +package tai + +import ( + "net/http" + "strings" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" +) + +// extractContainerID parses container ID from *path param. +// /{containerID}/ws → containerID +func extractContainerID(path string) string { + path = strings.TrimPrefix(path, "/") + path = strings.TrimSuffix(path, "/ws") + path = strings.TrimSuffix(path, "/") + if path == "" || path == "__host__" { + return "__host__" + } + return path +} + +// bridgeWebSocket copies messages bidirectionally between two WebSocket connections. +func bridgeWebSocket(client, target *websocket.Conn) { + done := make(chan struct{}, 2) + + go func() { + defer func() { done <- struct{}{} }() + for { + mt, data, err := client.ReadMessage() + if err != nil { + return + } + if err := target.WriteMessage(mt, data); err != nil { + return + } + } + }() + + go func() { + defer func() { done <- struct{}{} }() + for { + mt, data, err := target.ReadMessage() + if err != nil { + return + } + if err := client.WriteMessage(mt, data); err != nil { + return + } + } + }() + + <-done +} + +// reverseProxy forwards an HTTP request to targetURL and streams the response back. +func reverseProxy(c *gin.Context, targetURL string) { + req, err := http.NewRequestWithContext(c.Request.Context(), c.Request.Method, targetURL, c.Request.Body) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "create proxy request: " + err.Error()}) + return + } + for k, vv := range c.Request.Header { + for _, v := range vv { + req.Header.Add(k, v) + } + } + + resp, err := http.DefaultClient.Do(req) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": "proxy request failed: " + err.Error()}) + 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) + c.Writer.Flush() + + buf := make([]byte, 32*1024) + for { + n, readErr := resp.Body.Read(buf) + if n > 0 { + c.Writer.Write(buf[:n]) + c.Writer.Flush() + } + if readErr != nil { + return + } + } +} diff --git a/openapi/tai/vnc.go b/openapi/tai/vnc.go new file mode 100644 index 00000000..472aee39 --- /dev/null +++ b/openapi/tai/vnc.go @@ -0,0 +1,57 @@ +package tai + +import ( + "net/http" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + yaoTai "github.com/yaoapp/yao/tai" +) + +var wsUpgrader = websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, + Subprotocols: []string{"binary"}, +} + +// handleLocalVNC resolves the container's VNC address via Docker socket +// and proxies the WebSocket connection. +func handleLocalVNC(c *gin.Context, taiID string) { + containerID := extractContainerID(c.Param("path")) + if containerID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "missing container ID in path"}) + return + } + + res, ok := yaoTai.GetResources(taiID) + if !ok || res.VNC == nil { + c.JSON(http.StatusBadGateway, gin.H{"error": "VNC not available for node " + taiID}) + return + } + + targetURL, err := res.VNC.URL(c.Request.Context(), containerID) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": "resolve VNC target: " + err.Error()}) + return + } + + clientConn, err := wsUpgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + return + } + defer clientConn.Close() + + dialer := websocket.Dialer{ + Subprotocols: []string{"binary"}, + HandshakeTimeout: 5 * time.Second, + } + targetConn, _, err := dialer.Dial(targetURL, nil) + if err != nil { + clientConn.WriteMessage(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseInternalServerErr, "VNC connection failed")) + return + } + defer targetConn.Close() + + bridgeWebSocket(clientConn, targetConn) +}