zonat/e2e/handshake_test.go

// 转发链路"慢失败"测试：两处历史缺陷的回归锚——
//  1. 节点转发握手无上限：agent 会话半死(收得到帧但永不应答)时观看端无限挂起
//     → HandshakeTimeout 到点必须切断观看端, 让其快速失败可重试;
//  2. agent 拨内网目标无上限：目标离线且防火墙丢包(不回 RST)时裸 Dial 吊 20s+
//     → DialTimeout 到点必须向节点回执失败, 观看端立即断开。
package e2e

import (
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"testing"
	"time"

	"git.zeroonesoft.cn/golib/zonat/agent"
	"git.zeroonesoft.cn/golib/zonat/internal/node"

	"github.com/xtaci/smux"
)

// deafAgent 手工模拟"半死 agent"：完成登录+注册隧道后, 对节点打开的数据流
// (CmdTarget)永不应答。返回节点为该隧道分配的公网端口。
func deafAgent(t *testing.T, n *node.Node, tunnelId, targetIp string, targetPort int) int {
	t.Helper()
	conn, err := net.DialTimeout("tcp", n.Addr(), 2*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { conn.Close() })

	tok, err := node.SignToken(n.JwtSecret, "deaf-agent", time.Hour)
	if err != nil {
		t.Fatal(err)
	}
	lf := agent.NewFrame(agent.FrameVersion, agent.CmdLogin, 0)
	if err := lf.Marshal(agent.Login{AgentId: "deaf-agent", Token: tok}); err != nil {
		t.Fatal(err)
	}
	if err := agent.WriteFrame(conn, lf); err != nil {
		t.Fatal(err)
	}
	rf, err := agent.ReadFrame(conn)
	if err != nil {
		t.Fatal(err)
	}
	ret := agent.Ret{}
	if err := rf.Unmarshal(&ret); err != nil || ret.Code != 0 {
		t.Fatalf("登录失败: %+v", ret)
	}

	sess, err := smux.Client(conn, agent.SmuxConfig())
	if err != nil {
		t.Fatal(err)
	}
	control, err := sess.OpenStream()
	if err != nil {
		t.Fatal(err)
	}
	reg := agent.NewFrame(agent.FrameVersion, agent.CmdRegisterTunnel, 0)
	if err := reg.Marshal(agent.Tunnel{Id: tunnelId, TargetIp: targetIp, TargetPort: targetPort, TTLSec: 0}); err != nil {
		t.Fatal(err)
	}
	if err := agent.WriteFrame(control, reg); err != nil {
		t.Fatal(err)
	}
	rr, err := agent.ReadFrame(control)
	if err != nil {
		t.Fatal(err)
	}
	rret := agent.Ret{}
	if err := rr.Unmarshal(&rret); err != nil || rret.Code != 0 {
		t.Fatalf("注册失败: %+v", rret)
	}
	return rret.Port
}

// assertViewerCutBy 断言观看端在时限内被节点断开(读到 EOF/重置, 而非超时)。
func assertViewerCutBy(t *testing.T, port int, within time.Duration) {
	t.Helper()
	viewer, err := net.DialTimeout("tcp", net.JoinHostPort("127.0.0.1", strconv.Itoa(port)), 2*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer viewer.Close()
	_ = viewer.SetReadDeadline(time.Now().Add(within))
	buf := make([]byte, 16)
	_, err = viewer.Read(buf)
	if err == nil {
		t.Fatalf("观看端读到了数据, 不应发生")
	}
	var netErr net.Error
	if errors.As(err, &netErr) && netErr.Timeout() {
		t.Fatalf("观看端 %v 内未被节点断开(超时切断未生效)", within)
	}
}

// TestHandshakeTimeoutCutsViewer 半死 agent: 节点必须在握手超时内切断观看端
// (修复前观看端无限挂起, 只能等 smux keepalive 判死整个会话 ~40s)。
func TestHandshakeTimeoutCutsViewer(t *testing.T) {
	target := startEchoServer(t)
	n := node.New()
	n.JwtSecret = "test-jwt-secret"
	n.BindTunnel = "127.0.0.1"
	n.HandshakeTimeout = 500 * time.Millisecond
	if err := n.Start("127.0.0.1:0"); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(n.Stop)

	host, portStr, _ := net.SplitHostPort(target)
	portNum, _ := strconv.Atoi(portStr)
	listen := deafAgent(t, n, fmt.Sprintf("deaf-%d", time.Now().UnixNano()), host, portNum)

	assertViewerCutBy(t, listen, 3*time.Second)
}

// TestAgentDialTimeoutCutsViewer 目标黑洞(不回 RST): agent 拨号到点回执失败,
// 观看端立即断开(修复前裸 Dial 在 Windows 吊 20s+, 观看端干等)。
func TestAgentDialTimeoutCutsViewer(t *testing.T) {
	// 本机网络若把 TEST-NET-3 判定为可达(隧道/虚拟网卡), 黑洞前提不成立
	if c, err := net.DialTimeout("tcp", "203.0.113.1:5900", 300*time.Millisecond); err == nil {
		c.Close()
		t.Skip("本机可将 TEST-NET-3 判定连通, 黑洞拨号场景不成立")
	}

	n, a := startEnv(t)
	a.DialTimeout = 300 * time.Millisecond
	id := fmt.Sprintf("blackhole-%d", time.Now().UnixNano())
	port, err := a.RegisterTunnel(buildTunnel(t, n, a, id, "203.0.113.1:5900", 0))
	if err != nil {
		t.Fatal(err)
	}

	assertViewerCutBy(t, port, 3*time.Second)
}

// startCloseAfterEchoServer 回显一段即关闭的目标：模拟 tvnserver 进程退出。
func startCloseAfterEchoServer(t *testing.T) string {
	t.Helper()
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { l.Close() })
	go func() {
		for {
			c, err := l.Accept()
			if err != nil {
				return
			}
			go func(c net.Conn) {
				defer c.Close()
				buf := make([]byte, 512)
				n, _ := c.Read(buf)
				if n > 0 {
					c.Write(buf[:n]) // 回显一次随即断开
				}
			}(c)
		}
	}()
	return l.Addr().String()
}

// TestTargetCloseCutsViewer 目标端→观看端断开同步：目标(如 tvnserver)关闭后,
// 观看端必须在时限内被同步断开, 不允许残留悬挂连接。
func TestTargetCloseCutsViewer(t *testing.T) {
	n, a := startEnv(t)
	target := startCloseAfterEchoServer(t)
	_, port := mustRegister(t, n, a, target)

	viewer, err := net.DialTimeout("tcp", net.JoinHostPort("127.0.0.1", strconv.Itoa(port)), 2*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer viewer.Close()
	_ = viewer.SetDeadline(time.Now().Add(3 * time.Second))
	payload := []byte("probe")
	if _, err := viewer.Write(payload); err != nil {
		t.Fatal(err)
	}
	got := make([]byte, len(payload))
	if _, err := io.ReadFull(viewer, got); err != nil {
		t.Fatal(err)
	}
	// 回显已到, 目标随即关闭 → 观看端应读到 EOF/重置(而非超时)
	buf := make([]byte, 16)
	_, err = viewer.Read(buf)
	if err == nil {
		t.Fatal("目标已关闭, 观看端却读到数据")
	}
	var netErr net.Error
	if errors.As(err, &netErr) && netErr.Timeout() {
		t.Fatal("目标关闭 3s 内未同步断开观看端")
	}
}