package main

import (
	"context"
	"errors"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"syscall"
	"testing"
	"time"
)

// how a fakeTarget responds to a dial.
type dialMode int

const (
	modeUp   dialMode = iota // connect to the echo server
	modeDown                 // fail immediately
	modeHang                 // block until the caller gives up, like a routable but dead peer
)

// fakeTarget stands in for a tailnet target. Its echo server always runs; the
// mode only gates *new* dials. That models the failure this is meant to catch:
// a tailnet peer disappears, new dials fail or hang, but connections already
// established just sit there with no error from the userspace TCP stack.
type fakeTarget struct {
	ln    net.Listener
	dials atomic.Int64 // every dial, probe or forwarded alike

	mu       sync.Mutex
	mode     dialMode
	sink     bool // swallow input and never reply, instead of echoing
	accepted []net.Conn
}

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, mode: modeUp}
	t.Cleanup(func() { ln.Close() })
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			ft.mu.Lock()
			ft.accepted = append(ft.accepted, c)
			sink := ft.sink
			ft.mu.Unlock()
			go func() {
				if sink {
					io.Copy(io.Discard, c) // read forever, never reply
				} else {
					io.Copy(c, c)
				}
				c.Close()
			}()
		}
	}()
	return ft
}

// killAccepted hard-closes every connection the target has accepted, as if the
// service process were killed. reset picks RST (SO_LINGER 0) over a clean FIN.
func (f *fakeTarget) killAccepted(reset bool) {
	f.mu.Lock()
	conns := f.accepted
	f.accepted = nil
	f.mu.Unlock()
	for _, c := range conns {
		if reset {
			c.(*net.TCPConn).SetLinger(0)
		}
		c.Close()
	}
}

// setSink makes the target read and never reply, so traffic flows one way only.
func (f *fakeTarget) setSink(v bool) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.sink = v
}

func (f *fakeTarget) setMode(m dialMode) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.mode = m
}

func (f *fakeTarget) setReachable(v bool) {
	if v {
		f.setMode(modeUp)
	} else {
		f.setMode(modeDown)
	}
}

func (f *fakeTarget) Dial(ctx context.Context, network, addr string) (net.Conn, error) {
	f.dials.Add(1)
	f.mu.Lock()
	mode := f.mode
	f.mu.Unlock()
	switch mode {
	case modeDown:
		return nil, errors.New("no route to host")
	case modeHang:
		<-ctx.Done()
		return nil, ctx.Err()
	}
	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()
	return startProxyWith(t, ft, local, 20*time.Millisecond, time.Second, 0)
}

// 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.
//
// Nothing polls while the target is up, so the idle timeout is what catches
// this: the target vanished without any TCP signal, and the connection simply
// goes quiet. The teardown must still be a reset, not a clean EOF.
func TestSilentTargetLossClosesIdleConnection(t *testing.T) {
	ft := newFakeTarget(t)
	local := freePort(t)
	startProxyWith(t, ft, local, 30*time.Second, time.Second, 250*time.Millisecond)
	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 there
	// is no TCP signal at all -- the connection just stops carrying traffic.
	ft.setReachable(false)

	// The client's blocking read must return, and as a reset: a plain EOF here
	// would tell the client the stream ended normally.
	c.SetReadDeadline(time.Now().Add(5 * time.Second))
	_, err = c.Read(buf)
	if err == nil {
		t.Fatal("read succeeded after target loss; want the connection closed")
	}
	if isTimeout(err) {
		t.Fatal("read blocked after target loss; the idle timeout never fired")
	}
	if !isReset(err) {
		t.Errorf("read err = %v, want a connection reset (a clean EOF would falsely signal a complete stream)", err)
	}
}

