fix(msg): close control connection on write failure (#5550)

* fix(msg): close control connection on write failure

* fix(msg): handle shutdown enqueue race and preserve constructor API

pkg/msg/handler.go

@@ -19,6 +19,7 @@ import (
 	"io"
 	"net"
 	"reflect"
+	"sync"
 
 	"github.com/fatedier/frp/pkg/proto/wire"
 )
@@ -111,14 +112,20 @@ type Dispatcher struct {
 
 	sendCh      chan Message
 	doneCh      chan struct{}
+	stopCh      chan struct{}
+	stopOnce    sync.Once
 	msgHandlers map[reflect.Type]func(Message)
 }
 
+// NewDispatcher uses Close, when rw implements io.Closer, to interrupt ReadMsg
+// after a write failure. Otherwise, the caller must arrange for ReadMsg to return
+// before Done can close.
 func NewDispatcher(rw ReadWriter) *Dispatcher {
 	return &Dispatcher{
 		rw:          rw,
 		sendCh:      make(chan Message, 100),
 		doneCh:      make(chan struct{}),
+		stopCh:      make(chan struct{}),
 		msgHandlers: make(map[reflect.Type]func(Message)),
 	}
 }
@@ -129,37 +136,75 @@ func (d *Dispatcher) Run() {
 	go d.readLoop()
 }
 
+// stop rejects new sends and interrupts ReadMsg when rw is closable. Done remains
+// owned by readLoop so the control owner cannot clean up while a handler is running.
+func (d *Dispatcher) stop() {
+	d.stopOnce.Do(func() {
+		close(d.stopCh)
+		if closer, ok := d.rw.(io.Closer); ok {
+			_ = closer.Close()
+		}
+	})
+}
+
 func (d *Dispatcher) sendLoop() {
 	for {
 		select {
-		case <-d.doneCh:
+		case <-d.stopCh:
 			return
 		case m := <-d.sendCh:
-			_ = d.rw.WriteMsg(m)
+			if err := d.rw.WriteMsg(m); err != nil {
+				d.stop()
+				return
+			}
 		}
 	}
 }
 
 func (d *Dispatcher) readLoop() {
+	defer close(d.doneCh)
+	defer d.stop()
 	for {
 		m, err := d.rw.ReadMsg()
 		if err != nil {
-			close(d.doneCh)
 			return
 		}
 
+		select {
+		case <-d.stopCh:
+			return
+		default:
+		}
+
 		if handler, ok := d.msgHandlers[reflect.TypeOf(m)]; ok {
 			handler(m)
 		}
 	}
 }
 
+// Send queues m for writing. A nil error does not guarantee delivery.
 func (d *Dispatcher) Send(m Message) error {
 	select {
-	case <-d.doneCh:
+	case <-d.stopCh:
+		return io.EOF
+	default:
+	}
+	return d.enqueue(m)
+}
+
+func (d *Dispatcher) enqueue(m Message) error {
+	select {
+	case <-d.stopCh:
 		return io.EOF
 	case d.sendCh <- m:
-		return nil
+		// Both cases can be ready if shutdown races with enqueueing. Do not
+		// report success when shutdown is already observable after the send.
+		select {
+		case <-d.stopCh:
+			return io.EOF
+		default:
+			return nil
+		}
 	}
 }