Apply backpressure to masque inbound packets

protocol/masque/client.go

@@ -205,6 +205,16 @@ func (c *ClientEndpoint) WriteInboundBuffers(packetBuffers []*buf.Buffer) error
 	return err
 }
 
+func (c *ClientEndpoint) WaitWriteInboundBuffers(packetBuffers []*buf.Buffer) error {
+	if !c.state.Load().configured {
+		buf.ReleaseMulti(packetBuffers)
+		return nil
+	}
+	err := c.device.WaitWriteInboundBuffers(packetBuffers)
+	buf.ReleaseMulti(packetBuffers)
+	return err
+}
+
 func (c *ClientEndpoint) FrontHeadroom() int {
 	return c.device.FrontHeadroom()
 }

protocol/masque/server.go

@@ -210,6 +210,12 @@ func (s *ServerEndpoint) WriteInboundBuffers(packetBuffers []*buf.Buffer) error
 	return err
 }
 
+func (s *ServerEndpoint) WaitWriteInboundBuffers(packetBuffers []*buf.Buffer) error {
+	err := s.device.WaitWriteInboundBuffers(packetBuffers)
+	buf.ReleaseMulti(packetBuffers)
+	return err
+}
+
 func (s *ServerEndpoint) FrontHeadroom() int {
 	return s.device.FrontHeadroom()
 }

transport/device/device.go

@@ -22,6 +22,7 @@ type Device interface {
 	Start() error
 	UpdateConfiguration(configuration Configuration) error
 	WriteInboundBuffers(packetBuffers []*buf.Buffer) error
+	WaitWriteInboundBuffers(packetBuffers []*buf.Buffer) error
 	SetPacketWriter(writer PacketWriter)
 	PortAddresses() (netip.Addr, netip.Addr)
 	PortMTU() uint32

transport/device/stack.go

@@ -73,11 +73,19 @@ func (d *stackDevice) WriteInboundBuffers(packetBuffers []*buf.Buffer) error {
 	return d.processInboundBuffers(packetBuffers, d.writeBuffers)
 }
 
+func (d *stackDevice) WaitWriteInboundBuffers(packetBuffers []*buf.Buffer) error {
+	return d.processInboundBuffers(packetBuffers, d.waitWriteBuffers)
+}
+
 func (d *stackDevice) writeBuffers(packetBuffers []*buf.Buffer) error {
 	_, err := d.memoryTun.WritePackets(common.Map(packetBuffers, (*buf.Buffer).Bytes))
 	return err
 }
 
+func (d *stackDevice) waitWriteBuffers(packetBuffers []*buf.Buffer) error {
+	return d.memoryTun.WaitWritePackets(common.Map(packetBuffers, (*buf.Buffer).Bytes))
+}
+
 func (d *stackDevice) DialContext(ctx context.Context, network string, destination M.Socksaddr) (net.Conn, error) {
 	bind, err := d.bindAddress(destination)
 	if err != nil {

transport/device/system_stack.go

@@ -55,10 +55,18 @@ func (d *systemStackDevice) UpdateConfiguration(configuration Configuration) err
 }
 
 func (d *systemStackDevice) WriteInboundBuffers(packetBuffers []*buf.Buffer) error {
-	return d.systemDevice.processInboundBuffers(packetBuffers, d.writeBuffers)
+	return d.systemDevice.processInboundBuffers(packetBuffers, func(unconsumedBuffers []*buf.Buffer) error {
+		return d.writeBuffers(unconsumedBuffers, d.stackDevice.writeBuffers)
+	})
 }
 
-func (d *systemStackDevice) writeBuffers(packetBuffers []*buf.Buffer) error {
+func (d *systemStackDevice) WaitWriteInboundBuffers(packetBuffers []*buf.Buffer) error {
+	return d.systemDevice.processInboundBuffers(packetBuffers, func(unconsumedBuffers []*buf.Buffer) error {
+		return d.writeBuffers(unconsumedBuffers, d.stackDevice.waitWriteBuffers)
+	})
+}
+
+func (d *systemStackDevice) writeBuffers(packetBuffers []*buf.Buffer, writeStackBuffers func(packetBuffers []*buf.Buffer) error) error {
 	addresses := d.systemDevice.configurationAddresses()
 	runStart := 0
 	runUsesSystemDevice := false
@@ -77,7 +85,7 @@ func (d *systemStackDevice) writeBuffers(packetBuffers []*buf.Buffer) error {
 			if runUsesSystemDevice {
 				err = d.systemDevice.writeBuffers(packetBuffers[runStart:i])
 			} else {
-				err = d.stackDevice.writeBuffers(packetBuffers[runStart:i])
+				err = writeStackBuffers(packetBuffers[runStart:i])
 			}
 			writeErr = E.Errors(writeErr, err)
 			runStart = i
@@ -92,7 +100,7 @@ func (d *systemStackDevice) writeBuffers(packetBuffers []*buf.Buffer) error {
 	if runUsesSystemDevice {
 		return E.Errors(writeErr, d.systemDevice.writeBuffers(packetBuffers[runStart:]))
 	}
-	return E.Errors(writeErr, d.stackDevice.writeBuffers(packetBuffers[runStart:]))
+	return E.Errors(writeErr, writeStackBuffers(packetBuffers[runStart:]))
 }
 
 func packetDestination(packet []byte) netip.Addr {

transport/masque/client.go

@@ -29,6 +29,7 @@ type Configuration struct {
 type ClientHandler interface {
 	UpdateConfiguration(configuration Configuration) error
 	WriteInboundBuffers(packetBuffers []*buf.Buffer) error
+	WaitWriteInboundBuffers(packetBuffers []*buf.Buffer) error
 	FrontHeadroom() int
 }
 
@@ -426,7 +427,7 @@ func (s *clientSession) handlePacket(buffer *buf.Buffer) {
 		}
 		return
 	}
-	err := s.client.handler.WriteInboundBuffers([]*buf.Buffer{buffer})
+	err := s.client.handler.WaitWriteInboundBuffers([]*buf.Buffer{buffer})
 	if err != nil {
 		s.client.logger.Debug(E.Cause(err, "write packet to device"))
 	}

transport/masque/server.go

@@ -22,6 +22,7 @@ import (
 
 type ServerHandler interface {
 	WriteInboundBuffers(packetBuffers []*buf.Buffer) error
+	WaitWriteInboundBuffers(packetBuffers []*buf.Buffer) error
 	FrontHeadroom() int
 	NewOutboundQueue(handler func(packetBuffers []*buf.Buffer)) *tun.OutboundQueue
 }
@@ -434,7 +435,7 @@ func (s *serverSession) handlePacket(buffer *buf.Buffer) {
 			break
 		}
 		if target == nil {
-			err := s.server.handler.WriteInboundBuffers([]*buf.Buffer{buffer})
+			err := s.server.handler.WaitWriteInboundBuffers([]*buf.Buffer{buffer})
 			if err != nil {
 				s.server.logger.DebugContext(s.ctx, E.Cause(err, "write packet to device"))
 			}