From 3abfbf50aebbe3ba57dc749165edb56be6715c0a Mon Sep 17 00:00:00 2001
From: Anton Tolchanov <anton@tailscale.com>
Date: Tue, 28 Jan 2025 12:10:28 +0000
Subject: [PATCH] tsnet: return from Accept when the listener gets closed

Fixes #14808

Signed-off-by: Anton Tolchanov <anton@tailscale.com>
---
 tsnet/tsnet.go      |  7 ++++---
 tsnet/tsnet_test.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/tsnet/tsnet.go b/tsnet/tsnet.go
index 3505c9453..23a9f9a98 100644
--- a/tsnet/tsnet.go
+++ b/tsnet/tsnet.go
@@ -1286,11 +1286,12 @@ type listener struct {
 }
 
 func (ln *listener) Accept() (net.Conn, error) {
-	c, ok := <-ln.conn
-	if !ok {
+	select {
+	case c := <-ln.conn:
+		return c, nil
+	case <-ln.closedc:
 		return nil, fmt.Errorf("tsnet: %w", net.ErrClosed)
 	}
-	return c, nil
 }
 
 func (ln *listener) Addr() net.Addr { return addr{ln} }
diff --git a/tsnet/tsnet_test.go b/tsnet/tsnet_test.go
index 552e8dbee..0f245b015 100644
--- a/tsnet/tsnet_test.go
+++ b/tsnet/tsnet_test.go
@@ -667,6 +667,37 @@ func TestFunnel(t *testing.T) {
 	}
 }
 
+func TestListenerClose(t *testing.T) {
+	ctx := context.Background()
+	controlURL, _ := startControl(t)
+
+	s1, _, _ := startServer(t, ctx, controlURL, "s1")
+
+	ln, err := s1.Listen("tcp", ":8080")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	errc := make(chan error, 1)
+	go func() {
+		c, err := ln.Accept()
+		if c != nil {
+			c.Close()
+		}
+		errc <- err
+	}()
+
+	ln.Close()
+	select {
+	case err := <-errc:
+		if !errors.Is(err, net.ErrClosed) {
+			t.Errorf("unexpected error: %v", err)
+		}
+	case <-time.After(10 * time.Second):
+		t.Fatal("timeout waiting for Accept to return")
+	}
+}
+
 func dialIngressConn(from, to *Server, target string) (net.Conn, error) {
 	toLC := must.Get(to.LocalClient())
 	toStatus := must.Get(toLC.StatusWithoutPeers(context.Background()))