// An idle connection is not evidence that the target is unhealthy, so reaping
// one must not take the local port down with it.
func TestIdleTimeoutLeavesPortOpen(t *testing.T) {
	ft := newFakeTarget(t)
	local := freePort(t)
	startProxyWith(t, ft, local, 30*time.Second, time.Second, 200*time.Millisecond)
	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()

	// Never send anything; let the watchdog reap it.
	c.SetReadDeadline(time.Now().Add(5 * time.Second))
	buf := make([]byte, 1)
	if _, err := c.Read(buf); err == nil || isTimeout(err) {
		t.Fatalf("idle connection was not reaped (err: %v)", err)
	}

	// The target is perfectly healthy, so the port must still be accepting.
	// A long probe interval means a wrongly-closed port would stay closed.
	if localPortRefused(local) {
		t.Fatal("port closed after an idle reap; an unused connection says nothing about the target")
	}
}

// Traffic in one direction must keep the connection alive even though the other
// direction is silent, or every upload and every long download would be reaped
// mid-stream. The target here never replies, so client->target is the only
// activity there is -- an idle check that watched just one direction would kill
// this connection.
func TestOneWayTrafficIsNotIdle(t *testing.T) {
	ft := newFakeTarget(t)
	ft.setSink(true)
	local := freePort(t)
	idle := 200 * time.Millisecond
	startProxyWith(t, ft, local, 30*time.Second, time.Second, idle)
	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()

	// Write steadily for well over the idle window without ever reading.
	deadline := time.Now().Add(4 * idle)
	for time.Now().Before(deadline) {
		c.SetWriteDeadline(time.Now().Add(time.Second))
		if _, err := c.Write([]byte("x")); err != nil {
			t.Fatalf("connection died under one-way traffic within %s: %v", idle, err)
		}
		time.Sleep(idle / 8)
	}

	// Still writable, i.e. still alive after 4x the idle window of one-way use.
	c.SetWriteDeadline(time.Now().Add(time.Second))
	if _, err := c.Write([]byte("x")); err != nil {
		t.Fatalf("connection unusable after sustained one-way traffic: %v", err)
	}
}

// The point of the change: a healthy target must never be dialled again after
// the one probe that decides whether to bind at startup.
func TestNoProbingWhileTargetIsUp(t *testing.T) {
	ft := newFakeTarget(t)
	local := freePort(t)
	interval := 20 * time.Millisecond
	startProxyWith(t, ft, local, interval, time.Second, 0)
	waitFor(t, "local port to accept", func() bool { return !localPortRefused(local) })

	// The readiness check above opens a real connection, which forwards and so
	// dials the target. Let that finish before snapshotting, or its dial lands
	// inside the measurement window and looks like a probe.
	time.Sleep(5 * interval)

	settled := ft.dials.Load()
	time.Sleep(20 * interval) // twenty probe intervals' worth of opportunity
	if got := ft.dials.Load(); got != settled {
		t.Errorf("target was dialled %d more times while up; want 0 (probing must stop once up)", got-settled)
	}
	t.Logf("%d dials total while up over %s", ft.dials.Load(), 20*interval)
}

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

func isReset(err error) bool {
	return errors.Is(err, syscall.ECONNRESET)
}

// A target process that dies drops its connections with a FIN or an RST. The
// local client must see that within a round trip -- it must NOT have to wait
// for the reachability probe, which is the slow fallback for a peer that
// vanishes without any TCP signal at all.
func TestTargetHardCloseIsNoticedImmediately(t *testing.T) {
	for _, tc := range []struct {
		name  string
		reset bool
	}{
		{"fin", false},
		{"rst", true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			ft := newFakeTarget(t)
			local := freePort(t)
			// Probe interval far longer than the assertion window, so a pass
			// can only mean the per-connection path noticed.
			p := &proxy{
				dial:     ft,
				local:    local,
				target:   "target:1234",
				interval: 30 * time.Second,
				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) }()
			defer func() { cancel(); <-done }()

			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)
			}

			start := time.Now()
			ft.killAccepted(tc.reset)

			c.SetReadDeadline(time.Now().Add(5 * time.Second))
			_, err = c.Read(buf)
			elapsed := time.Since(start)
			if err == nil {
				t.Fatal("read succeeded after target hard-close")
			}
			if isTimeout(err) {
				t.Fatalf("client never noticed target hard-close (blocked %s)", elapsed)
			}
			t.Logf("client noticed after %s (err: %v)", elapsed, err)
			if elapsed > time.Second {
				t.Errorf("client took %s to notice; want sub-second", elapsed)
			}
		})
	}
}

