yao/tai/dial.go
Max 42b45f9357 refactor(sandbox): streamline sandbox configuration loading and enhance local execution capabilities
- Replaced the previous file extension checks with a unified parsing function for sandbox configuration, improving code clarity and maintainability.
- Introduced a new HostExecConfig structure to manage local execution settings, allowing for more granular control over command execution permissions.
- Removed deprecated Moapi API files and related functionality, simplifying the codebase and reducing maintenance overhead.
- Updated the Tai node registration process to ensure local capabilities are accurately reflected based on the environment, enhancing overall system robustness.
2026-03-18 11:22:08 +08:00

400 lines
11 KiB
Go

package tai
import (
"context"
"fmt"
"net"
"net/http"
"time"
yaoconfig "github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/tai/hostexec"
hepb "github.com/yaoapp/yao/tai/hostexec/pb"
"github.com/yaoapp/yao/tai/proxy"
"github.com/yaoapp/yao/tai/registry"
"github.com/yaoapp/yao/tai/runtime"
sipb "github.com/yaoapp/yao/tai/serverinfo/pb"
"github.com/yaoapp/yao/tai/types"
"github.com/yaoapp/yao/tai/vnc"
"github.com/yaoapp/yao/tai/volume"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/keepalive"
)
// DialRemote establishes connections to a remote Tai node via gRPC (direct mode).
// Does NOT interact with the registry. Caller must call ConnResources.Close().
func DialRemote(host string, ports types.Ports, opts ...DialOption) (*ConnResources, error) {
cfg := &dialConfig{ports: mergedPorts(ports)}
for _, o := range opts {
o.applyDial(cfg)
}
grpcAddr := fmt.Sprintf("%s:%d", host, cfg.ports.GRPC)
conn, err := dialGRPC(grpcAddr)
if err != nil {
return nil, fmt.Errorf("grpc dial %s: %w", grpcAddr, err)
}
return buildResources(conn, cfg, &remoteEnv{host: host, httpClient: cfg.httpClient})
}
// DialTunnel establishes connections to a Tai node through the WebSocket tunnel.
// Requires the node to already be registered in the registry (online).
// Does NOT call registry.SetResources. Caller must call ConnResources.Close().
func DialTunnel(taiID string, reg *registry.Registry, opts ...DialOption) (*ConnResources, error) {
node, ok := reg.Get(taiID)
if !ok || node.Status != "online" {
return nil, fmt.Errorf("tai node %s not online", taiID)
}
cfg := &dialConfig{
ports: types.Ports{
GRPC: intOr(node.Ports.GRPC, 19100),
HTTP: intOr(node.Ports.HTTP, 8099),
VNC: intOr(node.Ports.VNC, 16080),
Docker: intOr(node.Ports.Docker, 12375),
K8s: intOr(node.Ports.K8s, 16443),
},
}
for _, o := range opts {
o.applyDial(cfg)
}
grpcLn, err := reg.OpenLocalListener(taiID, cfg.ports.GRPC)
if err != nil {
return nil, fmt.Errorf("open grpc tunnel listener: %w", err)
}
conn, err := dialGRPC("passthrough:///" + grpcLn.Addr().String())
if err != nil {
grpcLn.Close()
return nil, fmt.Errorf("grpc dial tunnel %s: %w", grpcLn.Addr(), err)
}
env := &tunnelEnv{
taiID: taiID,
yaoBase: node.YaoBase,
reg: reg,
regCaps: node.Capabilities,
listeners: []net.Listener{grpcLn},
}
res, err := buildResources(conn, cfg, env)
if err != nil {
grpcLn.Close()
conn.Close()
return nil, err
}
res.Listeners = env.listeners
return res, nil
}
// DialLocal establishes connections to the local host as a Tai node.
// Docker is probed but not required — when unavailable the node still
// provides Volume (and optionally HostExec) capabilities.
// Does NOT interact with the registry. Caller must call ConnResources.Close().
func DialLocal(addr string, dataDir string, vol volume.Volume) (*ConnResources, error) {
sb, _ := runtime.NewLocal(addr) // Docker failure is non-fatal
res := &ConnResources{
DataDir: dataDir,
System: CollectSystemInfo(),
}
if sb != nil {
res.Runtime = sb
res.Image = runtime.NewDockerImage(runtime.DockerCli(sb))
res.Proxy = proxy.NewLocal(sb)
res.VNC = vnc.NewLocal(sb)
}
if yaoconfig.Conf.HostExec.Enabled {
res.HostExec = hostexec.NewLocalClient(dataDir, hostexec.Policy{
FullAccess: yaoconfig.Conf.HostExec.FullAccess,
AllowedCommands: yaoconfig.Conf.HostExec.AllowedCommands,
AllowedDirs: yaoconfig.Conf.HostExec.AllowedDirs,
DeniedDirs: yaoconfig.Conf.HostExec.DeniedDirs,
})
}
if vol != nil {
res.Volume = vol
} else {
if dataDir == "" {
dataDir = "/tmp/tai-volumes"
}
res.DataDir = dataDir
res.Volume = volume.NewLocal(dataDir)
}
return res, nil
}
// ---------------------------------------------------------------------------
// Shared build logic
// ---------------------------------------------------------------------------
// dialEnv abstracts the mode-specific differences (remote vs tunnel) that
// buildResources needs.
type dialEnv interface {
fallbackCaps() map[string]bool
mergeCaps(discovered map[string]bool) types.Capabilities
// listenAddr opens or formats a host:port address for the given port.
// Tunnel mode opens a local listener; remote mode formats host:port.
listenAddr(port int) (string, error)
newProxy(ports types.Ports) proxy.Proxy
newVNC(ports types.Ports) vnc.VNC
}
// buildResources constructs a ConnResources from an established gRPC
// connection. Shared by DialRemote and DialTunnel.
func buildResources(conn *grpc.ClientConn, cfg *dialConfig, env dialEnv) (*ConnResources, error) {
info, err := discoverInfo(conn, cfg)
if err != nil {
info = &discoveredInfo{Capabilities: env.fallbackCaps()}
}
caps := env.mergeCaps(info.Capabilities)
res := &ConnResources{
GRPCConn: conn,
HostExec: hepb.NewHostExecClient(conn),
Volume: volume.NewRemote(conn),
Caps: caps,
System: info.System,
Ports: cfg.ports,
Version: info.Version,
}
if cfg.runtime == types.K8s || (!caps.Docker && caps.K8s) {
if cfg.kubeConfig != "" {
k8sPort := cfg.ports.K8s
if k8sPort == 0 {
k8sPort = 16443
}
addr, err := env.listenAddr(k8sPort)
if err == nil {
sb, err := runtime.NewK8s(addr, runtime.K8sOption{
Namespace: cfg.namespace,
KubeConfig: cfg.kubeConfig,
})
if err == nil {
res.Runtime = sb
res.Image = runtime.NewK8sImage()
}
}
}
} else if caps.Docker {
dockerPort := cfg.ports.Docker
if dockerPort == 0 {
dockerPort = 12375
}
addr, err := env.listenAddr(dockerPort)
if err == nil {
sb, err := runtime.NewDocker("tcp://" + addr)
if err == nil {
res.Runtime = sb
res.Image = runtime.NewDockerImage(runtime.DockerCli(sb))
}
}
}
if res.Runtime != nil {
res.Proxy = env.newProxy(cfg.ports)
}
res.VNC = env.newVNC(cfg.ports)
return res, nil
}
// ---------------------------------------------------------------------------
// remoteEnv — direct TCP connections
// ---------------------------------------------------------------------------
type remoteEnv struct {
host string
httpClient *http.Client
}
func (e *remoteEnv) fallbackCaps() map[string]bool {
return map[string]bool{"docker": true}
}
func (e *remoteEnv) mergeCaps(discovered map[string]bool) types.Capabilities {
return types.Capabilities{
Docker: discovered["docker"],
K8s: discovered["k8s"],
HostExec: discovered["host_exec"],
}
}
func (e *remoteEnv) listenAddr(port int) (string, error) {
return fmt.Sprintf("%s:%d", e.host, port), nil
}
func (e *remoteEnv) newProxy(ports types.Ports) proxy.Proxy {
return proxy.NewRemote(e.host, ports.HTTP, e.httpClient)
}
func (e *remoteEnv) newVNC(ports types.Ports) vnc.VNC {
return vnc.NewRemote(e.host, ports.VNC, e.httpClient)
}
// ---------------------------------------------------------------------------
// tunnelEnv — connections via WebSocket tunnel
// ---------------------------------------------------------------------------
type tunnelEnv struct {
taiID string
yaoBase string
reg *registry.Registry
regCaps types.Capabilities
listeners []net.Listener
}
func (e *tunnelEnv) fallbackCaps() map[string]bool {
return make(map[string]bool)
}
func (e *tunnelEnv) mergeCaps(discovered map[string]bool) types.Capabilities {
return types.Capabilities{
Docker: discovered["docker"] || e.regCaps.Docker,
K8s: discovered["k8s"] || e.regCaps.K8s,
HostExec: discovered["host_exec"] || e.regCaps.HostExec,
}
}
func (e *tunnelEnv) listenAddr(port int) (string, error) {
ln, err := e.reg.OpenLocalListener(e.taiID, port)
if err != nil {
return "", err
}
e.listeners = append(e.listeners, ln)
return ln.Addr().String(), nil
}
func (e *tunnelEnv) newProxy(_ types.Ports) proxy.Proxy {
return proxy.NewTunnel(e.taiID, e.yaoBase)
}
func (e *tunnelEnv) newVNC(_ types.Ports) vnc.VNC {
return vnc.NewTunnel(e.taiID, e.yaoBase)
}
// ---------------------------------------------------------------------------
// Dial options
// ---------------------------------------------------------------------------
// DialOption configures a Dial* call.
type DialOption interface {
applyDial(*dialConfig)
}
type dialOptionFunc func(*dialConfig)
func (f dialOptionFunc) applyDial(c *dialConfig) { f(c) }
// WithDialRuntime selects the container runtime for the dial call.
func WithDialRuntime(rt types.Runtime) DialOption {
return dialOptionFunc(func(c *dialConfig) { c.runtime = rt })
}
// WithDialKubeConfig sets the kubeconfig for K8s runtime.
func WithDialKubeConfig(path string) DialOption {
return dialOptionFunc(func(c *dialConfig) { c.kubeConfig = path })
}
// WithDialNamespace sets the K8s namespace.
func WithDialNamespace(ns string) DialOption {
return dialOptionFunc(func(c *dialConfig) { c.namespace = ns })
}
// WithDialHTTPClient sets a custom HTTP client for proxy/VNC.
func WithDialHTTPClient(hc *http.Client) DialOption {
return dialOptionFunc(func(c *dialConfig) { c.httpClient = hc })
}
type dialConfig struct {
runtime types.Runtime
ports types.Ports
kubeConfig string
namespace string
httpClient *http.Client
userPorts types.Ports
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
func dialGRPC(target string) (*grpc.ClientConn, error) {
return grpc.NewClient(target,
grpc.WithTransportCredentials(insecure.NewCredentials()),
grpc.WithKeepaliveParams(keepalive.ClientParameters{
Time: 20 * time.Second,
Timeout: 5 * time.Second,
PermitWithoutStream: true,
}),
)
}
// ---------------------------------------------------------------------------
// ServerInfo discovery (shared by DialRemote / DialTunnel)
// ---------------------------------------------------------------------------
type discoveredInfo struct {
Capabilities map[string]bool
System types.SystemInfo
Version string
}
func discoverInfo(conn *grpc.ClientConn, cfg *dialConfig) (*discoveredInfo, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
client := sipb.NewServerInfoClient(conn)
resp, err := client.GetInfo(ctx, &sipb.GetInfoRequest{})
if err != nil {
return nil, err
}
up := cfg.userPorts
if p := int(resp.Ports["http"]); p > 0 && up.HTTP == 0 {
cfg.ports.HTTP = p
}
if p := int(resp.Ports["docker"]); p > 0 && up.Docker == 0 {
cfg.ports.Docker = p
}
if p := int(resp.Ports["vnc"]); p > 0 && up.VNC == 0 {
cfg.ports.VNC = p
}
if p := int(resp.Ports["k8s"]); p > 0 && up.K8s == 0 {
cfg.ports.K8s = p
}
caps := resp.Capabilities
if caps == nil {
caps = make(map[string]bool)
}
var sys types.SystemInfo
if s := resp.System; s != nil {
sys = types.SystemInfo{
OS: s.Os,
Arch: s.Arch,
Hostname: s.Hostname,
NumCPU: int(s.NumCpu),
TotalMem: s.TotalMem,
Shell: s.Shell,
TempDir: s.TempDir,
}
}
return &discoveredInfo{
Capabilities: caps,
System: sys,
Version: resp.Version,
}, nil
}