inp2p/internal/server/server_test.go

package server

import (
	"fmt"
	"log"
	"net/http"
	"testing"
	"time"

	"github.com/openp2p-cn/inp2p/pkg/config"
	"github.com/openp2p-cn/inp2p/pkg/nat"
	"github.com/openp2p-cn/inp2p/pkg/protocol"
	"github.com/openp2p-cn/inp2p/pkg/signal"
	"github.com/gorilla/websocket"
)

func TestLoginFlow(t *testing.T) {
	// Start server
	cfg := config.DefaultServerConfig()
	cfg.WSPort = 29300
	cfg.Token = 999

	srv := New(cfg)
	mux := http.NewServeMux()
	mux.HandleFunc("/ws", srv.HandleWS)
	go http.ListenAndServe(fmt.Sprintf(":%d", cfg.WSPort), mux)
	time.Sleep(200 * time.Millisecond)

	// Connect as client manually
	ws, _, err := websocket.DefaultDialer.Dial(fmt.Sprintf("ws://127.0.0.1:%d/ws", cfg.WSPort), nil)
	if err != nil {
		t.Fatal(err)
	}
	conn := signal.NewConn(ws)
	defer conn.Close()

	// Start read loop in background
	go conn.ReadLoop()

	// Send login
	loginReq := protocol.LoginReq{
		Node:    "testNode",
		Token:   999,
		Version: "test",
		NATType: protocol.NATCone,
	}

	rspData, err := conn.Request(
		protocol.MsgLogin, protocol.SubLoginReq, loginReq,
		protocol.MsgLogin, protocol.SubLoginRsp,
		5*time.Second,
	)
	if err != nil {
		t.Fatalf("login request failed: %v", err)
	}

	var rsp protocol.LoginRsp
	protocol.DecodePayload(rspData, &rsp)
	if rsp.Error != 0 {
		t.Fatalf("login error: %d %s", rsp.Error, rsp.Detail)
	}
	log.Printf("Login OK: node=%s", rsp.Node)

	// Verify node is registered
	time.Sleep(100 * time.Millisecond)
	nodes := srv.GetOnlineNodes()
	if len(nodes) != 1 {
		t.Fatalf("expected 1 node, got %d", len(nodes))
	}
	if nodes[0].Name != "testNode" {
		t.Fatalf("expected testNode, got %s", nodes[0].Name)
	}

	srv.Stop()
}

func TestTwoClientsWithSTUN(t *testing.T) {
	cfg := config.DefaultServerConfig()
	cfg.WSPort = 29301
	cfg.STUNUDP1 = 29382
	cfg.STUNUDP2 = 29384
	cfg.STUNTCP1 = 29380
	cfg.STUNTCP2 = 29381
	cfg.Token = 888

	// STUN
	stunQuit := make(chan struct{})
	defer close(stunQuit)
	go nat.ServeUDPSTUN(cfg.STUNUDP1, stunQuit)
	go nat.ServeUDPSTUN(cfg.STUNUDP2, stunQuit)
	go nat.ServeTCPSTUN(cfg.STUNTCP1, stunQuit)
	go nat.ServeTCPSTUN(cfg.STUNTCP2, stunQuit)

	srv := New(cfg)
	srv.StartCleanup()
	mux := http.NewServeMux()
	mux.HandleFunc("/ws", srv.HandleWS)
	go http.ListenAndServe(fmt.Sprintf(":%d", cfg.WSPort), mux)
	time.Sleep(300 * time.Millisecond)

	// NAT detect
	natResult := nat.Detect("127.0.0.1", cfg.STUNUDP1, cfg.STUNUDP2, cfg.STUNTCP1, cfg.STUNTCP2)
	log.Printf("NAT: type=%s publicIP=%s", natResult.Type, natResult.PublicIP)

	// Client A
	connectClient := func(name string, relay bool) *signal.Conn {
		ws, _, err := websocket.DefaultDialer.Dial(fmt.Sprintf("ws://127.0.0.1:%d/ws", cfg.WSPort), nil)
		if err != nil {
			t.Fatalf("dial %s: %v", name, err)
		}
		conn := signal.NewConn(ws)
		go conn.ReadLoop()

		rspData, err := conn.Request(
			protocol.MsgLogin, protocol.SubLoginReq,
			protocol.LoginReq{Node: name, Token: 888, Version: "test", NATType: natResult.Type, RelayEnabled: relay},
			protocol.MsgLogin, protocol.SubLoginRsp,
			5*time.Second,
		)
		if err != nil {
			t.Fatalf("login %s: %v", name, err)
		}
		var rsp protocol.LoginRsp
		protocol.DecodePayload(rspData, &rsp)
		if rsp.Error != 0 {
			t.Fatalf("login %s error: %s", name, rsp.Detail)
		}
		log.Printf("%s login ok", name)
		return conn
	}

	connA := connectClient("nodeA", true)
	defer connA.Close()
	connB := connectClient("nodeB", false)
	defer connB.Close()

	time.Sleep(200 * time.Millisecond)
	nodes := srv.GetOnlineNodes()
	if len(nodes) != 2 {
		t.Fatalf("expected 2 nodes, got %d", len(nodes))
	}

	// Test relay node discovery
	relays := srv.GetRelayNodes("", "nodeB")
	if len(relays) != 1 || relays[0].Name != "nodeA" {
		t.Fatalf("expected nodeA as relay, got %v", relays)
	}
	log.Printf("Relay nodes: %v", relays[0].Name)

	srv.Stop()
}