vibeproxy/internal/streams/tcp_test.go

package streams

import (
	"context"
	"io"
	"log/slog"
	"net"
	"testing"
	"time"

	"github.com/owen/vibeproxy/internal/models"
)

func TestTCPProxyEndToEnd(t *testing.T) {
	// Echo backend.
	backendLn, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer backendLn.Close()
	go func() {
		for {
			c, err := backendLn.Accept()
			if err != nil {
				return
			}
			go func() {
				defer c.Close()
				io.Copy(c, c)
			}()
		}
	}()

	backendAddr := backendLn.Addr().(*net.TCPAddr)
	model := models.Stream{
		Protocol: models.StreamTCP, ListenIP: "127.0.0.1", ListenPort: 0,
		DestinationHost: "127.0.0.1", DestinationPort: backendAddr.Port,
		ConnectTimeoutMS: 2000, IdleTimeoutMS: 5000,
	}

	// Bind an ephemeral port ourselves so we know it up front, matching
	// what NewTCPProxy will do.
	probe, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	model.ListenPort = probe.Addr().(*net.TCPAddr).Port
	probe.Close()

	log := slog.New(slog.NewTextHandler(io.Discard, nil))
	proxy := NewTCPProxy(model, log)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	if err := proxy.Start(ctx); err != nil {
		t.Fatal(err)
	}
	defer proxy.Stop()

	time.Sleep(50 * time.Millisecond)
	conn, err := net.Dial("tcp", net.JoinHostPort("127.0.0.1", itoa(model.ListenPort)))
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()

	if _, err := conn.Write([]byte("ping")); err != nil {
		t.Fatal(err)
	}
	conn.SetReadDeadline(time.Now().Add(2 * time.Second))
	buf := make([]byte, 4)
	if _, err := io.ReadFull(conn, buf); err != nil {
		t.Fatal(err)
	}
	if string(buf) != "ping" {
		t.Fatalf("expected echo of ping, got %q", buf)
	}
}

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	neg := n < 0
	if neg {
		n = -n
	}
	var b []byte
	for n > 0 {
		b = append([]byte{byte('0' + n%10)}, b...)
		n /= 10
	}
	if neg {
		b = append([]byte{'-'}, b...)
	}
	return string(b)
}