diff --git a/Makefile b/Makefile index 88a19449..cd762aa4 100644 --- a/Makefile +++ b/Makefile @@ -10,9 +10,9 @@ NOW := $(shell date +"%FT%T%z") OS := $(shell uname) # ROOT_DIR := $(shell dirname $(realpath $(firstword $(MAKEFILE_LIST)))) -TESTFOLDER := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|registry|agent/sandbox/v2' | awk '!/\/tests\// || /openapi\/tests/') +TESTFOLDER := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|registry|agent/sandbox/v2' | awk '!/\/tests\// || /openapi\/tests/' | grep -vE 'openapi/tests/(nodes|sandbox|workspace)') # Core tests (exclude AI-related: agent, aigc, openai, KB, sandbox, registry, grpc, and integrations which require external services) -TESTFOLDER_CORE := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|agent|kb|sandbox|integrations|registry|tai|grpc' | awk '!/\/tests\// || /openapi\/tests/') +TESTFOLDER_CORE := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|agent|kb|sandbox|integrations|registry|tai|grpc' | awk '!/\/tests\// || /openapi\/tests/' | grep -vE 'openapi/tests/(nodes|sandbox|workspace)') # Agent tests (agent, aigc) - exclude agent/search/handlers/web (requires external API keys), robot packages (tested in robot job), and agent/sandbox/v2 (WIP, has its own job) TESTFOLDER_AGENT := $(shell $(GO) list ./agent/... ./aigc/... | grep -vE 'agent/search/handlers/web|agent/robot/|agent/sandbox/v2') # KB tests (kb) diff --git a/agent/store/types/sandbox_v2.go b/agent/store/types/sandbox_v2.go index 88485574..9968d613 100644 --- a/agent/store/types/sandbox_v2.go +++ b/agent/store/types/sandbox_v2.go @@ -2,6 +2,7 @@ package types import ( "crypto/sha256" + "encoding/json" "fmt" "os" "path/filepath" @@ -72,11 +73,11 @@ func ToSandboxV2(v any) (*sandboxTypes.SandboxConfig, error) { func ComputeConfigHash(cfg *sandboxTypes.SandboxConfig, mcpServers []MCPServerConfig, skillsDir string) string { h := sha256.New() - raw, _ := jsoniter.Marshal(cfg) + raw, _ := json.Marshal(cfg) h.Write(raw) if len(mcpServers) > 0 { - mcpRaw, _ := jsoniter.Marshal(mcpServers) + mcpRaw, _ := json.Marshal(mcpServers) h.Write(mcpRaw) } diff --git a/openapi/nodes/nodes.go b/openapi/nodes/nodes.go new file mode 100644 index 00000000..7d548ae8 --- /dev/null +++ b/openapi/nodes/nodes.go @@ -0,0 +1,138 @@ +package nodes + +import ( + "net/http" + "time" + + "github.com/gin-gonic/gin" + "github.com/yaoapp/yao/openapi/oauth/authorized" + "github.com/yaoapp/yao/openapi/oauth/types" + "github.com/yaoapp/yao/openapi/response" + "github.com/yaoapp/yao/tai/registry" +) + +// Attach registers Tai node endpoints on the given group. +// - GET / — list nodes (filtered by team/user from token) +// - GET /:id — get single node (owner check) +func Attach(group *gin.RouterGroup, oauth types.OAuth) { + group.Use(oauth.Guard) + group.GET("", handleList) + group.GET("/:id", handleGet) +} + +type nodeResponse struct { + TaiID string `json:"tai_id"` + MachineID string `json:"machine_id,omitempty"` + Version string `json:"version,omitempty"` + DisplayName string `json:"display_name,omitempty"` + Mode string `json:"mode"` + Addr string `json:"addr,omitempty"` + Status string `json:"status"` + System systemResponse `json:"system"` + Capabilities map[string]bool `json:"capabilities,omitempty"` + Ports map[string]int `json:"ports,omitempty"` + ConnectedAt *time.Time `json:"connected_at,omitempty"` + LastPing *time.Time `json:"last_ping,omitempty"` +} + +type systemResponse struct { + OS string `json:"os"` + Arch string `json:"arch"` + Hostname string `json:"hostname"` + NumCPU int `json:"num_cpu"` + TotalMem int64 `json:"total_mem,omitempty"` + Shell string `json:"shell,omitempty"` +} + +func snapToResponse(s registry.NodeSnapshot) nodeResponse { + r := nodeResponse{ + TaiID: s.TaiID, + MachineID: s.MachineID, + Version: s.Version, + DisplayName: s.DisplayName, + Mode: s.Mode, + Addr: s.Addr, + Status: s.Status, + Capabilities: s.Capabilities, + Ports: s.Ports, + System: systemResponse{ + OS: s.System.OS, + Arch: s.System.Arch, + Hostname: s.System.Hostname, + NumCPU: s.System.NumCPU, + TotalMem: s.System.TotalMem, + Shell: s.System.Shell, + }, + } + if !s.ConnectedAt.IsZero() { + r.ConnectedAt = &s.ConnectedAt + } + if !s.LastPing.IsZero() { + r.LastPing = &s.LastPing + } + return r +} + +// nodeOwnedBy checks whether a node belongs to the caller. +// TeamID match → true; no team and UserID match → true. +func nodeOwnedBy(snap *registry.NodeSnapshot, authInfo *types.AuthorizedInfo) bool { + if authInfo == nil { + return true + } + + if authInfo.TeamID != "" { + return snap.Auth.TeamID == authInfo.TeamID + } + if authInfo.UserID != "" { + return snap.Auth.TeamID == "" && snap.Auth.UserID == authInfo.UserID + } + return true +} + +func handleList(c *gin.Context) { + reg := registry.Global() + if reg == nil { + response.RespondWithSuccess(c, http.StatusOK, []nodeResponse{}) + return + } + + authInfo := authorized.GetInfo(c) + + var snaps []registry.NodeSnapshot + if authInfo != nil && authInfo.TeamID != "" { + snaps = reg.ListByTeam(authInfo.TeamID) + } else if authInfo != nil && authInfo.UserID != "" { + snaps = reg.ListByUser(authInfo.UserID) + } else { + snaps = reg.List() + } + + result := make([]nodeResponse, 0, len(snaps)) + for _, s := range snaps { + result = append(result, snapToResponse(s)) + } + response.RespondWithSuccess(c, http.StatusOK, result) +} + +func handleGet(c *gin.Context) { + reg := registry.Global() + if reg == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "node registry not available"}) + return + } + + id := c.Param("id") + snap, ok := reg.Get(id) + if !ok { + c.JSON(http.StatusNotFound, gin.H{"error": "node not found"}) + return + } + + authInfo := authorized.GetInfo(c) + if !nodeOwnedBy(snap, authInfo) { + c.JSON(http.StatusForbidden, gin.H{"error": "no permission to access this node"}) + return + } + + response.RespondWithSuccess(c, http.StatusOK, snapToResponse(*snap)) +} diff --git a/openapi/oauth.go b/openapi/oauth.go index 52f19b1e..94c23cd5 100644 --- a/openapi/oauth.go +++ b/openapi/oauth.go @@ -2,6 +2,7 @@ package openapi import ( "encoding/base64" + "fmt" "net/http" "strings" @@ -585,11 +586,32 @@ func (openapi *OpenAPI) oauthDeviceAuthorize(c *gin.Context) { if extraClaims == nil { extraClaims = make(map[string]interface{}) } - if tokenClaims.TeamID != "" { - extraClaims["team_id"] = tokenClaims.TeamID + + teamID := tokenClaims.TeamID + if teamID == "" { + switch v := extraClaims["team_id"].(type) { + case string: + teamID = v + case float64: + teamID = fmt.Sprintf("%.0f", v) + } } - if tokenClaims.TenantID != "" { - extraClaims["tenant_id"] = tokenClaims.TenantID + if teamID != "" { + extraClaims["team_id"] = teamID + } + + tenantID := tokenClaims.TenantID + if tenantID == "" { + if v, ok := extraClaims["tenant_id"].(string); ok { + tenantID = v + } + } + if tenantID != "" { + extraClaims["tenant_id"] = tenantID + } + + if tokenClaims.ClientID != "" { + extraClaims["authorizer_client_id"] = tokenClaims.ClientID } userCode := c.PostForm("user_code") diff --git a/openapi/oauth/authenticate.go b/openapi/oauth/authenticate.go index 288ff2ad..16937e77 100644 --- a/openapi/oauth/authenticate.go +++ b/openapi/oauth/authenticate.go @@ -189,13 +189,29 @@ func (s *Service) refreshTokenDirect(refreshToken string, expiredClaims *types.T // buildAuthInfo constructs AuthorizedInfo directly from token claims, // equivalent to the SetInfo+GetInfo round-trip through gin.Context. func (s *Service) buildAuthInfo(claims *types.TokenClaims, sessionID string) *types.AuthorizedInfo { + teamID := claims.TeamID + tenantID := claims.TenantID + + if claims.Extra != nil { + if teamID == "" { + if v, ok := claims.Extra["team_id"].(string); ok && v != "" { + teamID = v + } + } + if tenantID == "" { + if v, ok := claims.Extra["tenant_id"].(string); ok && v != "" { + tenantID = v + } + } + } + info := &types.AuthorizedInfo{ Subject: claims.Subject, ClientID: claims.ClientID, Scope: claims.Scope, SessionID: sessionID, - TeamID: claims.TeamID, - TenantID: claims.TenantID, + TeamID: teamID, + TenantID: tenantID, } userID, err := s.UserID(claims.ClientID, claims.Subject) @@ -203,5 +219,14 @@ func (s *Service) buildAuthInfo(claims *types.TokenClaims, sessionID string) *ty info.UserID = userID } + if info.UserID == "" && claims.Extra != nil { + if authorizerClientID, ok := claims.Extra["authorizer_client_id"].(string); ok && authorizerClientID != "" { + if uid, err := s.UserID(authorizerClientID, claims.Subject); err == nil && uid != "" { + info.UserID = uid + s.copyFingerprint(authorizerClientID, claims.ClientID, claims.Subject) + } + } + } + return info } diff --git a/openapi/oauth/core.go b/openapi/oauth/core.go index 9862a5ee..f6e4fc34 100644 --- a/openapi/oauth/core.go +++ b/openapi/oauth/core.go @@ -241,9 +241,11 @@ func (s *Service) RefreshToken(ctx context.Context, refreshToken string, scope . finalScope = requestedScope } + extraClaims := extractExtraClaims(tokenInfo) + // Generate new access token with final scope expiresIn := int(s.config.Token.AccessTokenLifetime.Seconds()) - newAccessToken, err := s.generateAccessTokenWithScope(clientID, finalScope, originalSubject, expiresIn, nil) + newAccessToken, err := s.generateAccessTokenWithScope(clientID, finalScope, originalSubject, expiresIn, extraClaims) if err != nil { return nil, &types.ErrorResponse{ Code: types.ErrorServerError, @@ -340,9 +342,11 @@ func (s *Service) RotateRefreshToken(ctx context.Context, oldToken string, reque finalScope = scope } + extraClaims := extractExtraClaims(tokenInfo) + // Generate new tokens with final scope and original subject expiresIn := int(s.config.Token.AccessTokenLifetime.Seconds()) - newAccessToken, err := s.generateAccessTokenWithScope(clientID, finalScope, originalSubject, expiresIn, nil) + newAccessToken, err := s.generateAccessTokenWithScope(clientID, finalScope, originalSubject, expiresIn, extraClaims) if err != nil { return nil, &types.ErrorResponse{ Code: types.ErrorServerError, @@ -350,7 +354,7 @@ func (s *Service) RotateRefreshToken(ctx context.Context, oldToken string, reque } } - newRefreshToken, err := s.generateRefreshToken(clientID, finalScope, originalSubject, 0, nil) + newRefreshToken, err := s.generateRefreshToken(clientID, finalScope, originalSubject, 0, extraClaims) if err != nil { return nil, &types.ErrorResponse{ Code: types.ErrorServerError, @@ -497,20 +501,12 @@ func (s *Service) handleRefreshTokenGrant(ctx context.Context, client *types.Cli return nil, err } - // Extract scope and subject from refresh token if available - scope := "" - if scopeVal, ok := refreshTokenInfo["scope"].(string); ok { - scope = scopeVal - } + scope, _ := refreshTokenInfo["scope"].(string) + subject, _ := refreshTokenInfo["subject"].(string) + extraClaims := extractExtraClaims(refreshTokenInfo) - subject := "" - if subjectVal, ok := refreshTokenInfo["subject"].(string); ok { - subject = subjectVal - } - - // Generate and store new access token with proper scope and subject expiresIn := int(s.config.Token.AccessTokenLifetime.Seconds()) - accessToken, err := s.generateAccessTokenWithScope(client.ClientID, scope, subject, expiresIn, nil) + accessToken, err := s.generateAccessTokenWithScope(client.ClientID, scope, subject, expiresIn, extraClaims) if err != nil { return nil, &types.ErrorResponse{ Code: types.ErrorServerError, @@ -524,9 +520,8 @@ func (s *Service) handleRefreshTokenGrant(ctx context.Context, client *types.Cli ExpiresIn: expiresIn, } - // Include refresh token if rotation is enabled if s.config.Features.RefreshTokenRotationEnabled { - newRefreshToken, err := s.generateRefreshToken(client.ClientID, scope, subject, 0, nil) + newRefreshToken, err := s.generateRefreshToken(client.ClientID, scope, subject, 0, extraClaims) if err != nil { return nil, &types.ErrorResponse{ Code: types.ErrorServerError, @@ -534,11 +529,8 @@ func (s *Service) handleRefreshTokenGrant(ctx context.Context, client *types.Cli } } token.RefreshToken = newRefreshToken - - // Revoke old refresh token s.revokeRefreshToken(refreshToken) } else { - // Reuse the same refresh token token.RefreshToken = refreshToken } @@ -713,3 +705,23 @@ func (s *Service) handleDeviceCodeGrant(ctx context.Context, client *types.Clien } } } + +// extractExtraClaims pulls non-reserved fields from a token info map so they +// can be propagated into newly generated access/refresh tokens. +func extractExtraClaims(tokenInfo map[string]interface{}) map[string]interface{} { + reserved := map[string]bool{ + "client_id": true, "scope": true, "subject": true, + "type": true, "issued_at": true, "expires_at": true, + } + var extra map[string]interface{} + for k, v := range tokenInfo { + if reserved[k] { + continue + } + if extra == nil { + extra = make(map[string]interface{}) + } + extra[k] = v + } + return extra +} diff --git a/openapi/oauth/token.go b/openapi/oauth/token.go index 5d8b52d3..bb4f3428 100644 --- a/openapi/oauth/token.go +++ b/openapi/oauth/token.go @@ -320,6 +320,22 @@ func (s *Service) UserID(clientID, subject string) (string, error) { return userIDStr, nil } +// copyFingerprint copies the subject→userID fingerprint mapping from one +// clientID to another so that tokens issued under a different clientID +// (e.g. Device Flow) can resolve the same userID. +func (s *Service) copyFingerprint(srcClientID, dstClientID, subject string) { + srcKey := s.userFingerprintKey(srcClientID, subject) + userID, exists := s.store.Get(srcKey) + if !exists { + return + } + dstKey := s.userFingerprintKey(dstClientID, subject) + if _, already := s.store.Get(dstKey); already { + return + } + s.store.Set(dstKey, userID, 0) +} + // MakeAuthorizationCode generates a new authorization code with specific parameters and stores it // ============================================================================ diff --git a/openapi/openapi.go b/openapi/openapi.go index 23e9c108..dd5c6c9d 100644 --- a/openapi/openapi.go +++ b/openapi/openapi.go @@ -19,6 +19,7 @@ import ( "github.com/yaoapp/yao/openapi/llm" "github.com/yaoapp/yao/openapi/mcp" "github.com/yaoapp/yao/openapi/messenger" + "github.com/yaoapp/yao/openapi/nodes" "github.com/yaoapp/yao/openapi/oauth" "github.com/yaoapp/yao/openapi/oauth/acl" "github.com/yaoapp/yao/openapi/oauth/types" @@ -28,6 +29,7 @@ import ( "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" ) @@ -173,9 +175,17 @@ func (openapi *OpenAPI) Attach(router *gin.Engine) { // OTP handlers (passwordless authentication) otp.Attach(group.Group("/otp"), openapi.OAuth) - // Sandbox handlers (VNC proxy for visual browser automation) + // Sandbox handlers (VNC proxy + management CRUD) sandbox.SetPathPrefix(baseURL) - sandbox.Attach(group.Group("/sandbox"), openapi.OAuth) + sandboxGroup := group.Group("/sandbox") + sandbox.Attach(sandboxGroup, openapi.OAuth) + sandbox.AttachManage(sandboxGroup) + + // Workspace handlers + openapiWorkspace.Attach(group.Group("/workspace"), openapi.OAuth) + + // Tai nodes handlers + nodes.Attach(group.Group("/nodes"), openapi.OAuth) // Tai tunnel WebSocket and reverse proxy routes group.GET("/ws/tai", taitunnel.HandleControl) diff --git a/openapi/sandbox/manage.go b/openapi/sandbox/manage.go new file mode 100644 index 00000000..74057930 --- /dev/null +++ b/openapi/sandbox/manage.go @@ -0,0 +1,366 @@ +package sandbox + +import ( + "context" + "net/http" + "time" + + "github.com/gin-gonic/gin" + "github.com/yaoapp/yao/openapi/oauth/authorized" + "github.com/yaoapp/yao/openapi/oauth/types" + "github.com/yaoapp/yao/openapi/response" + sandboxv2 "github.com/yaoapp/yao/sandbox/v2" +) + +// AttachManage registers sandbox management CRUD routes on the given group. +// oauth.Guard is already applied by the parent Attach call on the same group. +// - GET / — list sandboxes (filtered by owner) +// - POST / — create sandbox (owner from token) +// - GET /:id — get sandbox (owner check) +// - DELETE /:id — remove sandbox (owner check) +// - POST /:id/exec — execute command (owner check) +// - POST /:id/heartbeat — heartbeat (owner check) +func AttachManage(group *gin.RouterGroup) { + group.GET("", handleList) + group.POST("", handleCreate) + group.GET("/:id", handleGet) + group.DELETE("/:id", handleRemove) + group.POST("/:id/exec", handleExec) + group.POST("/:id/heartbeat", handleHeartbeat) +} + +// resolveOwner returns TeamID if present, otherwise UserID. +func resolveOwner(authInfo *types.AuthorizedInfo) string { + if authInfo != nil && authInfo.TeamID != "" { + return authInfo.TeamID + } + if authInfo != nil { + return authInfo.UserID + } + return "" +} + +// --- request / response types --- + +type createSandboxRequest struct { + ID string `json:"id,omitempty"` + NodeID string `json:"node_id"` + Image string `json:"image"` + WorkDir string `json:"work_dir,omitempty"` + User string `json:"user,omitempty"` + Env map[string]string `json:"env,omitempty"` + Memory int64 `json:"memory,omitempty"` + CPUs float64 `json:"cpus,omitempty"` + VNC bool `json:"vnc,omitempty"` + Policy string `json:"policy,omitempty"` + Labels map[string]string `json:"labels,omitempty"` + WorkspaceID string `json:"workspace_id,omitempty"` + MountMode string `json:"mount_mode,omitempty"` + MountPath string `json:"mount_path,omitempty"` +} + +type execRequest struct { + Cmd []string `json:"cmd" binding:"required"` + WorkDir string `json:"work_dir,omitempty"` + Env map[string]string `json:"env,omitempty"` + Timeout int `json:"timeout,omitempty"` +} + +type heartbeatRequest struct { + Active bool `json:"active"` + ProcessCount int `json:"process_count"` +} + +type sandboxSystemInfo struct { + OS string `json:"os"` + Arch string `json:"arch"` + Hostname string `json:"hostname"` + NumCPU int `json:"num_cpu"` + TotalMem int64 `json:"total_mem,omitempty"` + Shell string `json:"shell,omitempty"` + TempDir string `json:"temp_dir,omitempty"` +} + +type sandboxResponse struct { + ID string `json:"id"` + ContainerID string `json:"container_id"` + NodeID string `json:"node_id"` + Owner string `json:"owner"` + Status string `json:"status"` + Policy string `json:"policy"` + Labels map[string]string `json:"labels,omitempty"` + Image string `json:"image"` + VNC bool `json:"vnc"` + CreatedAt time.Time `json:"created_at"` + LastActive time.Time `json:"last_active"` + ProcessCount int `json:"process_count"` + System sandboxSystemInfo `json:"system"` + WorkspaceID string `json:"workspace_id,omitempty"` +} + +func boxToResponse(b *sandboxv2.Box) sandboxResponse { + snap := b.Snapshot() + info := b.ComputerInfo() + return sandboxResponse{ + ID: snap.ID, + ContainerID: snap.ContainerID, + NodeID: snap.NodeID, + Owner: snap.Owner, + Status: snap.Status, + Policy: string(snap.Policy), + Labels: snap.Labels, + Image: snap.Image, + VNC: snap.VNC, + CreatedAt: snap.CreatedAt, + LastActive: snap.LastActive, + ProcessCount: snap.ProcessCount, + WorkspaceID: b.WorkspaceID(), + System: sandboxSystemInfo{ + OS: info.System.OS, + Arch: info.System.Arch, + Hostname: info.System.Hostname, + NumCPU: info.System.NumCPU, + TotalMem: info.System.TotalMem, + Shell: info.System.Shell, + TempDir: info.System.TempDir, + }, + } +} + +func getManager(c *gin.Context) *sandboxv2.Manager { + defer func() { recover() }() + return sandboxv2.M() +} + +// checkBoxOwner verifies the caller owns the sandbox. +func checkBoxOwner(c *gin.Context, box *sandboxv2.Box, owner string) bool { + if owner == "" { + return true + } + info := box.ComputerInfo() + if info.Owner != "" && info.Owner != owner { + c.JSON(http.StatusForbidden, gin.H{"error": "no permission to access this sandbox"}) + return false + } + return true +} + +// --- handlers --- + +func handleList(c *gin.Context) { + mgr := getManager(c) + if mgr == nil { + response.RespondWithSuccess(c, http.StatusOK, []sandboxResponse{}) + return + } + + authInfo := authorized.GetInfo(c) + owner := resolveOwner(authInfo) + + boxes, err := mgr.List(context.Background(), sandboxv2.ListOptions{ + Owner: owner, + NodeID: c.Query("node_id"), + }) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + result := make([]sandboxResponse, 0, len(boxes)) + for _, b := range boxes { + result = append(result, boxToResponse(b)) + } + response.RespondWithSuccess(c, http.StatusOK, result) +} + +func handleCreate(c *gin.Context) { + mgr := getManager(c) + if mgr == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "sandbox service not available"}) + return + } + + var req createSandboxRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + if req.Image == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "image is required"}) + return + } + + authInfo := authorized.GetInfo(c) + owner := resolveOwner(authInfo) + + opts := sandboxv2.CreateOptions{ + ID: req.ID, + Owner: owner, + NodeID: req.NodeID, + Image: req.Image, + WorkDir: req.WorkDir, + User: req.User, + Env: req.Env, + Memory: req.Memory, + CPUs: req.CPUs, + VNC: req.VNC, + Policy: sandboxv2.LifecyclePolicy(req.Policy), + Labels: req.Labels, + WorkspaceID: req.WorkspaceID, + MountMode: req.MountMode, + MountPath: req.MountPath, + } + + box, err := mgr.Create(context.Background(), opts) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + response.RespondWithSuccess(c, http.StatusCreated, boxToResponse(box)) +} + +func handleGet(c *gin.Context) { + mgr := getManager(c) + if mgr == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "sandbox service not available"}) + return + } + + id := c.Param("id") + box, err := mgr.Get(context.Background(), id) + if err != nil { + if err == sandboxv2.ErrNotFound { + c.JSON(http.StatusNotFound, gin.H{"error": "sandbox not found"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + authInfo := authorized.GetInfo(c) + if !checkBoxOwner(c, box, resolveOwner(authInfo)) { + return + } + + response.RespondWithSuccess(c, http.StatusOK, boxToResponse(box)) +} + +func handleRemove(c *gin.Context) { + mgr := getManager(c) + if mgr == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "sandbox service not available"}) + return + } + + id := c.Param("id") + box, err := mgr.Get(context.Background(), id) + if err != nil { + if err == sandboxv2.ErrNotFound { + c.JSON(http.StatusNotFound, gin.H{"error": "sandbox not found"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + authInfo := authorized.GetInfo(c) + if !checkBoxOwner(c, box, resolveOwner(authInfo)) { + return + } + + if err := mgr.Remove(context.Background(), id); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.Status(http.StatusNoContent) +} + +func handleExec(c *gin.Context) { + mgr := getManager(c) + if mgr == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "sandbox service not available"}) + return + } + + id := c.Param("id") + box, err := mgr.Get(context.Background(), id) + if err != nil { + if err == sandboxv2.ErrNotFound { + c.JSON(http.StatusNotFound, gin.H{"error": "sandbox not found"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + authInfo := authorized.GetInfo(c) + if !checkBoxOwner(c, box, resolveOwner(authInfo)) { + return + } + + var req execRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + var opts []sandboxv2.ExecOption + if req.WorkDir != "" { + opts = append(opts, sandboxv2.WithWorkDir(req.WorkDir)) + } + if len(req.Env) > 0 { + opts = append(opts, sandboxv2.WithEnv(req.Env)) + } + if req.Timeout > 0 { + opts = append(opts, sandboxv2.WithTimeout(time.Duration(req.Timeout)*time.Second)) + } + + result, err := box.Exec(context.Background(), req.Cmd, opts...) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + response.RespondWithSuccess(c, http.StatusOK, result) +} + +func handleHeartbeat(c *gin.Context) { + mgr := getManager(c) + if mgr == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "sandbox service not available"}) + return + } + + id := c.Param("id") + + box, err := mgr.Get(context.Background(), id) + if err != nil { + if err == sandboxv2.ErrNotFound { + c.JSON(http.StatusNotFound, gin.H{"error": "sandbox not found"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + authInfo := authorized.GetInfo(c) + if !checkBoxOwner(c, box, resolveOwner(authInfo)) { + return + } + + var req heartbeatRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + if err := mgr.Heartbeat(id, req.Active, req.ProcessCount); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.Status(http.StatusNoContent) +} diff --git a/openapi/tests/nodes/nodes_test.go b/openapi/tests/nodes/nodes_test.go new file mode 100644 index 00000000..b1dfc19c --- /dev/null +++ b/openapi/tests/nodes/nodes_test.go @@ -0,0 +1,91 @@ +package openapi_test + +import ( + "encoding/json" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +func TestNodesListAuthenticated(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + client := testutils.RegisterTestClient(t, "Nodes Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + req, err := http.NewRequest("GET", serverURL+baseURL+"/nodes", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var result []map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&result) + assert.NoError(t, err) + assert.NotNil(t, result) + t.Logf("Nodes list returned %d items", len(result)) + + for _, node := range result { + assert.NotEmpty(t, node["tai_id"], "node should have tai_id") + assert.NotEmpty(t, node["mode"], "node should have mode") + assert.NotEmpty(t, node["status"], "node should have status") + t.Logf("Node: tai_id=%s, mode=%s, status=%s", node["tai_id"], node["mode"], node["status"]) + } +} + +func TestNodesListUnauthorized(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + resp, err := http.Get(serverURL + baseURL + "/nodes") + assert.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) +} + +func TestNodesGetNotFound(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + client := testutils.RegisterTestClient(t, "Nodes Get Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + req, err := http.NewRequest("GET", serverURL+baseURL+"/nodes/nonexistent-node", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + + assert.True(t, resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusServiceUnavailable, + "expected 404 or 503, got %d", resp.StatusCode) +} diff --git a/openapi/tests/sandbox/sandbox_test.go b/openapi/tests/sandbox/sandbox_test.go new file mode 100644 index 00000000..7f3fedca --- /dev/null +++ b/openapi/tests/sandbox/sandbox_test.go @@ -0,0 +1,150 @@ +package openapi_test + +import ( + "encoding/json" + "net/http" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +func TestSandboxListPublicDenied(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + resp, err := http.Get(serverURL + baseURL + "/sandbox") + assert.NoError(t, err) + defer resp.Body.Close() + + // Without auth, sandbox list returns 200 (scopes.yml allows GET /sandbox/*) + // but since /sandbox (no trailing wildcard match) could be denied or allowed, + // check that a response is returned. + assert.True(t, resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusUnauthorized, + "expected 200 or 401, got %d", resp.StatusCode) +} + +func TestSandboxListAuthenticated(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + client := testutils.RegisterTestClient(t, "Sandbox Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + req, err := http.NewRequest("GET", serverURL+baseURL+"/sandbox", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var result []map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&result) + assert.NoError(t, err) + assert.NotNil(t, result) + t.Logf("Sandbox list returned %d items", len(result)) +} + +func TestSandboxGetNotFound(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + client := testutils.RegisterTestClient(t, "Sandbox NotFound Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + req, err := http.NewRequest("GET", serverURL+baseURL+"/sandbox/nonexistent-id", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + + // Either 404 (sandbox not found) or 503 (sandbox service not available) is acceptable + assert.True(t, resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusServiceUnavailable, + "expected 404 or 503, got %d", resp.StatusCode) +} + +func TestSandboxCreateMissingImage(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + client := testutils.RegisterTestClient(t, "Sandbox Create Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + body := `{"node_id": "local"}` + req, err := http.NewRequest("POST", serverURL+baseURL+"/sandbox", jsonBody(body)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + + // Should return 400 (image required) or 503 (service unavailable) + assert.True(t, resp.StatusCode == http.StatusBadRequest || resp.StatusCode == http.StatusServiceUnavailable, + "expected 400 or 503, got %d", resp.StatusCode) +} + +func TestSandboxDeleteNotFound(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + client := testutils.RegisterTestClient(t, "Sandbox Delete Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + req, err := http.NewRequest("DELETE", serverURL+baseURL+"/sandbox/nonexistent-id", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + + // Either 404 or 503 is acceptable + assert.True(t, resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusServiceUnavailable, + "expected 404 or 503, got %d", resp.StatusCode) +} + +func jsonBody(s string) *strings.Reader { + return strings.NewReader(s) +} diff --git a/openapi/tests/workspace/workspace_test.go b/openapi/tests/workspace/workspace_test.go new file mode 100644 index 00000000..d637f90a --- /dev/null +++ b/openapi/tests/workspace/workspace_test.go @@ -0,0 +1,110 @@ +package openapi_test + +import ( + "encoding/json" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +func TestWorkspaceListAuthenticated(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + client := testutils.RegisterTestClient(t, "Workspace Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + req, err := http.NewRequest("GET", serverURL+baseURL+"/workspace", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var result []map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&result) + assert.NoError(t, err) + assert.NotNil(t, result) + t.Logf("Workspace list returned %d items", len(result)) +} + +func TestWorkspaceListUnauthorized(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + resp, err := http.Get(serverURL + baseURL + "/workspace") + assert.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) +} + +func TestWorkspaceGetNotFound(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + client := testutils.RegisterTestClient(t, "Workspace Get Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + req, err := http.NewRequest("GET", serverURL+baseURL+"/workspace/nonexistent-ws", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + + assert.True(t, resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusServiceUnavailable, + "expected 404 or 503, got %d", resp.StatusCode) +} + +func TestWorkspaceDeleteNotFound(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + client := testutils.RegisterTestClient(t, "Workspace Delete Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + req, err := http.NewRequest("DELETE", serverURL+baseURL+"/workspace/nonexistent-ws", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + + assert.True(t, resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusServiceUnavailable, + "expected 404 or 503, got %d", resp.StatusCode) +} diff --git a/openapi/workspace/workspace.go b/openapi/workspace/workspace.go new file mode 100644 index 00000000..5a550199 --- /dev/null +++ b/openapi/workspace/workspace.go @@ -0,0 +1,388 @@ +package workspace + +import ( + "context" + "encoding/base64" + "fmt" + "io" + "mime" + "net/http" + "path/filepath" + + "github.com/gin-gonic/gin" + "github.com/yaoapp/yao/openapi/oauth/authorized" + "github.com/yaoapp/yao/openapi/oauth/types" + "github.com/yaoapp/yao/openapi/response" + ws "github.com/yaoapp/yao/workspace" +) + +// Attach registers workspace management routes on the given group. +// - GET / — list workspaces (filtered by owner from token) +// - POST / — create workspace (owner from token) +// - GET /:id — get workspace (owner check) +// - PUT /:id — update workspace (owner check) +// - DELETE /:id — delete workspace (owner check) +// - GET /:id/files — list files +// - GET /:id/files/*path — read file +// - PUT /:id/files/*path — write file +// - DELETE /:id/files/*path — delete file +// - POST /:id/mkdir — create directory +// - POST /:id/rename — rename file/directory +func Attach(group *gin.RouterGroup, oauth types.OAuth) { + group.Use(oauth.Guard) + + group.GET("", handleList) + group.POST("", handleCreate) + group.GET("/:id", handleGet) + group.PUT("/:id", handleUpdate) + group.DELETE("/:id", handleDelete) + + group.GET("/:id/files", handleListFiles) + group.GET("/:id/files/*path", handleReadFile) + group.PUT("/:id/files/*path", handleWriteFile) + group.DELETE("/:id/files/*path", handleDeleteFile) + group.POST("/:id/mkdir", handleMkdir) + group.POST("/:id/rename", handleRename) +} + +// resolveOwner returns TeamID if present, otherwise UserID. +func resolveOwner(authInfo *types.AuthorizedInfo) string { + if authInfo != nil && authInfo.TeamID != "" { + return authInfo.TeamID + } + if authInfo != nil { + return authInfo.UserID + } + return "" +} + +// checkWSOwner verifies the caller owns the workspace. +func checkWSOwner(c *gin.Context, w *ws.Workspace, owner string) bool { + if owner == "" { + return true + } + if w.Owner != "" && w.Owner != owner { + c.JSON(http.StatusForbidden, gin.H{"error": "no permission to access this workspace"}) + return false + } + return true +} + +// --- request / response types --- + +type createRequest struct { + ID string `json:"id,omitempty"` + Name string `json:"name" binding:"required"` + Node string `json:"node" binding:"required"` + Labels map[string]string `json:"labels,omitempty"` +} + +type updateRequest struct { + Name *string `json:"name,omitempty"` + Labels map[string]string `json:"labels,omitempty"` +} + +type mkdirRequest struct { + Path string `json:"path" binding:"required"` +} + +type renameRequest struct { + OldPath string `json:"old_path" binding:"required"` + NewPath string `json:"new_path" binding:"required"` +} + +type workspaceResponse struct { + ID string `json:"id"` + Name string `json:"name"` + Owner string `json:"owner"` + Node string `json:"node"` + Labels map[string]string `json:"labels,omitempty"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` +} + +func toResponse(w *ws.Workspace) workspaceResponse { + return workspaceResponse{ + ID: w.ID, + Name: w.Name, + Owner: w.Owner, + Node: w.Node, + Labels: w.Labels, + CreatedAt: w.CreatedAt.Format("2006-01-02T15:04:05Z"), + UpdatedAt: w.UpdatedAt.Format("2006-01-02T15:04:05Z"), + } +} + +func mgr() *ws.Manager { + return ws.M() +} + +// resolveAndCheckWS fetches the workspace and verifies owner permission. +func resolveAndCheckWS(c *gin.Context) (*ws.Workspace, bool) { + m := mgr() + if m == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "workspace service not available"}) + return nil, false + } + + w, err := m.Get(context.Background(), c.Param("id")) + if err != nil { + if err == ws.ErrNotFound { + c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found"}) + return nil, false + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return nil, false + } + + authInfo := authorized.GetInfo(c) + if !checkWSOwner(c, w, resolveOwner(authInfo)) { + return nil, false + } + return w, true +} + +// --- handlers --- + +func handleList(c *gin.Context) { + m := mgr() + if m == nil { + response.RespondWithSuccess(c, http.StatusOK, []workspaceResponse{}) + return + } + + authInfo := authorized.GetInfo(c) + owner := resolveOwner(authInfo) + + list, err := m.List(context.Background(), ws.ListOptions{ + Owner: owner, + Node: c.Query("node"), + }) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + result := make([]workspaceResponse, 0, len(list)) + for _, w := range list { + result = append(result, toResponse(w)) + } + response.RespondWithSuccess(c, http.StatusOK, result) +} + +func handleCreate(c *gin.Context) { + m := mgr() + if m == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "workspace service not available"}) + return + } + + var req createRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + authInfo := authorized.GetInfo(c) + owner := resolveOwner(authInfo) + + w, err := m.Create(context.Background(), ws.CreateOptions{ + ID: req.ID, + Name: req.Name, + Owner: owner, + Node: req.Node, + Labels: req.Labels, + }) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + response.RespondWithSuccess(c, http.StatusCreated, toResponse(w)) +} + +func handleGet(c *gin.Context) { + w, ok := resolveAndCheckWS(c) + if !ok { + return + } + response.RespondWithSuccess(c, http.StatusOK, toResponse(w)) +} + +func handleUpdate(c *gin.Context) { + _, ok := resolveAndCheckWS(c) + if !ok { + return + } + + var req updateRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + w, err := mgr().Update(context.Background(), c.Param("id"), ws.UpdateOptions{ + Name: req.Name, + Labels: req.Labels, + }) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + response.RespondWithSuccess(c, http.StatusOK, toResponse(w)) +} + +func handleDelete(c *gin.Context) { + _, ok := resolveAndCheckWS(c) + if !ok { + return + } + + force := c.Query("force") == "true" + if err := mgr().Delete(context.Background(), c.Param("id"), force); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.Status(http.StatusNoContent) +} + +func handleListFiles(c *gin.Context) { + _, ok := resolveAndCheckWS(c) + if !ok { + return + } + + dir := c.DefaultQuery("path", ".") + entries, err := mgr().ListDir(context.Background(), c.Param("id"), dir) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + response.RespondWithSuccess(c, http.StatusOK, entries) +} + +func handleReadFile(c *gin.Context) { + _, ok := resolveAndCheckWS(c) + if !ok { + return + } + + path := c.Param("path") + if len(path) > 0 && path[0] == '/' { + path = path[1:] + } + + fmt.Printf("[workspace] handleReadFile id=%s path=%q\n", c.Param("id"), path) + + data, err := mgr().ReadFile(context.Background(), c.Param("id"), path) + if err != nil { + fmt.Printf("[workspace] ReadFile error: %v\n", err) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + fmt.Printf("[workspace] ReadFile ok, size=%d, encoding=%q\n", len(data), c.Query("encoding")) + + if c.Query("encoding") == "base64" { + response.RespondWithSuccess(c, http.StatusOK, gin.H{ + "content": base64.StdEncoding.EncodeToString(data), + "encoding": "base64", + }) + return + } + + ext := filepath.Ext(path) + mimeType := mime.TypeByExtension(ext) + if mimeType == "" { + mimeType = "application/octet-stream" + } + fmt.Printf("[workspace] serving ext=%q mime=%q size=%d\n", ext, mimeType, len(data)) + c.Data(http.StatusOK, mimeType, data) +} + +func handleWriteFile(c *gin.Context) { + _, ok := resolveAndCheckWS(c) + if !ok { + return + } + + path := c.Param("path") + if len(path) > 0 && path[0] == '/' { + path = path[1:] + } + + data, err := io.ReadAll(c.Request.Body) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "failed to read body"}) + return + } + + if err := mgr().WriteFile(context.Background(), c.Param("id"), path, data, 0644); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.Status(http.StatusNoContent) +} + +func handleDeleteFile(c *gin.Context) { + _, ok := resolveAndCheckWS(c) + if !ok { + return + } + + path := c.Param("path") + if len(path) > 0 && path[0] == '/' { + path = path[1:] + } + + if err := mgr().Remove(context.Background(), c.Param("id"), path); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.Status(http.StatusNoContent) +} + +func handleMkdir(c *gin.Context) { + _, ok := resolveAndCheckWS(c) + if !ok { + return + } + + var req mkdirRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + if err := mgr().MkdirAll(context.Background(), c.Param("id"), req.Path); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.Status(http.StatusNoContent) +} + +func handleRename(c *gin.Context) { + _, ok := resolveAndCheckWS(c) + if !ok { + return + } + + var req renameRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + if err := mgr().Rename(context.Background(), c.Param("id"), req.OldPath, req.NewPath); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.Status(http.StatusNoContent) +} diff --git a/sandbox/v2/box.go b/sandbox/v2/box.go index 5cf4223e..9b25013f 100644 --- a/sandbox/v2/box.go +++ b/sandbox/v2/box.go @@ -198,6 +198,25 @@ func (b *Box) Workspace() workspace.FS { // WorkspaceID returns the workspace ID mounted to this sandbox, or empty string. func (b *Box) WorkspaceID() string { return b.workspaceID } +// Snapshot returns a local-only BoxInfo snapshot without any remote calls. +// Status is inferred from local state (not from the container runtime). +func (b *Box) Snapshot() BoxInfo { + return BoxInfo{ + ID: b.id, + ContainerID: b.containerID, + NodeID: b.nodeID, + Owner: b.owner, + Status: "running", + Policy: b.policy, + Labels: b.labels, + Image: b.image, + CreatedAt: b.createdAt, + LastActive: b.lastActiveTime(), + ProcessCount: int(b.processCount.Load()), + VNC: b.vnc, + } +} + // VNC returns the VNC WebSocket URL. func (b *Box) VNC(ctx context.Context) (string, error) { b.touch() diff --git a/sandbox/v2/box_workspace_test.go b/sandbox/v2/box_workspace_test.go index aff1053f..6f5edcb8 100644 --- a/sandbox/v2/box_workspace_test.go +++ b/sandbox/v2/box_workspace_test.go @@ -83,19 +83,27 @@ func TestWorkspace_InvalidID(t *testing.T) { for _, pc := range testNodes() { pc := pc t.Run(pc.Name, func(t *testing.T) { - sbm, _ := setupManagerWithWorkspace(t, &pc) + sbm, wsm := setupManagerWithWorkspace(t, &pc) ensureTestImage(t, sbm, pc.TaiID) - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() - _, err := sbm.Create(ctx, sandbox.CreateOptions{ + wsID := "nonexistent-workspace" + + box, err := sbm.Create(ctx, sandbox.CreateOptions{ Image: testImage(), Owner: "user", - WorkspaceID: "nonexistent-workspace", + WorkspaceID: wsID, }) - assert.Error(t, err) - assert.Contains(t, err.Error(), "resolve workspace") + + // With online nodes the manager auto-creates the workspace. + require.NoError(t, err) + require.NotNil(t, box) + defer box.Remove(context.Background()) + if wsm != nil { + defer wsm.Delete(context.Background(), wsID, true) + } }) } } diff --git a/tai/api/register.go b/tai/api/register.go index 9c3640cb..c5779de8 100644 --- a/tai/api/register.go +++ b/tai/api/register.go @@ -8,7 +8,9 @@ import ( "github.com/gin-gonic/gin" "github.com/yaoapp/yao/openapi/oauth" + tai "github.com/yaoapp/yao/tai" "github.com/yaoapp/yao/tai/registry" + "github.com/yaoapp/yao/tai/taiid" ) // authenticateBearer validates a Bearer token and returns the caller's identity. @@ -33,6 +35,44 @@ func authenticateBearerDefault(token string) (registry.AuthInfo, error) { info.TeamID = result.Info.TeamID info.TenantID = result.Info.TenantID } + + slog.Info("[auth] buildAuthInfo result", + "subject", info.Subject, "user_id", info.UserID, + "client_id", info.ClientID, "team_id", info.TeamID, + "scope", info.Scope) + + if result.Claims != nil { + slog.Info("[auth] claims", + "claims.TeamID", result.Claims.TeamID, + "claims.TenantID", result.Claims.TenantID, + "claims.ClientID", result.Claims.ClientID, + "claims.Subject", result.Claims.Subject) + if result.Claims.Extra != nil { + slog.Info("[auth] claims.Extra", "extra", fmt.Sprintf("%+v", result.Claims.Extra)) + } else { + slog.Info("[auth] claims.Extra is nil") + } + + if info.TeamID == "" { + switch v := result.Claims.Extra["team_id"].(type) { + case string: + info.TeamID = v + slog.Info("[auth] team_id from Extra (string)", "team_id", v) + case float64: + info.TeamID = fmt.Sprintf("%.0f", v) + slog.Info("[auth] team_id from Extra (float64)", "team_id", info.TeamID) + default: + slog.Info("[auth] team_id not found in Extra or unknown type", + "type", fmt.Sprintf("%T", result.Claims.Extra["team_id"]), + "value", fmt.Sprintf("%v", result.Claims.Extra["team_id"])) + } + } + if info.TenantID == "" { + if v, ok := result.Claims.Extra["tenant_id"].(string); ok { + info.TenantID = v + } + } + } return info, nil } @@ -46,8 +86,10 @@ func extractBearer(r *http.Request) string { // registerRequest is the JSON body for POST /tai-nodes/register. type registerRequest struct { - TaiID string `json:"tai_id"` + NodeID string `json:"node_id,omitempty"` + ClientID string `json:"client_id,omitempty"` MachineID string `json:"machine_id"` + DisplayName string `json:"display_name,omitempty"` Version string `json:"version"` Addr string `json:"addr"` Ports map[string]int `json:"ports"` @@ -87,31 +129,63 @@ func HandleRegister(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"}) return } - if req.TaiID == "" { - c.JSON(http.StatusBadRequest, gin.H{"error": "tai_id is required"}) + + if req.NodeID == "" || req.MachineID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "node_id and machine_id are required"}) return } + resolvedTaiID, err := taiid.Generate(req.MachineID, req.NodeID) + if err != nil { + slog.Warn("taiid generation failed", "err", err) + c.JSON(http.StatusBadRequest, gin.H{"error": "failed to generate tai_id"}) + return + } + + remoteIP := c.ClientIP() + addr := req.Addr + if addr == "" && remoteIP != "" { + grpcPort := req.Ports["grpc"] + if grpcPort > 0 { + addr = fmt.Sprintf("tai://%s:%d", remoteIP, grpcPort) + } else { + addr = remoteIP + } + } + node := ®istry.TaiNode{ - TaiID: req.TaiID, + TaiID: resolvedTaiID, MachineID: req.MachineID, Version: req.Version, + DisplayName: req.DisplayName, Auth: authInfo, System: req.System, Mode: "direct", - Addr: req.Addr, + Addr: addr, Ports: req.Ports, Capabilities: req.Capabilities, } reg.Register(node) + slog.Info("[register] node registered via API", + "tai_id", resolvedTaiID, "addr", addr, "remote_ip", remoteIP, + "user_id", authInfo.UserID, "team_id", authInfo.TeamID) - remoteIP := c.ClientIP() - slog.Info("tai node registered via API", - "tai_id", req.TaiID, "remote_ip", remoteIP, "user_id", authInfo.UserID) + allBefore := reg.List() + slog.Info("[register] registry snapshot after Register", + "total", len(allBefore)) + for _, s := range allBefore { + slog.Info("[register] node", "tai_id", s.TaiID, "mode", s.Mode, "addr", s.Addr) + } + + if strings.HasPrefix(addr, "tai://") { + slog.Info("[register] launching connectRegisteredNode goroutine", + "tai_id", resolvedTaiID, "addr", addr) + go connectRegisteredNode(resolvedTaiID, addr, reg) + } c.JSON(http.StatusOK, gin.H{ "status": "registered", - "tai_id": req.TaiID, + "tai_id": resolvedTaiID, "remote_ip": remoteIP, }) } @@ -203,3 +277,46 @@ func HandleUnregister(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"status": "unregistered"}) } + +// connectRegisteredNode dials the self-registered Tai node via gRPC, +// creates a tai.Client, and binds it to the node's TaiID in the registry. +// initRemote internally registers a redundant "host-port" entry; we remove +// it so that the registry contains only the canonical taiID. +func connectRegisteredNode(taiID, addr string, reg *registry.Registry) { + slog.Info("[connect] start", "tai_id", taiID, "addr", addr) + + client, err := tai.New(addr) + if err != nil { + slog.Warn("[connect] tai.New FAILED", + "tai_id", taiID, "addr", addr, "err", err) + + allAfterFail := reg.List() + slog.Info("[connect] registry after tai.New failure", "total", len(allAfterFail)) + for _, s := range allAfterFail { + slog.Info("[connect] node", "tai_id", s.TaiID, "mode", s.Mode, "addr", s.Addr) + } + return + } + + autoID := client.TaiID() + slog.Info("[connect] tai.New OK", "tai_id", taiID, "autoID", autoID) + + allAfterNew := reg.List() + slog.Info("[connect] registry after tai.New", "total", len(allAfterNew)) + for _, s := range allAfterNew { + slog.Info("[connect] node", "tai_id", s.TaiID, "mode", s.Mode, "addr", s.Addr) + } + + if autoID != "" && autoID != taiID { + slog.Info("[connect] removing redundant autoID", "autoID", autoID) + reg.Unregister(autoID) + } + reg.SetClient(taiID, client) + + allFinal := reg.List() + slog.Info("[connect] registry FINAL", "total", len(allFinal)) + for _, s := range allFinal { + slog.Info("[connect] node", "tai_id", s.TaiID, "mode", s.Mode, "addr", s.Addr) + } + slog.Info("[connect] done", "tai_id", taiID) +} diff --git a/tai/api/register_test.go b/tai/api/register_test.go index d027469a..3957735d 100644 --- a/tai/api/register_test.go +++ b/tai/api/register_test.go @@ -46,9 +46,10 @@ func TestHandleRegister_Success(t *testing.T) { defer teardown() body := registerRequest{ - TaiID: "tai-abc123", + NodeID: "9100", MachineID: "m-001", Version: "0.2.0", + DisplayName: "My Dev Machine", Addr: "192.168.1.100", Ports: map[string]int{"grpc": 19100, "http": 8099}, Capabilities: map[string]bool{"docker": true, "host_exec": false}, @@ -74,14 +75,15 @@ func TestHandleRegister_Success(t *testing.T) { if resp["status"] != "registered" { t.Errorf("status = %v, want registered", resp["status"]) } - if resp["tai_id"] != "tai-abc123" { - t.Errorf("tai_id = %v, want tai-abc123", resp["tai_id"]) + taiID, _ := resp["tai_id"].(string) + if taiID == "" || len(taiID) < 5 || taiID[:4] != "tai-" { + t.Errorf("tai_id = %v, want server-generated tai-xxx", resp["tai_id"]) } if _, ok := resp["remote_ip"]; !ok { t.Error("response missing remote_ip") } - snap, ok := registry.Global().Get("tai-abc123") + snap, ok := registry.Global().Get(taiID) if !ok { t.Fatal("node not found in registry after register") } @@ -94,6 +96,74 @@ func TestHandleRegister_Success(t *testing.T) { if snap.Auth.UserID != "user-alice" { t.Errorf("Auth.UserID = %q, want user-alice", snap.Auth.UserID) } + if snap.DisplayName != "My Dev Machine" { + t.Errorf("DisplayName = %q, want %q", snap.DisplayName, "My Dev Machine") + } +} + +func TestHandleRegister_ServerGeneratedTaiID(t *testing.T) { + teardown := setupTest() + defer teardown() + + body := registerRequest{ + NodeID: "19100", + ClientID: "local-uuid-001", + MachineID: "m-001", + Version: "0.2.0", + DisplayName: "Generated ID Node", + Addr: "192.168.1.200", + Ports: map[string]int{"grpc": 19100}, + Capabilities: map[string]bool{"docker": true}, + System: registry.SystemInfo{OS: "darwin", Arch: "arm64", Hostname: "mac-01", NumCPU: 12}, + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/tai-nodes/register", jsonBody(body)) + c.Request.Header.Set("Authorization", "Bearer test-token") + c.Request.Header.Set("Content-Type", "application/json") + + HandleRegister(c) + + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want %d; body = %s", w.Code, http.StatusOK, w.Body.String()) + } + + var resp map[string]interface{} + json.Unmarshal(w.Body.Bytes(), &resp) + + generatedID, ok := resp["tai_id"].(string) + if !ok || generatedID == "" { + t.Fatal("response missing tai_id") + } + if generatedID == "19100" { + t.Error("tai_id should be server-generated, not the raw node_id") + } + if len(generatedID) != 26 { + t.Errorf("tai_id length = %d, want 26 (tai- + 22 base62); got %q", len(generatedID), generatedID) + } + + snap, ok2 := registry.Global().Get(generatedID) + if !ok2 { + t.Fatalf("node %q not found in registry", generatedID) + } + if snap.DisplayName != "Generated ID Node" { + t.Errorf("DisplayName = %q, want %q", snap.DisplayName, "Generated ID Node") + } + + // Deterministic: same inputs produce same ID + w2 := httptest.NewRecorder() + c2, _ := gin.CreateTestContext(w2) + c2.Request = httptest.NewRequest("POST", "/tai-nodes/register", jsonBody(body)) + c2.Request.Header.Set("Authorization", "Bearer test-token") + c2.Request.Header.Set("Content-Type", "application/json") + HandleRegister(c2) + + var resp2 map[string]interface{} + json.Unmarshal(w2.Body.Bytes(), &resp2) + if resp2["tai_id"] != generatedID { + t.Errorf("not deterministic: %v != %v", resp2["tai_id"], generatedID) + } } func TestHandleRegister_MissingAuth(t *testing.T) { @@ -102,7 +172,7 @@ func TestHandleRegister_MissingAuth(t *testing.T) { w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) - c.Request = httptest.NewRequest("POST", "/tai-nodes/register", jsonBody(registerRequest{TaiID: "x"})) + c.Request = httptest.NewRequest("POST", "/tai-nodes/register", jsonBody(registerRequest{NodeID: "x", MachineID: "m1"})) c.Request.Header.Set("Content-Type", "application/json") HandleRegister(c) @@ -112,7 +182,7 @@ func TestHandleRegister_MissingAuth(t *testing.T) { } } -func TestHandleRegister_MissingTaiID(t *testing.T) { +func TestHandleRegister_MissingTaiIDAndClientID(t *testing.T) { teardown := setupTest() defer teardown() diff --git a/tai/registry/registry.go b/tai/registry/registry.go index 52115787..a4683299 100644 --- a/tai/registry/registry.go +++ b/tai/registry/registry.go @@ -273,6 +273,21 @@ func (r *Registry) SetClient(taiID string, c any) { } } +// FindTaiIDByAuthClient returns the TaiID of the first node whose +// Auth.ClientID matches the given OAuth client ID. Returns "" if not found. +// This is needed because Tai's data channel authenticates with its OAuth +// ClientID, which may differ from the server-assigned TaiID. +func (r *Registry) FindTaiIDByAuthClient(clientID string) string { + r.mu.RLock() + defer r.mu.RUnlock() + for _, n := range r.nodes { + if n.Auth.ClientID == clientID { + return n.TaiID + } + } + return "" +} + // ListByTeam returns snapshots of all nodes belonging to the given team. func (r *Registry) ListByTeam(teamID string) []NodeSnapshot { r.mu.RLock() @@ -286,6 +301,20 @@ func (r *Registry) ListByTeam(teamID string) []NodeSnapshot { return result } +// ListByUser returns snapshots of all nodes registered by the given user +// that are NOT associated with any team. +func (r *Registry) ListByUser(userID string) []NodeSnapshot { + r.mu.RLock() + defer r.mu.RUnlock() + var result []NodeSnapshot + for _, n := range r.nodes { + if n.Auth.TeamID == "" && n.Auth.UserID == userID { + result = append(result, n.snapshot()) + } + } + return result +} + // StartHealthCheck runs a background goroutine that periodically checks // direct-mode nodes for heartbeat timeout. Nodes whose LastPing exceeds // timeout are marked offline. Nodes that remain offline longer than diff --git a/tai/tai.go b/tai/tai.go index a2663a13..e1170b85 100644 --- a/tai/tai.go +++ b/tai/tai.go @@ -334,6 +334,7 @@ func (c *Client) initTunnel(cfg *config) (*Client, error) { HTTP: nodePort(node.Ports, "http", 8099), VNC: nodePort(node.Ports, "vnc", 16080), Docker: nodePort(node.Ports, "docker", 12375), + K8s: nodePort(node.Ports, "k8s", 16443), } grpcLn, err := reg.OpenLocalListener(taiID, c.ports.GRPC) @@ -359,15 +360,30 @@ func (c *Client) initTunnel(cfg *config) (*Client, error) { } hasDocker := info.Capabilities["docker"] + hasK8s := info.Capabilities["k8s"] hasHostExec := info.Capabilities["host_exec"] - if !hasDocker && !hasHostExec { + if !hasDocker && !hasK8s && !hasHostExec { c.closeTunnelListeners() conn.Close() - return nil, fmt.Errorf("tai %s: no capabilities available via tunnel", taiID) + return nil, fmt.Errorf("tai %s: no capabilities available via tunnel (docker/k8s/host_exec all false)", taiID) } - if hasDocker && c.ports.Docker > 0 { + if cfg.runtime == K8s || (!hasDocker && hasK8s) { + k8sLn, err := reg.OpenLocalListener(taiID, c.ports.K8s) + if err == nil { + c.tunnelListeners = append(c.tunnelListeners, k8sLn) + sbAddr := k8sLn.Addr().String() + sb, err := sandbox.NewK8s(sbAddr, sandbox.K8sOption{ + Namespace: cfg.namespace, + KubeConfig: cfg.kubeConfig, + }) + if err == nil { + c.sb = sb + c.img = sandbox.NewK8sImage() + } + } + } else if hasDocker && c.ports.Docker > 0 { dockerLn, err := reg.OpenLocalListener(taiID, c.ports.Docker) if err == nil { c.tunnelListeners = append(c.tunnelListeners, dockerLn) diff --git a/tai/taiid/taiid.go b/tai/taiid/taiid.go new file mode 100644 index 00000000..c23f6180 --- /dev/null +++ b/tai/taiid/taiid.go @@ -0,0 +1,37 @@ +package taiid + +import ( + "crypto/sha256" + "fmt" + "math/big" +) + +const base62Chars = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" + +// Generate produces a deterministic tai_id from a machine ID and a node ID. +// The result is "tai-" followed by a Base62-encoded truncated SHA-256 hash. +// Both machineID and nodeID must be non-empty. +func Generate(machineID, nodeID string) (string, error) { + if machineID == "" || nodeID == "" { + return "", fmt.Errorf("machineID and nodeID are required") + } + h := sha256.Sum256([]byte(machineID + ":" + nodeID)) + return "tai-" + base62Encode(h[:16]), nil +} + +func base62Encode(data []byte) string { + num := new(big.Int).SetBytes(data) + base := big.NewInt(62) + zero := big.NewInt(0) + mod := new(big.Int) + + var encoded []byte + for num.Cmp(zero) > 0 { + num.DivMod(num, base, mod) + encoded = append([]byte{base62Chars[mod.Int64()]}, encoded...) + } + if len(encoded) == 0 { + return "0" + } + return string(encoded) +} diff --git a/tai/taiid/taiid_test.go b/tai/taiid/taiid_test.go new file mode 100644 index 00000000..66d6d2af --- /dev/null +++ b/tai/taiid/taiid_test.go @@ -0,0 +1,47 @@ +package taiid + +import ( + "testing" +) + +func TestGenerate_Deterministic(t *testing.T) { + id1, err := Generate("machine-abc", "9100") + if err != nil { + t.Fatalf("Generate: %v", err) + } + id2, err := Generate("machine-abc", "9100") + if err != nil { + t.Fatalf("Generate: %v", err) + } + if id1 != id2 { + t.Errorf("same inputs produced different results: %q vs %q", id1, id2) + } + if len(id1) < 5 || id1[:4] != "tai-" { + t.Errorf("result should start with 'tai-', got %q", id1) + } +} + +func TestGenerate_DifferentInputs(t *testing.T) { + id1, _ := Generate("machine-abc", "9100") + id2, _ := Generate("machine-abc", "9200") + id3, _ := Generate("machine-xyz", "9100") + + if id1 == id2 { + t.Errorf("different nodeID should produce different results: %q == %q", id1, id2) + } + if id1 == id3 { + t.Errorf("different machineID should produce different results: %q == %q", id1, id3) + } +} + +func TestGenerate_EmptyInputs(t *testing.T) { + if _, err := Generate("", "9100"); err == nil { + t.Error("empty machineID should return error") + } + if _, err := Generate("machine-abc", ""); err == nil { + t.Error("empty nodeID should return error") + } + if _, err := Generate("", ""); err == nil { + t.Error("both empty should return error") + } +} diff --git a/tai/tunnel/server.go b/tai/tunnel/server.go index 0429e3b8..3f0bb002 100644 --- a/tai/tunnel/server.go +++ b/tai/tunnel/server.go @@ -13,7 +13,9 @@ import ( "github.com/gin-gonic/gin" "github.com/gorilla/websocket" oauth "github.com/yaoapp/yao/openapi/oauth" + tai "github.com/yaoapp/yao/tai" "github.com/yaoapp/yao/tai/registry" + "github.com/yaoapp/yao/tai/taiid" ) var upgrader = websocket.Upgrader{ @@ -62,19 +64,32 @@ func HandleControl(c *gin.Context) { conn.Close() return } - if regMsg.TaiID == "" { - logger.Error("register message missing tai_id") + if regMsg.NodeID == "" || regMsg.MachineID == "" { + logger.Error("register message missing node_id or machine_id") + conn.Close() + return + } + resolvedTaiID, err := taiid.Generate(regMsg.MachineID, regMsg.NodeID) + if err != nil { + logger.Error("taiid generation failed", "err", err) conn.Close() return } + addr := "" + if host, _, err := net.SplitHostPort(c.Request.RemoteAddr); err == nil { + addr = "tunnel://" + host + } + node := ®istry.TaiNode{ - TaiID: regMsg.TaiID, + TaiID: resolvedTaiID, MachineID: regMsg.MachineID, Version: regMsg.Version, + DisplayName: regMsg.DisplayName, Auth: authInfo, System: regMsg.System, Mode: "tunnel", + Addr: addr, YaoBase: regMsg.Server, Ports: regMsg.Ports, Capabilities: regMsg.Capabilities, @@ -82,16 +97,18 @@ func HandleControl(c *gin.Context) { } reg.Register(node) defer func() { - reg.Unregister(regMsg.TaiID) - logger.Info("tai tunnel disconnected", "tai_id", regMsg.TaiID) + reg.Unregister(resolvedTaiID) + logger.Info("tai tunnel disconnected", "tai_id", resolvedTaiID) }() - if err := reg.WriteControlJSON(regMsg.TaiID, map[string]string{"type": "registered", "tai_id": regMsg.TaiID}); err != nil { + if err := reg.WriteControlJSON(resolvedTaiID, map[string]string{"type": "registered", "tai_id": resolvedTaiID}); err != nil { logger.Error("write registered response", "err", err) return } - logger.Info("tai tunnel connected", "tai_id", regMsg.TaiID, "version", regMsg.Version) + logger.Info("tai tunnel connected", "tai_id", resolvedTaiID, "version", regMsg.Version) + + go connectTunnelNode(resolvedTaiID, reg, logger) for { var msg controlMsg @@ -104,8 +121,8 @@ func HandleControl(c *gin.Context) { switch msg.Type { case "ping": - reg.UpdatePing(regMsg.TaiID) - if err := reg.WriteControlJSON(regMsg.TaiID, map[string]string{"type": "pong"}); err != nil { + reg.UpdatePing(resolvedTaiID) + if err := reg.WriteControlJSON(resolvedTaiID, map[string]string{"type": "pong"}); err != nil { logger.Debug("pong write failed", "err", err) return } @@ -150,9 +167,15 @@ func HandleData(c *gin.Context) { return } + resolvedTaiID := reg.FindTaiIDByAuthClient(authInfo.ClientID) + if resolvedTaiID == "" { + resolvedTaiID = authInfo.ClientID + } + wsConn := newWSConn(conn) - if err := reg.AcceptDataChannel(channelID, authInfo.ClientID, wsConn); err != nil { - logger.Debug("accept data channel failed", "channel_id", channelID, "err", err) + if err := reg.AcceptDataChannel(channelID, resolvedTaiID, wsConn); err != nil { + logger.Debug("accept data channel failed", "channel_id", channelID, "err", err, + "auth_client_id", authInfo.ClientID, "resolved_tai_id", resolvedTaiID) conn.Close() return } @@ -161,8 +184,10 @@ func HandleData(c *gin.Context) { // registerMessage is the JSON structure for Tai's register message. type registerMessage struct { Type string `json:"type"` - TaiID string `json:"tai_id"` + NodeID string `json:"node_id,omitempty"` + ClientID string `json:"client_id,omitempty"` MachineID string `json:"machine_id"` + DisplayName string `json:"display_name,omitempty"` Version string `json:"version"` Server string `json:"server"` Ports map[string]int `json:"ports"` @@ -207,6 +232,37 @@ func authenticateBearerDefault(token string) (registry.AuthInfo, error) { info.TeamID = result.Info.TeamID info.TenantID = result.Info.TenantID } + + slog.Info("[tunnel-auth] info from token", + "subject", info.Subject, "user_id", info.UserID, + "client_id", info.ClientID, "team_id", info.TeamID, + "scope", info.Scope) + + if result.Claims != nil { + slog.Info("[tunnel-auth] claims", + "claims.TeamID", result.Claims.TeamID, + "claims.ClientID", result.Claims.ClientID, + "extra", fmt.Sprintf("%+v", result.Claims.Extra)) + + if info.TeamID == "" && result.Claims.TeamID != "" { + info.TeamID = result.Claims.TeamID + } + if info.TeamID == "" { + switch v := result.Claims.Extra["team_id"].(type) { + case string: + info.TeamID = v + case float64: + info.TeamID = fmt.Sprintf("%.0f", v) + } + } + if info.TenantID == "" { + if v, ok := result.Claims.Extra["tenant_id"].(string); ok { + info.TenantID = v + } + } + } + + slog.Info("[tunnel-auth] final", "team_id", info.TeamID, "client_id", info.ClientID) return info, nil } @@ -267,3 +323,15 @@ func (c *wsConn) SetDeadline(t time.Time) error { 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) } + +// connectTunnelNode creates a tai.Client through the tunnel and binds it to the taiID. +func connectTunnelNode(taiID string, reg *registry.Registry, logger *slog.Logger) { + client, err := tai.New("tunnel://" + taiID) + if err != nil { + logger.Warn("failed to connect tunnel node", + "tai_id", taiID, "err", err) + return + } + _ = client // initTunnel already calls reg.SetClient(taiID, c) + logger.Info("tai client created for tunnel node", "tai_id", taiID) +} diff --git a/tai/tunnel/server_test.go b/tai/tunnel/server_test.go index 186e13b5..dfab4fde 100644 --- a/tai/tunnel/server_test.go +++ b/tai/tunnel/server_test.go @@ -281,7 +281,7 @@ func TestHandleControl_RegisterAndPing(t *testing.T) { regMsg := registerMessage{ Type: "register", - TaiID: "tai-001", + NodeID: "9100", MachineID: "m-test", Version: "2.0", Ports: map[string]int{"grpc": 9100}, @@ -297,11 +297,12 @@ func TestHandleControl_RegisterAndPing(t *testing.T) { 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"]) + gotTaiID := registered["tai_id"] + if gotTaiID == "" || len(gotTaiID) < 5 || gotTaiID[:4] != "tai-" { + t.Errorf("response tai_id = %q, want server-generated tai-xxx", gotTaiID) } - snap, ok := reg.Get("tai-001") + snap, ok := reg.Get(gotTaiID) if !ok { t.Fatal("node not found in registry after register") } @@ -332,15 +333,24 @@ func TestHandleControl_RegisterAndPing(t *testing.T) { t.Fatalf("write ping: %v", err) } - var pong map[string]string - if err := conn.ReadJSON(&pong); err != nil { - t.Fatalf("read pong: %v", err) + // Read messages until we get the pong; connectTunnelNode may inject + // "open" messages (with numeric fields) before our pong arrives. + var gotPong bool + for i := 0; i < 10; i++ { + var msg map[string]interface{} + if err := conn.ReadJSON(&msg); err != nil { + t.Fatalf("read message: %v", err) + } + if msg["type"] == "pong" { + gotPong = true + break + } } - if pong["type"] != "pong" { - t.Errorf("pong type = %q, want pong", pong["type"]) + if !gotPong { + t.Error("did not receive pong after ping") } - snap2, _ := reg.Get("tai-001") + snap2, _ := reg.Get(gotTaiID) if !snap2.LastPing.After(snap.LastPing) { t.Error("LastPing should be updated after ping") } @@ -529,9 +539,10 @@ func TestHandleControl_OpenChannelAndBridge(t *testing.T) { defer ctrlConn.Close() ctrlConn.WriteJSON(registerMessage{ - Type: "register", - TaiID: "tai-001", - Ports: map[string]int{"grpc": 9100}, + Type: "register", + NodeID: "9100", + MachineID: "m-test", + Ports: map[string]int{"grpc": 9100}, }) var registered map[string]string if err := ctrlConn.ReadJSON(®istered); err != nil { @@ -540,6 +551,7 @@ func TestHandleControl_OpenChannelAndBridge(t *testing.T) { if registered["type"] != "registered" { t.Fatalf("expected registered, got %v", registered) } + taiID := registered["tai_id"] var wg sync.WaitGroup wg.Add(1) @@ -547,7 +559,7 @@ func TestHandleControl_OpenChannelAndBridge(t *testing.T) { var channelConn net.Conn go func() { defer wg.Done() - _, resultCh, err := reg.RequestChannel("tai-001", 9100) + _, resultCh, err := reg.RequestChannel(taiID, 9100) if err != nil { requestErr = err return