tsproxy/main_test.go

package main

import (
	"context"
	"errors"
	"io"
	"net"
	"sync"
	"testing"
	"time"
)

// fakeTarget stands in for a tailnet target. Its echo server always runs; the
// reachable flag only gates *new* dials. That models the failure this is meant
// to catch: a tailnet peer disappears, new dials fail, but connections already
// established just hang forever with no error from the userspace TCP stack.
type fakeTarget struct {
	ln net.Listener

	mu        sync.Mutex
	reachable bool
}

func newFakeTarget(t *testing.T) *fakeTarget {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("fake target listen: %v", err)
	}
	ft := &fakeTarget{ln: ln, reachable: true}
	t.Cleanup(func() { ln.Close() })
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			go func() { io.Copy(c, c); c.Close() }()
		}
	}()
	return ft
}

func (f *fakeTarget) setReachable(v bool) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.reachable = v
}

func (f *fakeTarget) Dial(ctx context.Context, network, addr string) (net.Conn, error) {
	f.mu.Lock()
	ok := f.reachable
	f.mu.Unlock()
	if !ok {
		return nil, errors.New("no route to host")
	}
	var d net.Dialer
	return d.DialContext(ctx, network, f.ln.Addr().String())
}

// freePort returns an address that is currently bindable.
func freePort(t *testing.T) string {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("free port: %v", err)
	}
	addr := ln.Addr().String()
	ln.Close()
	return addr
}

func startProxy(t *testing.T, ft *fakeTarget, local string) *proxy {
	t.Helper()
	p := &proxy{
		dial:     ft,
		local:    local,
		target:   "target:1234",
		interval: 20 * time.Millisecond,
		timeout:  time.Second,
		recheck:  make(chan struct{}, 1),
		conns:    make(map[net.Conn]struct{}),
	}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() { defer close(done); p.run(ctx) }()
	t.Cleanup(func() {
		cancel()
		<-done
	})
	return p
}

// waitFor polls cond until it holds or the deadline passes.
func waitFor(t *testing.T, what string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s", what)
}

func localPortRefused(local string) bool {
	c, err := net.DialTimeout("tcp", local, 500*time.Millisecond)
	if err != nil {
		return true
	}
	c.Close()
	return false
}

// An unreachable target must mean a refused connect, not a successful connect
// followed by an immediate EOF.
func TestUnreachableTargetRefusesConnections(t *testing.T) {
	ft := newFakeTarget(t)
	ft.setReachable(false)
	local := freePort(t)
	startProxy(t, ft, local)

	// Give the probe loop several ticks; the port must never come up.
	time.Sleep(200 * time.Millisecond)
	if !localPortRefused(local) {
		t.Fatal("connect succeeded while target was unreachable; want refused")
	}

	// And it must start accepting once the target comes back.
	ft.setReachable(true)
	waitFor(t, "local port to accept after target recovery", func() bool {
		return !localPortRefused(local)
	})
}

func TestForwardsData(t *testing.T) {
	ft := newFakeTarget(t)
	local := freePort(t)
	startProxy(t, ft, local)
	waitFor(t, "local port to accept", func() bool { return !localPortRefused(local) })

	c, err := net.Dial("tcp", local)
	if err != nil {
		t.Fatalf("dial local: %v", err)
	}
	defer c.Close()

	if _, err := c.Write([]byte("ping")); err != nil {
		t.Fatalf("write: %v", err)
	}
	buf := make([]byte, 4)
	c.SetReadDeadline(time.Now().Add(3 * time.Second))
	if _, err := io.ReadFull(c, buf); err != nil {
		t.Fatalf("read echo: %v", err)
	}
	if string(buf) != "ping" {
		t.Fatalf("echo = %q, want %q", buf, "ping")
	}
}

// The reported bug: when the proxy loses the target, an already-established
// local socket must be closed so the client's read fails and it reconnects,
// rather than blocking forever on a half-dead connection.
func TestTargetLossClosesLocalConnection(t *testing.T) {
	ft := newFakeTarget(t)
	local := freePort(t)
	startProxy(t, ft, local)
	waitFor(t, "local port to accept", func() bool { return !localPortRefused(local) })

	c, err := net.Dial("tcp", local)
	if err != nil {
		t.Fatalf("dial local: %v", err)
	}
	defer c.Close()

	// Establish that the connection is live and forwarding.
	if _, err := c.Write([]byte("ping")); err != nil {
		t.Fatalf("write: %v", err)
	}
	buf := make([]byte, 4)
	c.SetReadDeadline(time.Now().Add(3 * time.Second))
	if _, err := io.ReadFull(c, buf); err != nil {
		t.Fatalf("read echo: %v", err)
	}

	// Target vanishes. The echo server still holds the far end open, so
	// nothing but the probe loop can notice.
	ft.setReachable(false)

	// The client's blocking read must return, and reasonably promptly.
	c.SetReadDeadline(time.Now().Add(3 * time.Second))
	if _, err := c.Read(buf); err == nil {
		t.Fatal("read succeeded after target loss; want the connection closed")
	} else if errors.Is(err, io.EOF) || isTimeout(err) {
		if isTimeout(err) {
			t.Fatal("read blocked after target loss; local socket was never closed")
		}
	}

	// New connects must be refused too.
	waitFor(t, "local port to be refused", func() bool { return localPortRefused(local) })
}

func isTimeout(err error) bool {
	var ne net.Error
	return errors.As(err, &ne) && ne.Timeout()
}