diff --git a/protocol/masque/client.go b/protocol/masque/client.go index 8403f4c53..f87fc3b86 100644 --- a/protocol/masque/client.go +++ b/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() } diff --git a/protocol/masque/server.go b/protocol/masque/server.go index f400a10a8..bd37395a7 100644 --- a/protocol/masque/server.go +++ b/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() } diff --git a/transport/device/device.go b/transport/device/device.go index 46c5041ec..c8074dcb2 100644 --- a/transport/device/device.go +++ b/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 diff --git a/transport/device/stack.go b/transport/device/stack.go index d92e66516..2c3cc6d10 100644 --- a/transport/device/stack.go +++ b/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 { diff --git a/transport/device/system_stack.go b/transport/device/system_stack.go index 960c0daaa..bcc30d3b8 100644 --- a/transport/device/system_stack.go +++ b/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 { diff --git a/transport/masque/client.go b/transport/masque/client.go index b7d0c3ff4..494ec8f3b 100644 --- a/transport/masque/client.go +++ b/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")) } diff --git a/transport/masque/server.go b/transport/masque/server.go index 605482ef0..422549cfc 100644 --- a/transport/masque/server.go +++ b/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")) }