feat(tai): enhance node registration and management

- Updated the registration process to utilize NodeID instead of TaiID, allowing for server-generated TaiIDs.
- Implemented additional fields in the registration request, including DisplayName and improved error handling for missing parameters.
- Enhanced the connection logic for registered nodes, ensuring proper client binding and logging for better traceability.
- Introduced new methods for extracting and managing extra claims in OAuth tokens, improving the flexibility of user identification.

Made-with: Cursor
This commit is contained in:
Max 2026-03-10 15:40:41 +08:00
parent 3f390c223c
commit 2dc4307175
23 changed files with 1834 additions and 82 deletions

View file

@ -10,9 +10,9 @@ NOW := $(shell date +"%FT%T%z")
OS := $(shell uname) OS := $(shell uname)
# ROOT_DIR := $(shell dirname $(realpath $(firstword $(MAKEFILE_LIST)))) # 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) # 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) # 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') TESTFOLDER_AGENT := $(shell $(GO) list ./agent/... ./aigc/... | grep -vE 'agent/search/handlers/web|agent/robot/|agent/sandbox/v2')
# KB tests (kb) # KB tests (kb)

View file

@ -2,6 +2,7 @@ package types
import ( import (
"crypto/sha256" "crypto/sha256"
"encoding/json"
"fmt" "fmt"
"os" "os"
"path/filepath" "path/filepath"
@ -72,11 +73,11 @@ func ToSandboxV2(v any) (*sandboxTypes.SandboxConfig, error) {
func ComputeConfigHash(cfg *sandboxTypes.SandboxConfig, mcpServers []MCPServerConfig, skillsDir string) string { func ComputeConfigHash(cfg *sandboxTypes.SandboxConfig, mcpServers []MCPServerConfig, skillsDir string) string {
h := sha256.New() h := sha256.New()
raw, _ := jsoniter.Marshal(cfg) raw, _ := json.Marshal(cfg)
h.Write(raw) h.Write(raw)
if len(mcpServers) > 0 { if len(mcpServers) > 0 {
mcpRaw, _ := jsoniter.Marshal(mcpServers) mcpRaw, _ := json.Marshal(mcpServers)
h.Write(mcpRaw) h.Write(mcpRaw)
} }

138
openapi/nodes/nodes.go Normal file
View file

@ -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))
}

View file

@ -2,6 +2,7 @@ package openapi
import ( import (
"encoding/base64" "encoding/base64"
"fmt"
"net/http" "net/http"
"strings" "strings"
@ -585,11 +586,32 @@ func (openapi *OpenAPI) oauthDeviceAuthorize(c *gin.Context) {
if extraClaims == nil { if extraClaims == nil {
extraClaims = make(map[string]interface{}) 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") userCode := c.PostForm("user_code")

View file

@ -189,13 +189,29 @@ func (s *Service) refreshTokenDirect(refreshToken string, expiredClaims *types.T
// buildAuthInfo constructs AuthorizedInfo directly from token claims, // buildAuthInfo constructs AuthorizedInfo directly from token claims,
// equivalent to the SetInfo+GetInfo round-trip through gin.Context. // equivalent to the SetInfo+GetInfo round-trip through gin.Context.
func (s *Service) buildAuthInfo(claims *types.TokenClaims, sessionID string) *types.AuthorizedInfo { 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{ info := &types.AuthorizedInfo{
Subject: claims.Subject, Subject: claims.Subject,
ClientID: claims.ClientID, ClientID: claims.ClientID,
Scope: claims.Scope, Scope: claims.Scope,
SessionID: sessionID, SessionID: sessionID,
TeamID: claims.TeamID, TeamID: teamID,
TenantID: claims.TenantID, TenantID: tenantID,
} }
userID, err := s.UserID(claims.ClientID, claims.Subject) 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 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 return info
} }

View file

@ -241,9 +241,11 @@ func (s *Service) RefreshToken(ctx context.Context, refreshToken string, scope .
finalScope = requestedScope finalScope = requestedScope
} }
extraClaims := extractExtraClaims(tokenInfo)
// Generate new access token with final scope // Generate new access token with final scope
expiresIn := int(s.config.Token.AccessTokenLifetime.Seconds()) 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 { if err != nil {
return nil, &types.ErrorResponse{ return nil, &types.ErrorResponse{
Code: types.ErrorServerError, Code: types.ErrorServerError,
@ -340,9 +342,11 @@ func (s *Service) RotateRefreshToken(ctx context.Context, oldToken string, reque
finalScope = scope finalScope = scope
} }
extraClaims := extractExtraClaims(tokenInfo)
// Generate new tokens with final scope and original subject // Generate new tokens with final scope and original subject
expiresIn := int(s.config.Token.AccessTokenLifetime.Seconds()) 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 { if err != nil {
return nil, &types.ErrorResponse{ return nil, &types.ErrorResponse{
Code: types.ErrorServerError, 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 { if err != nil {
return nil, &types.ErrorResponse{ return nil, &types.ErrorResponse{
Code: types.ErrorServerError, Code: types.ErrorServerError,
@ -497,20 +501,12 @@ func (s *Service) handleRefreshTokenGrant(ctx context.Context, client *types.Cli
return nil, err return nil, err
} }
// Extract scope and subject from refresh token if available scope, _ := refreshTokenInfo["scope"].(string)
scope := "" subject, _ := refreshTokenInfo["subject"].(string)
if scopeVal, ok := refreshTokenInfo["scope"].(string); ok { extraClaims := extractExtraClaims(refreshTokenInfo)
scope = scopeVal
}
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()) 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 { if err != nil {
return nil, &types.ErrorResponse{ return nil, &types.ErrorResponse{
Code: types.ErrorServerError, Code: types.ErrorServerError,
@ -524,9 +520,8 @@ func (s *Service) handleRefreshTokenGrant(ctx context.Context, client *types.Cli
ExpiresIn: expiresIn, ExpiresIn: expiresIn,
} }
// Include refresh token if rotation is enabled
if s.config.Features.RefreshTokenRotationEnabled { 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 { if err != nil {
return nil, &types.ErrorResponse{ return nil, &types.ErrorResponse{
Code: types.ErrorServerError, Code: types.ErrorServerError,
@ -534,11 +529,8 @@ func (s *Service) handleRefreshTokenGrant(ctx context.Context, client *types.Cli
} }
} }
token.RefreshToken = newRefreshToken token.RefreshToken = newRefreshToken
// Revoke old refresh token
s.revokeRefreshToken(refreshToken) s.revokeRefreshToken(refreshToken)
} else { } else {
// Reuse the same refresh token
token.RefreshToken = refreshToken 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
}

View file

@ -320,6 +320,22 @@ func (s *Service) UserID(clientID, subject string) (string, error) {
return userIDStr, nil 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 // MakeAuthorizationCode generates a new authorization code with specific parameters and stores it
// ============================================================================ // ============================================================================

View file

@ -19,6 +19,7 @@ import (
"github.com/yaoapp/yao/openapi/llm" "github.com/yaoapp/yao/openapi/llm"
"github.com/yaoapp/yao/openapi/mcp" "github.com/yaoapp/yao/openapi/mcp"
"github.com/yaoapp/yao/openapi/messenger" "github.com/yaoapp/yao/openapi/messenger"
"github.com/yaoapp/yao/openapi/nodes"
"github.com/yaoapp/yao/openapi/oauth" "github.com/yaoapp/yao/openapi/oauth"
"github.com/yaoapp/yao/openapi/oauth/acl" "github.com/yaoapp/yao/openapi/oauth/acl"
"github.com/yaoapp/yao/openapi/oauth/types" "github.com/yaoapp/yao/openapi/oauth/types"
@ -28,6 +29,7 @@ import (
"github.com/yaoapp/yao/openapi/team" "github.com/yaoapp/yao/openapi/team"
openapiTrace "github.com/yaoapp/yao/openapi/trace" openapiTrace "github.com/yaoapp/yao/openapi/trace"
"github.com/yaoapp/yao/openapi/user" "github.com/yaoapp/yao/openapi/user"
openapiWorkspace "github.com/yaoapp/yao/openapi/workspace"
taiapi "github.com/yaoapp/yao/tai/api" taiapi "github.com/yaoapp/yao/tai/api"
taitunnel "github.com/yaoapp/yao/tai/tunnel" taitunnel "github.com/yaoapp/yao/tai/tunnel"
) )
@ -173,9 +175,17 @@ func (openapi *OpenAPI) Attach(router *gin.Engine) {
// OTP handlers (passwordless authentication) // OTP handlers (passwordless authentication)
otp.Attach(group.Group("/otp"), openapi.OAuth) 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.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 // Tai tunnel WebSocket and reverse proxy routes
group.GET("/ws/tai", taitunnel.HandleControl) group.GET("/ws/tai", taitunnel.HandleControl)

366
openapi/sandbox/manage.go Normal file
View file

@ -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)
}

View file

@ -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)
}

View file

@ -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)
}

View file

@ -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)
}

View file

@ -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)
}

View file

@ -198,6 +198,25 @@ func (b *Box) Workspace() workspace.FS {
// WorkspaceID returns the workspace ID mounted to this sandbox, or empty string. // WorkspaceID returns the workspace ID mounted to this sandbox, or empty string.
func (b *Box) WorkspaceID() string { return b.workspaceID } 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. // VNC returns the VNC WebSocket URL.
func (b *Box) VNC(ctx context.Context) (string, error) { func (b *Box) VNC(ctx context.Context) (string, error) {
b.touch() b.touch()

View file

@ -83,19 +83,27 @@ func TestWorkspace_InvalidID(t *testing.T) {
for _, pc := range testNodes() { for _, pc := range testNodes() {
pc := pc pc := pc
t.Run(pc.Name, func(t *testing.T) { t.Run(pc.Name, func(t *testing.T) {
sbm, _ := setupManagerWithWorkspace(t, &pc) sbm, wsm := setupManagerWithWorkspace(t, &pc)
ensureTestImage(t, sbm, pc.TaiID) 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() defer cancel()
_, err := sbm.Create(ctx, sandbox.CreateOptions{ wsID := "nonexistent-workspace"
box, err := sbm.Create(ctx, sandbox.CreateOptions{
Image: testImage(), Image: testImage(),
Owner: "user", 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)
}
}) })
} }
} }

View file

@ -8,7 +8,9 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/yaoapp/yao/openapi/oauth" "github.com/yaoapp/yao/openapi/oauth"
tai "github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/registry"
"github.com/yaoapp/yao/tai/taiid"
) )
// authenticateBearer validates a Bearer token and returns the caller's identity. // 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.TeamID = result.Info.TeamID
info.TenantID = result.Info.TenantID 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 return info, nil
} }
@ -46,8 +86,10 @@ func extractBearer(r *http.Request) string {
// registerRequest is the JSON body for POST /tai-nodes/register. // registerRequest is the JSON body for POST /tai-nodes/register.
type registerRequest struct { 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"` MachineID string `json:"machine_id"`
DisplayName string `json:"display_name,omitempty"`
Version string `json:"version"` Version string `json:"version"`
Addr string `json:"addr"` Addr string `json:"addr"`
Ports map[string]int `json:"ports"` 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"}) c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
return 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 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 := &registry.TaiNode{ node := &registry.TaiNode{
TaiID: req.TaiID, TaiID: resolvedTaiID,
MachineID: req.MachineID, MachineID: req.MachineID,
Version: req.Version, Version: req.Version,
DisplayName: req.DisplayName,
Auth: authInfo, Auth: authInfo,
System: req.System, System: req.System,
Mode: "direct", Mode: "direct",
Addr: req.Addr, Addr: addr,
Ports: req.Ports, Ports: req.Ports,
Capabilities: req.Capabilities, Capabilities: req.Capabilities,
} }
reg.Register(node) 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() allBefore := reg.List()
slog.Info("tai node registered via API", slog.Info("[register] registry snapshot after Register",
"tai_id", req.TaiID, "remote_ip", remoteIP, "user_id", authInfo.UserID) "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{ c.JSON(http.StatusOK, gin.H{
"status": "registered", "status": "registered",
"tai_id": req.TaiID, "tai_id": resolvedTaiID,
"remote_ip": remoteIP, "remote_ip": remoteIP,
}) })
} }
@ -203,3 +277,46 @@ func HandleUnregister(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"status": "unregistered"}) 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)
}

View file

@ -46,9 +46,10 @@ func TestHandleRegister_Success(t *testing.T) {
defer teardown() defer teardown()
body := registerRequest{ body := registerRequest{
TaiID: "tai-abc123", NodeID: "9100",
MachineID: "m-001", MachineID: "m-001",
Version: "0.2.0", Version: "0.2.0",
DisplayName: "My Dev Machine",
Addr: "192.168.1.100", Addr: "192.168.1.100",
Ports: map[string]int{"grpc": 19100, "http": 8099}, Ports: map[string]int{"grpc": 19100, "http": 8099},
Capabilities: map[string]bool{"docker": true, "host_exec": false}, Capabilities: map[string]bool{"docker": true, "host_exec": false},
@ -74,14 +75,15 @@ func TestHandleRegister_Success(t *testing.T) {
if resp["status"] != "registered" { if resp["status"] != "registered" {
t.Errorf("status = %v, want registered", resp["status"]) t.Errorf("status = %v, want registered", resp["status"])
} }
if resp["tai_id"] != "tai-abc123" { taiID, _ := resp["tai_id"].(string)
t.Errorf("tai_id = %v, want tai-abc123", resp["tai_id"]) 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 { if _, ok := resp["remote_ip"]; !ok {
t.Error("response missing remote_ip") t.Error("response missing remote_ip")
} }
snap, ok := registry.Global().Get("tai-abc123") snap, ok := registry.Global().Get(taiID)
if !ok { if !ok {
t.Fatal("node not found in registry after register") 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" { if snap.Auth.UserID != "user-alice" {
t.Errorf("Auth.UserID = %q, want user-alice", snap.Auth.UserID) 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) { func TestHandleRegister_MissingAuth(t *testing.T) {
@ -102,7 +172,7 @@ func TestHandleRegister_MissingAuth(t *testing.T) {
w := httptest.NewRecorder() w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w) 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") c.Request.Header.Set("Content-Type", "application/json")
HandleRegister(c) 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() teardown := setupTest()
defer teardown() defer teardown()

View file

@ -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. // ListByTeam returns snapshots of all nodes belonging to the given team.
func (r *Registry) ListByTeam(teamID string) []NodeSnapshot { func (r *Registry) ListByTeam(teamID string) []NodeSnapshot {
r.mu.RLock() r.mu.RLock()
@ -286,6 +301,20 @@ func (r *Registry) ListByTeam(teamID string) []NodeSnapshot {
return result 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 // StartHealthCheck runs a background goroutine that periodically checks
// direct-mode nodes for heartbeat timeout. Nodes whose LastPing exceeds // direct-mode nodes for heartbeat timeout. Nodes whose LastPing exceeds
// timeout are marked offline. Nodes that remain offline longer than // timeout are marked offline. Nodes that remain offline longer than

View file

@ -334,6 +334,7 @@ func (c *Client) initTunnel(cfg *config) (*Client, error) {
HTTP: nodePort(node.Ports, "http", 8099), HTTP: nodePort(node.Ports, "http", 8099),
VNC: nodePort(node.Ports, "vnc", 16080), VNC: nodePort(node.Ports, "vnc", 16080),
Docker: nodePort(node.Ports, "docker", 12375), Docker: nodePort(node.Ports, "docker", 12375),
K8s: nodePort(node.Ports, "k8s", 16443),
} }
grpcLn, err := reg.OpenLocalListener(taiID, c.ports.GRPC) grpcLn, err := reg.OpenLocalListener(taiID, c.ports.GRPC)
@ -359,15 +360,30 @@ func (c *Client) initTunnel(cfg *config) (*Client, error) {
} }
hasDocker := info.Capabilities["docker"] hasDocker := info.Capabilities["docker"]
hasK8s := info.Capabilities["k8s"]
hasHostExec := info.Capabilities["host_exec"] hasHostExec := info.Capabilities["host_exec"]
if !hasDocker && !hasHostExec { if !hasDocker && !hasK8s && !hasHostExec {
c.closeTunnelListeners() c.closeTunnelListeners()
conn.Close() 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) dockerLn, err := reg.OpenLocalListener(taiID, c.ports.Docker)
if err == nil { if err == nil {
c.tunnelListeners = append(c.tunnelListeners, dockerLn) c.tunnelListeners = append(c.tunnelListeners, dockerLn)

37
tai/taiid/taiid.go Normal file
View file

@ -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)
}

47
tai/taiid/taiid_test.go Normal file
View file

@ -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")
}
}

View file

@ -13,7 +13,9 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
oauth "github.com/yaoapp/yao/openapi/oauth" oauth "github.com/yaoapp/yao/openapi/oauth"
tai "github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/registry"
"github.com/yaoapp/yao/tai/taiid"
) )
var upgrader = websocket.Upgrader{ var upgrader = websocket.Upgrader{
@ -62,19 +64,32 @@ func HandleControl(c *gin.Context) {
conn.Close() conn.Close()
return return
} }
if regMsg.TaiID == "" { if regMsg.NodeID == "" || regMsg.MachineID == "" {
logger.Error("register message missing tai_id") 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() conn.Close()
return return
} }
addr := ""
if host, _, err := net.SplitHostPort(c.Request.RemoteAddr); err == nil {
addr = "tunnel://" + host
}
node := &registry.TaiNode{ node := &registry.TaiNode{
TaiID: regMsg.TaiID, TaiID: resolvedTaiID,
MachineID: regMsg.MachineID, MachineID: regMsg.MachineID,
Version: regMsg.Version, Version: regMsg.Version,
DisplayName: regMsg.DisplayName,
Auth: authInfo, Auth: authInfo,
System: regMsg.System, System: regMsg.System,
Mode: "tunnel", Mode: "tunnel",
Addr: addr,
YaoBase: regMsg.Server, YaoBase: regMsg.Server,
Ports: regMsg.Ports, Ports: regMsg.Ports,
Capabilities: regMsg.Capabilities, Capabilities: regMsg.Capabilities,
@ -82,16 +97,18 @@ func HandleControl(c *gin.Context) {
} }
reg.Register(node) reg.Register(node)
defer func() { defer func() {
reg.Unregister(regMsg.TaiID) reg.Unregister(resolvedTaiID)
logger.Info("tai tunnel disconnected", "tai_id", regMsg.TaiID) 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) logger.Error("write registered response", "err", err)
return 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 { for {
var msg controlMsg var msg controlMsg
@ -104,8 +121,8 @@ func HandleControl(c *gin.Context) {
switch msg.Type { switch msg.Type {
case "ping": case "ping":
reg.UpdatePing(regMsg.TaiID) reg.UpdatePing(resolvedTaiID)
if err := reg.WriteControlJSON(regMsg.TaiID, map[string]string{"type": "pong"}); err != nil { if err := reg.WriteControlJSON(resolvedTaiID, map[string]string{"type": "pong"}); err != nil {
logger.Debug("pong write failed", "err", err) logger.Debug("pong write failed", "err", err)
return return
} }
@ -150,9 +167,15 @@ func HandleData(c *gin.Context) {
return return
} }
resolvedTaiID := reg.FindTaiIDByAuthClient(authInfo.ClientID)
if resolvedTaiID == "" {
resolvedTaiID = authInfo.ClientID
}
wsConn := newWSConn(conn) wsConn := newWSConn(conn)
if err := reg.AcceptDataChannel(channelID, authInfo.ClientID, wsConn); err != nil { if err := reg.AcceptDataChannel(channelID, resolvedTaiID, wsConn); err != nil {
logger.Debug("accept data channel failed", "channel_id", channelID, "err", err) logger.Debug("accept data channel failed", "channel_id", channelID, "err", err,
"auth_client_id", authInfo.ClientID, "resolved_tai_id", resolvedTaiID)
conn.Close() conn.Close()
return return
} }
@ -161,8 +184,10 @@ func HandleData(c *gin.Context) {
// registerMessage is the JSON structure for Tai's register message. // registerMessage is the JSON structure for Tai's register message.
type registerMessage struct { type registerMessage struct {
Type string `json:"type"` 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"` MachineID string `json:"machine_id"`
DisplayName string `json:"display_name,omitempty"`
Version string `json:"version"` Version string `json:"version"`
Server string `json:"server"` Server string `json:"server"`
Ports map[string]int `json:"ports"` Ports map[string]int `json:"ports"`
@ -207,6 +232,37 @@ func authenticateBearerDefault(token string) (registry.AuthInfo, error) {
info.TeamID = result.Info.TeamID info.TeamID = result.Info.TeamID
info.TenantID = result.Info.TenantID 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 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) SetReadDeadline(t time.Time) error { return c.ws.SetReadDeadline(t) }
func (c *wsConn) SetWriteDeadline(t time.Time) error { return c.ws.SetWriteDeadline(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)
}

View file

@ -281,7 +281,7 @@ func TestHandleControl_RegisterAndPing(t *testing.T) {
regMsg := registerMessage{ regMsg := registerMessage{
Type: "register", Type: "register",
TaiID: "tai-001", NodeID: "9100",
MachineID: "m-test", MachineID: "m-test",
Version: "2.0", Version: "2.0",
Ports: map[string]int{"grpc": 9100}, Ports: map[string]int{"grpc": 9100},
@ -297,11 +297,12 @@ func TestHandleControl_RegisterAndPing(t *testing.T) {
if registered["type"] != "registered" { if registered["type"] != "registered" {
t.Errorf("response type = %q, want registered", registered["type"]) t.Errorf("response type = %q, want registered", registered["type"])
} }
if registered["tai_id"] != "tai-001" { gotTaiID := registered["tai_id"]
t.Errorf("response tai_id = %q, want tai-001", 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 { if !ok {
t.Fatal("node not found in registry after register") 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) t.Fatalf("write ping: %v", err)
} }
var pong map[string]string // Read messages until we get the pong; connectTunnelNode may inject
if err := conn.ReadJSON(&pong); err != nil { // "open" messages (with numeric fields) before our pong arrives.
t.Fatalf("read pong: %v", err) 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 pong["type"] != "pong" { if msg["type"] == "pong" {
t.Errorf("pong type = %q, want pong", pong["type"]) gotPong = true
break
}
}
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) { if !snap2.LastPing.After(snap.LastPing) {
t.Error("LastPing should be updated after ping") t.Error("LastPing should be updated after ping")
} }
@ -530,7 +540,8 @@ func TestHandleControl_OpenChannelAndBridge(t *testing.T) {
ctrlConn.WriteJSON(registerMessage{ ctrlConn.WriteJSON(registerMessage{
Type: "register", Type: "register",
TaiID: "tai-001", NodeID: "9100",
MachineID: "m-test",
Ports: map[string]int{"grpc": 9100}, Ports: map[string]int{"grpc": 9100},
}) })
var registered map[string]string var registered map[string]string
@ -540,6 +551,7 @@ func TestHandleControl_OpenChannelAndBridge(t *testing.T) {
if registered["type"] != "registered" { if registered["type"] != "registered" {
t.Fatalf("expected registered, got %v", registered) t.Fatalf("expected registered, got %v", registered)
} }
taiID := registered["tai_id"]
var wg sync.WaitGroup var wg sync.WaitGroup
wg.Add(1) wg.Add(1)
@ -547,7 +559,7 @@ func TestHandleControl_OpenChannelAndBridge(t *testing.T) {
var channelConn net.Conn var channelConn net.Conn
go func() { go func() {
defer wg.Done() defer wg.Done()
_, resultCh, err := reg.RequestChannel("tai-001", 9100) _, resultCh, err := reg.RequestChannel(taiID, 9100)
if err != nil { if err != nil {
requestErr = err requestErr = err
return return