// startProxyWith runs a proxy with explicit timings, for tests that need the
// probe loop held back so a pass can only come from the connection path.
func startProxyWith(t *testing.T, ft *fakeTarget, local string, interval, timeout, idle time.Duration) *proxy {
	t.Helper()
	p := &proxy{
		dial:     ft,
		local:    local,
		target:   "target:1234",
		interval: interval,
		timeout:  timeout,
		idle:     idle,
		recheck:  make(chan struct{}, 1),
		conns:    make(map[net.Conn]struct{}),
		trace:    true, // so a failing test leaves a usable log behind
	}
	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
}

// A target that is routable but dead makes the dial block rather than fail.
// The client must not be left parked on a socket whose handler is stuck in
// Dial -- nothing else would ever reclaim it, because the probe loop only
// tears down live connections when a probe FAILS, and a target that comes back
// makes the probe succeed.
func TestHangingDialDoesNotStrandClient(t *testing.T) {
	ft := newFakeTarget(t)
	local := freePort(t)
	// Probe interval far beyond the test, so the probe loop cannot be what
	// rescues the connection.
	startProxyWith(t, ft, local, 30*time.Second, 300*time.Millisecond, 0)
	waitFor(t, "local port to accept", func() bool { return !localPortRefused(local) })

	ft.setMode(modeHang)

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

	start := time.Now()
	c.SetReadDeadline(time.Now().Add(5 * time.Second))
	buf := make([]byte, 1)
	_, err = c.Read(buf)
	elapsed := time.Since(start)
	if err == nil {
		t.Fatal("read succeeded against a hanging target")
	}
	if isTimeout(err) {
		t.Fatalf("client stranded on a hanging dial (blocked %s); dial must be bounded", elapsed)
	}
	t.Logf("client released after %s (err: %v)", elapsed, err)
	if elapsed > 2*time.Second {
		t.Errorf("client held for %s; want release near the probe timeout", elapsed)
	}
	// Recovery is deliberately not asserted here: the 30s interval that makes
	// this test meaningful is also long enough that rebinding would be a race
	// against the probe loop's next tick. TestUnreachableTargetRefusesConnections
	// and TestPortReturnsAfterTransientFailure cover it.
}

// What the user hit: a client reconnects the instant its connection breaks,
// beating the probe. It must be refused, not accepted into a dead forward.
func TestFailedConnectionClosesPortBeforeProbe(t *testing.T) {
	ft := newFakeTarget(t)
	local := freePort(t)
	// Long interval and a long probe timeout: once the target starts hanging,
	// a probe cannot complete within the assertion window, so the port closing
	// can only be the connection path doing it.
	startProxyWith(t, ft, local, 30*time.Second, 5*time.Second, 0)
	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)
	}

	ft.setMode(modeHang)
	start := time.Now()
	ft.killAccepted(true) // target drops the connection, as on process death

	deadline := time.Now().Add(2 * time.Second)
	for !localPortRefused(local) {
		if time.Now().After(deadline) {
			t.Fatal("local port still accepting after the target dropped a connection")
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Logf("port closed %s after target dropped the connection", time.Since(start))
}

// Suspending on one failure must not strand the port: if the target is in fact
// fine, the probe that suspension schedules has to bring it straight back.
func TestPortReturnsAfterTransientFailure(t *testing.T) {
	ft := newFakeTarget(t)
	local := freePort(t)
	// Long interval: recovery must come from the probe suspension schedules,
	// not from the next scheduled tick.
	startProxyWith(t, ft, local, 30*time.Second, time.Second, 0)
	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)
	}
	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)
	}

	// The target stays healthy; this one connection just breaks.
	ft.killAccepted(true)
	c.Close()

	waitFor(t, "local port to come back without waiting for a tick", func() bool {
		return !localPortRefused(local)
	})
}