Fix leaks and add test

transport/wireguard/client_bind.go

@@ -20,6 +20,7 @@ type ClientBind struct {
 	peerAddr   M.Socksaddr
 	connAccess sync.Mutex
 	conn       *wireConn
+	done       chan struct{}
 }
 
 func NewClientBind(ctx context.Context, dialer N.Dialer, peerAddr M.Socksaddr) *ClientBind {
@@ -63,6 +64,12 @@ func (c *ClientBind) connect() (*wireConn, error) {
 }
 
 func (c *ClientBind) Open(port uint16) (fns []conn.ReceiveFunc, actualPort uint16, err error) {
+	select {
+	case <-c.done:
+		err = net.ErrClosed
+		return
+	default:
+	}
 	return []conn.ReceiveFunc{c.receive}, 0, nil
 }
 
@@ -75,7 +82,12 @@ func (c *ClientBind) receive(b []byte) (n int, ep conn.Endpoint, err error) {
 	n, err = udpConn.Read(b)
 	if err != nil {
 		udpConn.Close()
-		err = &wireError{err}
+		select {
+		case <-c.done:
+		default:
+			err = &wireError{err}
+		}
+		return
 	}
 	ep = Endpoint(c.peerAddr)
 	return
@@ -85,6 +97,16 @@ func (c *ClientBind) Close() error {
 	c.connAccess.Lock()
 	defer c.connAccess.Unlock()
 	common.Close(common.PtrOrNil(c.conn))
+	if c.done == nil {
+		c.done = make(chan struct{})
+		return nil
+	}
+	select {
+	case <-c.done:
+		return net.ErrClosed
+	default:
+		close(c.done)
+	}
 	return nil
 }
 

transport/wireguard/device_stack.go

@@ -35,7 +35,6 @@ type StackDevice struct {
 	events     chan tun.Event
 	outbound   chan *stack.PacketBuffer
 	dispatcher stack.NetworkDispatcher
-	done       chan struct{}
 	addr4      tcpip.Address
 	addr6      tcpip.Address
 }
@@ -51,7 +50,6 @@ func NewStackDevice(localAddresses []netip.Prefix, mtu uint32) (*StackDevice, er
 		mtu:      mtu,
 		events:   make(chan tun.Event, 1),
 		outbound: make(chan *stack.PacketBuffer, 256),
-		done:     make(chan struct{}),
 	}
 	err := ipStack.CreateNIC(defaultNIC, (*wireEndpoint)(tunDevice))
 	if err != nil {
@@ -193,11 +191,11 @@ func (w *StackDevice) Events() chan tun.Event {
 
 func (w *StackDevice) Close() error {
 	select {
-	case <-w.done:
+	case <-w.events:
 		return os.ErrClosed
 	default:
+		close(w.events)
 	}
-	close(w.done)
 	w.stack.Close()
 	for _, endpoint := range w.stack.CleanupEndpoints() {
 		endpoint.Abort()

transport/wireguard/device_system.go

@@ -106,5 +106,11 @@ func (w *SystemDevice) Events() chan wgTun.Event {
 }
 
 func (w *SystemDevice) Close() error {
+	select {
+	case <-w.events:
+		return os.ErrClosed
+	default:
+		close(w.events)
+	}
 	return w.device.Close()
 }