6408c140fc
- Introduce a new "tunnel" scheme for the Tai client, allowing connections through Yao's reverse proxy. - Enhance the Tai registry to manage tunnel-connected nodes and their ports. - Add WebSocket and reverse proxy routes for tunnel connections in the OpenAPI server. - Implement tunnel-specific proxy and VNC handling to facilitate communication with containerized environments. - Update gRPC environment variable handling to support tunnel connections. These changes improve the Tai service's flexibility and connectivity options, enabling better integration with remote and containerized environments.
499 lines
12 KiB
Go
499 lines
12 KiB
Go
package registry
|
|
|
|
import (
|
|
"log/slog"
|
|
"net"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
// newTestRegistry creates a standalone registry for testing (bypasses global singleton).
|
|
func newTestRegistry() *Registry {
|
|
return &Registry{
|
|
nodes: make(map[string]*TaiNode),
|
|
pending: make(map[string]*pendingChannel),
|
|
logger: slog.Default(),
|
|
}
|
|
}
|
|
|
|
func TestRegister_SetsFieldsAndOnline(t *testing.T) {
|
|
r := newTestRegistry()
|
|
node := &TaiNode{
|
|
TaiID: "tai-001",
|
|
MachineID: "m-abc",
|
|
Version: "1.0.0",
|
|
Mode: "tunnel",
|
|
Ports: map[string]int{"grpc": 9100},
|
|
}
|
|
r.Register(node)
|
|
|
|
snap, ok := r.Get("tai-001")
|
|
if !ok {
|
|
t.Fatal("expected node to exist after Register")
|
|
}
|
|
if snap.Status != "online" {
|
|
t.Errorf("Status = %q, want online", snap.Status)
|
|
}
|
|
if snap.MachineID != "m-abc" {
|
|
t.Errorf("MachineID = %q, want m-abc", snap.MachineID)
|
|
}
|
|
if snap.ConnectedAt.IsZero() {
|
|
t.Error("ConnectedAt should be set")
|
|
}
|
|
}
|
|
|
|
func TestRegister_Overwrite(t *testing.T) {
|
|
r := newTestRegistry()
|
|
r.Register(&TaiNode{TaiID: "tai-001", Version: "1.0"})
|
|
r.Register(&TaiNode{TaiID: "tai-001", Version: "2.0"})
|
|
|
|
snap, ok := r.Get("tai-001")
|
|
if !ok {
|
|
t.Fatal("node should exist")
|
|
}
|
|
if snap.Version != "2.0" {
|
|
t.Errorf("Version = %q, want 2.0 after re-register", snap.Version)
|
|
}
|
|
}
|
|
|
|
func TestUnregister_RemovesNode(t *testing.T) {
|
|
r := newTestRegistry()
|
|
r.Register(&TaiNode{TaiID: "tai-001"})
|
|
r.Unregister("tai-001")
|
|
|
|
if _, ok := r.Get("tai-001"); ok {
|
|
t.Error("expected node to be removed after Unregister")
|
|
}
|
|
}
|
|
|
|
func TestUnregister_Nonexistent(t *testing.T) {
|
|
r := newTestRegistry()
|
|
r.Unregister("ghost")
|
|
}
|
|
|
|
func TestGet_NotFound(t *testing.T) {
|
|
r := newTestRegistry()
|
|
if _, ok := r.Get("missing"); ok {
|
|
t.Error("expected false for missing node")
|
|
}
|
|
}
|
|
|
|
func TestList_Empty(t *testing.T) {
|
|
r := newTestRegistry()
|
|
if got := r.List(); len(got) != 0 {
|
|
t.Errorf("List() = %d items, want 0", len(got))
|
|
}
|
|
}
|
|
|
|
func TestList_MultipleNodes(t *testing.T) {
|
|
r := newTestRegistry()
|
|
r.Register(&TaiNode{TaiID: "a"})
|
|
r.Register(&TaiNode{TaiID: "b"})
|
|
r.Register(&TaiNode{TaiID: "c"})
|
|
|
|
list := r.List()
|
|
if len(list) != 3 {
|
|
t.Errorf("List() = %d items, want 3", len(list))
|
|
}
|
|
|
|
ids := map[string]bool{}
|
|
for _, snap := range list {
|
|
ids[snap.TaiID] = true
|
|
}
|
|
for _, id := range []string{"a", "b", "c"} {
|
|
if !ids[id] {
|
|
t.Errorf("missing node %q in List()", id)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestSnapshot_DeepCopy(t *testing.T) {
|
|
r := newTestRegistry()
|
|
r.Register(&TaiNode{
|
|
TaiID: "tai-001",
|
|
Ports: map[string]int{"grpc": 9100, "http": 8080},
|
|
})
|
|
|
|
snap, _ := r.Get("tai-001")
|
|
snap.Ports["grpc"] = 0
|
|
|
|
snap2, _ := r.Get("tai-001")
|
|
if snap2.Ports["grpc"] != 9100 {
|
|
t.Error("snapshot modification leaked into registry node")
|
|
}
|
|
}
|
|
|
|
func TestUpdatePing(t *testing.T) {
|
|
r := newTestRegistry()
|
|
r.Register(&TaiNode{TaiID: "tai-001"})
|
|
time.Sleep(10 * time.Millisecond)
|
|
|
|
r.UpdatePing("tai-001")
|
|
snap, _ := r.Get("tai-001")
|
|
if snap.LastPing.Before(snap.ConnectedAt) {
|
|
t.Error("LastPing should be after ConnectedAt")
|
|
}
|
|
}
|
|
|
|
func TestUpdatePing_NonexistentNode(t *testing.T) {
|
|
r := newTestRegistry()
|
|
r.UpdatePing("ghost")
|
|
}
|
|
|
|
func TestWriteControlJSON_NoNode(t *testing.T) {
|
|
r := newTestRegistry()
|
|
err := r.WriteControlJSON("missing", map[string]string{"type": "test"})
|
|
if err == nil {
|
|
t.Fatal("expected error for missing node")
|
|
}
|
|
}
|
|
|
|
func TestWriteControlJSON_NilConn(t *testing.T) {
|
|
r := newTestRegistry()
|
|
r.Register(&TaiNode{TaiID: "tai-001"})
|
|
err := r.WriteControlJSON("tai-001", map[string]string{"type": "test"})
|
|
if err == nil {
|
|
t.Fatal("expected error for nil ControlConn")
|
|
}
|
|
}
|
|
|
|
func TestRequestChannel_NotFound(t *testing.T) {
|
|
r := newTestRegistry()
|
|
_, _, err := r.RequestChannel("ghost", 9100)
|
|
if err == nil {
|
|
t.Fatal("expected error for missing node")
|
|
}
|
|
}
|
|
|
|
func TestRequestChannel_DirectMode(t *testing.T) {
|
|
r := newTestRegistry()
|
|
r.Register(&TaiNode{TaiID: "tai-001", Mode: "direct"})
|
|
_, _, err := r.RequestChannel("tai-001", 9100)
|
|
if err == nil {
|
|
t.Fatal("expected error for direct-mode node")
|
|
}
|
|
}
|
|
|
|
func TestAcceptDataChannel_NotPending(t *testing.T) {
|
|
r := newTestRegistry()
|
|
pipe1, pipe2 := net.Pipe()
|
|
defer pipe1.Close()
|
|
defer pipe2.Close()
|
|
|
|
err := r.AcceptDataChannel("unknown-channel", "tai-001", pipe1)
|
|
if err == nil {
|
|
t.Fatal("expected error for non-pending channel")
|
|
}
|
|
}
|
|
|
|
func TestAcceptDataChannel_TaiIDMismatch(t *testing.T) {
|
|
r := newTestRegistry()
|
|
|
|
resultCh := make(chan net.Conn, 1)
|
|
timer := time.AfterFunc(5*time.Second, func() {})
|
|
r.mu.Lock()
|
|
r.pending["ch-001"] = &pendingChannel{taiID: "tai-owner", result: resultCh, timer: timer}
|
|
r.mu.Unlock()
|
|
|
|
pipe1, pipe2 := net.Pipe()
|
|
defer pipe1.Close()
|
|
defer pipe2.Close()
|
|
|
|
err := r.AcceptDataChannel("ch-001", "tai-intruder", pipe1)
|
|
if err == nil {
|
|
t.Fatal("expected error for tai_id mismatch")
|
|
}
|
|
}
|
|
|
|
func TestAcceptDataChannel_Success(t *testing.T) {
|
|
r := newTestRegistry()
|
|
|
|
resultCh := make(chan net.Conn, 1)
|
|
timer := time.AfterFunc(5*time.Second, func() {})
|
|
r.mu.Lock()
|
|
r.pending["ch-002"] = &pendingChannel{taiID: "tai-001", result: resultCh, timer: timer}
|
|
r.mu.Unlock()
|
|
|
|
pipe1, pipe2 := net.Pipe()
|
|
defer pipe2.Close()
|
|
|
|
if err := r.AcceptDataChannel("ch-002", "tai-001", pipe1); err != nil {
|
|
t.Fatalf("AcceptDataChannel: %v", err)
|
|
}
|
|
|
|
select {
|
|
case conn := <-resultCh:
|
|
if conn == nil {
|
|
t.Fatal("expected non-nil conn")
|
|
}
|
|
conn.Close()
|
|
case <-time.After(time.Second):
|
|
t.Fatal("timeout waiting for conn on resultCh")
|
|
}
|
|
}
|
|
|
|
func TestGenerateChannelID_Unique(t *testing.T) {
|
|
seen := make(map[string]bool)
|
|
for i := 0; i < 100; i++ {
|
|
id, err := generateChannelID()
|
|
if err != nil {
|
|
t.Fatalf("generateChannelID: %v", err)
|
|
}
|
|
if len(id) != 64 {
|
|
t.Errorf("len = %d, want 64 hex chars", len(id))
|
|
}
|
|
if seen[id] {
|
|
t.Fatalf("duplicate channel ID: %s", id)
|
|
}
|
|
seen[id] = true
|
|
}
|
|
}
|
|
|
|
func TestBridgeTCP(t *testing.T) {
|
|
a1, a2 := net.Pipe()
|
|
b1, b2 := net.Pipe()
|
|
|
|
go bridgeTCP(a2, b1)
|
|
|
|
msg := []byte("hello tunnel")
|
|
go func() {
|
|
a1.Write(msg)
|
|
a1.Close()
|
|
}()
|
|
|
|
buf := make([]byte, 64)
|
|
n, _ := b2.Read(buf)
|
|
if string(buf[:n]) != "hello tunnel" {
|
|
t.Errorf("got %q, want %q", buf[:n], "hello tunnel")
|
|
}
|
|
b2.Close()
|
|
}
|
|
|
|
func TestConcurrentRegisterGet(t *testing.T) {
|
|
r := newTestRegistry()
|
|
var wg sync.WaitGroup
|
|
|
|
for i := 0; i < 50; i++ {
|
|
wg.Add(2)
|
|
id := "tai-" + string(rune('A'+i%26))
|
|
|
|
go func() {
|
|
defer wg.Done()
|
|
r.Register(&TaiNode{TaiID: id, Mode: "tunnel"})
|
|
}()
|
|
|
|
go func() {
|
|
defer wg.Done()
|
|
r.Get(id)
|
|
r.List()
|
|
}()
|
|
}
|
|
|
|
wg.Wait()
|
|
}
|
|
|
|
func TestWriteControlJSON_Success(t *testing.T) {
|
|
done := make(chan map[string]string, 1)
|
|
|
|
srv := newWSServer(func(conn *websocket.Conn) {
|
|
var msg map[string]string
|
|
conn.ReadJSON(&msg)
|
|
done <- msg
|
|
conn.Close()
|
|
})
|
|
defer srv.Close()
|
|
|
|
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
|
|
wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
|
if err != nil {
|
|
t.Fatalf("dial: %v", err)
|
|
}
|
|
|
|
r := newTestRegistry()
|
|
r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn})
|
|
|
|
payload := map[string]string{"type": "test", "data": "hello"}
|
|
if err := r.WriteControlJSON("tai-001", payload); err != nil {
|
|
t.Fatalf("WriteControlJSON: %v", err)
|
|
}
|
|
|
|
select {
|
|
case got := <-done:
|
|
if got["type"] != "test" {
|
|
t.Errorf("type = %q, want test", got["type"])
|
|
}
|
|
if got["data"] != "hello" {
|
|
t.Errorf("data = %q, want hello", got["data"])
|
|
}
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("timeout waiting for server to receive message")
|
|
}
|
|
}
|
|
|
|
func TestRequestChannel_Success(t *testing.T) {
|
|
openCh := make(chan map[string]interface{}, 1)
|
|
|
|
srv := newWSServer(func(conn *websocket.Conn) {
|
|
var msg map[string]interface{}
|
|
conn.ReadJSON(&msg)
|
|
openCh <- msg
|
|
time.Sleep(time.Second)
|
|
conn.Close()
|
|
})
|
|
defer srv.Close()
|
|
|
|
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
|
|
wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
|
if err != nil {
|
|
t.Fatalf("dial: %v", err)
|
|
}
|
|
|
|
r := newTestRegistry()
|
|
r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn})
|
|
|
|
channelID, resultCh, err := r.RequestChannel("tai-001", 9100)
|
|
if err != nil {
|
|
t.Fatalf("RequestChannel: %v", err)
|
|
}
|
|
if channelID == "" {
|
|
t.Fatal("channelID should not be empty")
|
|
}
|
|
if len(channelID) != 64 {
|
|
t.Errorf("channelID len = %d, want 64", len(channelID))
|
|
}
|
|
if resultCh == nil {
|
|
t.Fatal("resultCh should not be nil")
|
|
}
|
|
|
|
select {
|
|
case cmd := <-openCh:
|
|
if cmd["type"] != "open" {
|
|
t.Errorf("cmd type = %v, want open", cmd["type"])
|
|
}
|
|
if cmd["channel_id"] != channelID {
|
|
t.Errorf("cmd channel_id = %v, want %s", cmd["channel_id"], channelID)
|
|
}
|
|
if tp, ok := cmd["target_port"].(float64); !ok || int(tp) != 9100 {
|
|
t.Errorf("cmd target_port = %v, want 9100", cmd["target_port"])
|
|
}
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("timeout waiting for open command")
|
|
}
|
|
}
|
|
|
|
func TestRequestChannel_NoControlConn(t *testing.T) {
|
|
r := newTestRegistry()
|
|
r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel"})
|
|
|
|
_, _, err := r.RequestChannel("tai-001", 9100)
|
|
if err == nil {
|
|
t.Fatal("expected error for nil ControlConn")
|
|
}
|
|
}
|
|
|
|
func TestOpenLocalListener_Success(t *testing.T) {
|
|
r := newTestRegistry()
|
|
|
|
controlCh := make(chan map[string]interface{}, 1)
|
|
srv := newWSServer(func(conn *websocket.Conn) {
|
|
for {
|
|
var msg map[string]interface{}
|
|
if err := conn.ReadJSON(&msg); err != nil {
|
|
return
|
|
}
|
|
controlCh <- msg
|
|
}
|
|
})
|
|
defer srv.Close()
|
|
|
|
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
|
|
wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
|
if err != nil {
|
|
t.Fatalf("dial: %v", err)
|
|
}
|
|
|
|
r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn})
|
|
|
|
ln, err := r.OpenLocalListener("tai-001", 9100)
|
|
if err != nil {
|
|
t.Fatalf("OpenLocalListener: %v", err)
|
|
}
|
|
defer ln.Close()
|
|
|
|
addr := ln.Addr().String()
|
|
if addr == "" {
|
|
t.Fatal("listener address should not be empty")
|
|
}
|
|
if !strings.HasPrefix(addr, "127.0.0.1:") {
|
|
t.Errorf("addr = %q, want 127.0.0.1:*", addr)
|
|
}
|
|
|
|
conn, err := net.DialTimeout("tcp", addr, time.Second)
|
|
if err != nil {
|
|
t.Fatalf("connect to local listener: %v", err)
|
|
}
|
|
defer conn.Close()
|
|
|
|
select {
|
|
case cmd := <-controlCh:
|
|
if cmd["type"] != "open" {
|
|
t.Errorf("open cmd type = %v, want open", cmd["type"])
|
|
}
|
|
if _, ok := cmd["channel_id"].(string); !ok {
|
|
t.Error("open cmd missing channel_id")
|
|
}
|
|
if tp, ok := cmd["target_port"].(float64); !ok || int(tp) != 9100 {
|
|
t.Errorf("target_port = %v, want 9100", cmd["target_port"])
|
|
}
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("timeout waiting for open command from local listener")
|
|
}
|
|
}
|
|
|
|
func TestOpenLocalListener_NodeNotFound(t *testing.T) {
|
|
r := newTestRegistry()
|
|
_, err := r.OpenLocalListener("ghost", 9100)
|
|
if err == nil {
|
|
t.Fatal("expected error for missing node")
|
|
}
|
|
}
|
|
|
|
func newWSServer(handler func(*websocket.Conn)) *httptest.Server {
|
|
up := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
|
|
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
conn, err := up.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
handler(conn)
|
|
}))
|
|
}
|
|
|
|
func TestNodeSnapshot_AuthInfo(t *testing.T) {
|
|
r := newTestRegistry()
|
|
r.Register(&TaiNode{
|
|
TaiID: "tai-001",
|
|
Auth: AuthInfo{
|
|
Subject: "user123",
|
|
ClientID: "tai-001",
|
|
Scope: "tai:tunnel",
|
|
},
|
|
})
|
|
|
|
snap, ok := r.Get("tai-001")
|
|
if !ok {
|
|
t.Fatal("node not found")
|
|
}
|
|
if snap.Auth.Subject != "user123" {
|
|
t.Errorf("Auth.Subject = %q, want user123", snap.Auth.Subject)
|
|
}
|
|
if snap.Auth.Scope != "tai:tunnel" {
|
|
t.Errorf("Auth.Scope = %q, want tai:tunnel", snap.Auth.Scope)
|
|
}
|
|
}
|