fix: Only hold tailnet.*Conn.Close() for a short duration (#4015)

* fix: Only hold `tailnet.*Conn.Close()` for a short duration

The long duration could be cause to a test deadlock.

* Add closed chan to listener struct

tailnet/conn.go

@@ -387,16 +387,17 @@ func (c *Conn) Closed() <-chan struct{} {
 // Close shuts down the Wireguard connection.
 func (c *Conn) Close() error {
 	c.mutex.Lock()
-	defer c.mutex.Unlock()
 	select {
 	case <-c.closed:
+		c.mutex.Unlock()
 		return nil
 	default:
 	}
+	close(c.closed)
 	for _, l := range c.listeners {
 		_ = l.closeNoLock()
 	}
-	close(c.closed)
+	c.mutex.Unlock()
 	_ = c.dialer.Close()
 	_ = c.magicConn.Close()
 	_ = c.netStack.Close()
@@ -406,6 +407,15 @@ func (c *Conn) Close() error {
 	return nil
 }
 
+func (c *Conn) isClosed() bool {
+	select {
+	case <-c.closed:
+		return true
+	default:
+		return false
+	}
+}
+
 // This and below is taken _mostly_ verbatim from Tailscale:
 // https://github.com/tailscale/tailscale/blob/c88bd53b1b7b2fcf7ba302f2e53dd1ce8c32dad4/tsnet/tsnet.go#L459-L494
 
@@ -422,9 +432,14 @@ func (c *Conn) Listen(network, addr string) (net.Listener, error) {
 		key:  lk,
 		addr: addr,
 
-		conn: make(chan net.Conn),
+		closed: make(chan struct{}),
+		conn:   make(chan net.Conn),
 	}
 	c.mutex.Lock()
+	if c.isClosed() {
+		c.mutex.Unlock()
+		return nil, xerrors.New("closed")
+	}
 	if c.listeners == nil {
 		c.listeners = map[listenKey]*listener{}
 	}
@@ -460,9 +475,12 @@ func (c *Conn) forwardTCP(conn net.Conn, port uint16) {
 	defer t.Stop()
 	select {
 	case ln.conn <- conn:
+		return
+	case <-ln.closed:
+	case <-c.closed:
 	case <-t.C:
-		_ = conn.Close()
 	}
+	_ = conn.Close()
 }
 
 func (c *Conn) forwardTCPToLocal(conn net.Conn, port uint16) {
@@ -506,15 +524,18 @@ type listenKey struct {
 }
 
 type listener struct {
-	s    *Conn
-	key  listenKey
-	addr string
-	conn chan net.Conn
+	s      *Conn
+	key    listenKey
+	addr   string
+	conn   chan net.Conn
+	closed chan struct{}
 }
 
 func (ln *listener) Accept() (net.Conn, error) {
-	c, ok := <-ln.conn
-	if !ok {
+	var c net.Conn
+	select {
+	case c = <-ln.conn:
+	case <-ln.closed:
 		return nil, xerrors.Errorf("wgnet: %w", net.ErrClosed)
 	}
 	return c, nil
@@ -530,7 +551,7 @@ func (ln *listener) Close() error {
 func (ln *listener) closeNoLock() error {
 	if v, ok := ln.s.listeners[ln.key]; ok && v == ln {
 		delete(ln.s.listeners, ln.key)
-		close(ln.conn)
+		close(ln.closed)
 	}
 	return nil
 }