inp2p/pkg/relay/relay_test.go

package relay

import (
	"fmt"
	"net"
	"sync"
	"testing"
	"time"

	"github.com/openp2p-cn/inp2p/pkg/auth"
)

func TestRelayBridge(t *testing.T) {
	token := auth.MakeToken("test", "pass")
	mgr := NewManager(29700, true, false, 10, token, 10)
	if err := mgr.Start(); err != nil {
		t.Fatal(err)
	}
	defer mgr.Stop()

	sessionID := "test-session-1"
	totp := auth.GenTOTP(token, time.Now().Unix())

	var wg sync.WaitGroup
	var connA, connB net.Conn
	var errA, errB error

	// Peer A connects as "from"
	wg.Add(1)
	go func() {
		defer wg.Done()
		connA, errA = ConnectToRelay(
			fmt.Sprintf("127.0.0.1:%d", 29700),
			sessionID, "from", "nodeA", totp,
		)
	}()

	// Peer B connects as "to" after a short delay
	wg.Add(1)
	go func() {
		defer wg.Done()
		time.Sleep(200 * time.Millisecond)
		connB, errB = ConnectToRelay(
			fmt.Sprintf("127.0.0.1:%d", 29700),
			sessionID, "to", "nodeB", totp,
		)
	}()

	wg.Wait()

	if errA != nil {
		t.Fatalf("connA error: %v", errA)
	}
	if errB != nil {
		t.Fatalf("connB error: %v", errB)
	}
	defer connA.Close()
	defer connB.Close()

	// Test data flow: A → B
	msg := []byte("hello through relay")
	connA.Write(msg)

	buf := make([]byte, 256)
	connB.SetReadDeadline(time.Now().Add(3 * time.Second))
	n, err := connB.Read(buf)
	if err != nil {
		t.Fatalf("read from B: %v", err)
	}
	if string(buf[:n]) != string(msg) {
		t.Errorf("got %q, want %q", buf[:n], msg)
	}

	// Test data flow: B → A
	reply := []byte("relay pong")
	connB.Write(reply)

	connA.SetReadDeadline(time.Now().Add(3 * time.Second))
	n, err = connA.Read(buf)
	if err != nil {
		t.Fatalf("read from A: %v", err)
	}
	if string(buf[:n]) != string(reply) {
		t.Errorf("got %q, want %q", buf[:n], reply)
	}

	// Verify session count
	if mgr.ActiveSessions() != 1 {
		t.Errorf("active sessions: got %d want 1", mgr.ActiveSessions())
	}

	t.Logf("✅ Relay bridge OK: A↔B bidirectional, %d active sessions", mgr.ActiveSessions())
}

func TestRelayLargeData(t *testing.T) {
	token := auth.MakeToken("test", "pass")
	mgr := NewManager(29701, true, false, 10, token, 10)
	if err := mgr.Start(); err != nil {
		t.Fatal(err)
	}
	defer mgr.Stop()

	sessionID := "test-large-data"
	totp := auth.GenTOTP(token, time.Now().Unix())

	var wg sync.WaitGroup
	var connA, connB net.Conn

	wg.Add(2)
	go func() {
		defer wg.Done()
		var err error
		connA, err = ConnectToRelay("127.0.0.1:29701", sessionID, "from", "bigA", totp)
		if err != nil {
			t.Errorf("connA: %v", err)
		}
	}()
	go func() {
		defer wg.Done()
		time.Sleep(100 * time.Millisecond)
		var err error
		connB, err = ConnectToRelay("127.0.0.1:29701", sessionID, "to", "bigB", totp)
		if err != nil {
			t.Errorf("connB: %v", err)
		}
	}()
	wg.Wait()

	if connA == nil || connB == nil {
		t.Fatal("connection failed")
	}
	defer connA.Close()
	defer connB.Close()

	// Send 1MB through relay
	const dataSize = 1024 * 1024
	data := make([]byte, dataSize)
	for i := range data {
		data[i] = byte(i % 256)
	}

	wg.Add(1)
	go func() {
		defer wg.Done()
		connA.Write(data)
	}()

	// Read exact amount on B side
	received := make([]byte, dataSize)
	total := 0
	connB.SetReadDeadline(time.Now().Add(10 * time.Second))
	for total < dataSize {
		n, err := connB.Read(received[total:])
		if err != nil {
			t.Fatalf("read at %d: %v", total, err)
		}
		total += n
	}

	wg.Wait()

	if len(received) != len(data) {
		t.Fatalf("size mismatch: got %d want %d", len(received), len(data))
	}
	for i := 0; i < len(data); i++ {
		if received[i] != data[i] {
			t.Fatalf("data mismatch at byte %d", i)
			break
		}
	}
	t.Logf("✅ 1MB relay transfer OK")
}

func TestRelayAuthDenied(t *testing.T) {
	token := auth.MakeToken("real", "token")
	mgr := NewManager(29702, true, false, 10, token, 10)
	if err := mgr.Start(); err != nil {
		t.Fatal(err)
	}
	defer mgr.Stop()

	// Use wrong TOTP
	wrongToken := auth.GenTOTP(auth.MakeToken("wrong", "creds"), time.Now().Unix())
	_, err := ConnectToRelay("127.0.0.1:29702", "bad-session", "from", "badNode", wrongToken)
	if err == nil {
		t.Fatal("expected auth denied, got success")
	}
	t.Logf("✅ Auth denied correctly: %v", err)
}