Trojan-UoT & UDP-nameserver: Fix forgotten release buffer; UDP dispatcher: Simplified and optimized (#5050)

transport/internet/udp/dispatcher.go

@@ -22,8 +22,24 @@ type ResponseCallback func(ctx context.Context, packet *udp.Packet)
 
 type connEntry struct {
 	link   *transport.Link
-	timer  signal.ActivityUpdater
+	timer  *signal.ActivityTimer
 	cancel context.CancelFunc
+	closed bool
+}
+
+func (c *connEntry) Close() error {
+	c.timer.SetTimeout(0)
+	return nil
+}
+
+func (c *connEntry) terminate() {
+	if c.closed {
+		panic("terminate called more than once")
+	}
+	c.closed = true
+	c.cancel()
+	common.Interrupt(c.link.Reader)
+	common.Interrupt(c.link.Writer)
 }
 
 type Dispatcher struct {
@@ -32,6 +48,7 @@ type Dispatcher struct {
 	dispatcher routing.Dispatcher
 	callback   ResponseCallback
 	callClose  func() error
+	closed     bool
 }
 
 func NewDispatcher(dispatcher routing.Dispatcher, callback ResponseCallback) *Dispatcher {
@@ -44,13 +61,9 @@ func NewDispatcher(dispatcher routing.Dispatcher, callback ResponseCallback) *Di
 func (v *Dispatcher) RemoveRay() {
 	v.Lock()
 	defer v.Unlock()
-	v.removeRay()
-}
-
-func (v *Dispatcher) removeRay() {
+	v.closed = true
 	if v.conn != nil {
-		common.Interrupt(v.conn.link.Reader)
-		common.Close(v.conn.link.Writer)
+		v.conn.Close()
 		v.conn = nil
 	}
 }
@@ -59,35 +72,34 @@ func (v *Dispatcher) getInboundRay(ctx context.Context, dest net.Destination) (*
 	v.Lock()
 	defer v.Unlock()
 
+	if v.closed {
+		return nil, errors.New("dispatcher is closed")
+	}
+
 	if v.conn != nil {
-		return v.conn, nil
+		if v.conn.closed {
+			v.conn = nil
+		} else {
+			return v.conn, nil
+		}
 	}
 
 	errors.LogInfo(ctx, "establishing new connection for ", dest)
 
 	ctx, cancel := context.WithCancel(ctx)
-	entry := &connEntry{}
-	removeRay := func() {
-		v.Lock()
-		defer v.Unlock()
-		// sometimes the entry is already removed by others, don't close again
-		if entry == v.conn {
-			cancel()
-			v.removeRay()
-		}
-	}
-	timer := signal.CancelAfterInactivity(ctx, removeRay, time.Minute)
 
 	link, err := v.dispatcher.Dispatch(ctx, dest)
 	if err != nil {
+		cancel()
 		return nil, errors.New("failed to dispatch request to ", dest).Base(err)
 	}
 
-	*entry = connEntry{
+	entry := &connEntry{
 		link:   link,
-		timer:  timer,
-		cancel: removeRay,
+		cancel: cancel,
 	}
+
+	entry.timer = signal.CancelAfterInactivity(ctx, entry.terminate, time.Minute)
 	v.conn = entry
 	go handleInput(ctx, entry, dest, v.callback, v.callClose)
 	return entry, nil
@@ -106,7 +118,7 @@ func (v *Dispatcher) Dispatch(ctx context.Context, destination net.Destination,
 	if outputStream != nil {
 		if err := outputStream.WriteMultiBuffer(buf.MultiBuffer{payload}); err != nil {
 			errors.LogInfoInner(ctx, err, "failed to write first UDP payload")
-			conn.cancel()
+			conn.Close()
 			return
 		}
 	}
@@ -114,7 +126,7 @@ func (v *Dispatcher) Dispatch(ctx context.Context, destination net.Destination,
 
 func handleInput(ctx context.Context, conn *connEntry, dest net.Destination, callback ResponseCallback, callClose func() error) {
 	defer func() {
-		conn.cancel()
+		conn.Close()
 		if callClose != nil {
 			callClose()
 		}

transport/pipe/impl.go

@@ -200,16 +200,19 @@ func (p *pipe) Interrupt() {
 	p.Lock()
 	defer p.Unlock()
 
+	if !p.data.IsEmpty() {
+		buf.ReleaseMulti(p.data)
+		p.data = nil
+		if p.state == closed {
+			p.state = errord
+		}
+	}
+
 	if p.state == closed || p.state == errord {
 		return
 	}
 
 	p.state = errord
 
-	if !p.data.IsEmpty() {
-		buf.ReleaseMulti(p.data)
-		p.data = nil
-	}
-
 	common.Must(p.done.Close())
 }