store read size within the buffer carrier

illotum · illotum · commit be7596d4fb5d · 2026-03-18T14:51:29.000-07:00
Signed-off-by: Alex Valiushko &lt;alexvaliushko@tailscale.com&gt;
Change-Id: I48217f8f461b17a01901cee9ab64d45e6a6a6964
diff --git a/buffer/buffer.go b/buffer/buffer.go
@@ -46,6 +46,8 @@ type Recycler interface {
 type Buffer struct {
 	data     []byte
 	recycler Recycler
+
+	Size int // size of the read, excluding offset
 }
 
 // New creates a standalone Buffer.
diff --git a/conn/bind_std.go b/conn/bind_std.go
@@ -235,7 +235,6 @@ func (s *StdNetBind) receiveIP(
 	conn *net.UDPConn,
 	rxOffload bool,
 	bufs []*buffer.Buffer,
-	sizes []int,
 	eps []Endpoint,
 ) (n int, err error) {
 	msgs := s.getMessages()
@@ -288,8 +287,8 @@ func (s *StdNetBind) receiveIP(
 	}
 	for i := 0; i < numMsgs; i++ {
 		msg := &(*msgs)[i]
-		sizes[i] = msg.N
-		if sizes[i] == 0 {
+		bufs[i].Size = msg.N
+		if bufs[i].Size == 0 {
 			continue
 		}
 		addrPort := msg.Addr.(*net.UDPAddr).AddrPort()
@@ -301,14 +300,14 @@ func (s *StdNetBind) receiveIP(
 }
 
 func (s *StdNetBind) makeReceiveIPv4(pc *ipv4.PacketConn, conn *net.UDPConn, rxOffload bool) ReceiveFunc {
-	return func(bufs []*buffer.Buffer, sizes []int, eps []Endpoint) (n int, err error) {
-		return s.receiveIP(pc, conn, rxOffload, bufs, sizes, eps)
+	return func(bufs []*buffer.Buffer, eps []Endpoint) (n int, err error) {
+		return s.receiveIP(pc, conn, rxOffload, bufs, eps)
 	}
 }
 
 func (s *StdNetBind) makeReceiveIPv6(pc *ipv6.PacketConn, conn *net.UDPConn, rxOffload bool) ReceiveFunc {
-	return func(bufs []*buffer.Buffer, sizes []int, eps []Endpoint) (n int, err error) {
-		return s.receiveIP(pc, conn, rxOffload, bufs, sizes, eps)
+	return func(bufs []*buffer.Buffer, eps []Endpoint) (n int, err error) {
+		return s.receiveIP(pc, conn, rxOffload, bufs, eps)
 	}
 }
 
diff --git a/conn/bind_std_test.go b/conn/bind_std_test.go
@@ -19,13 +19,12 @@ func TestStdNetBindReceiveFuncAfterClose(t *testing.T) {
 	bind.Close()
 	bufs := make([]*buffer.Buffer, IdealBatchSize)
 	bufs[0] = buffer.New(make([]byte, buffer.MaxMessageSize), nil)
-	sizes := make([]int, IdealBatchSize)
 	eps := make([]Endpoint, IdealBatchSize)
 	for _, fn := range fns {
 		// The ReceiveFuncs must not access conn-related fields on StdNetBind
 		// unguarded. Close() nils the conn-related fields resulting in a panic
 		// if they violate the mutex.
-		fn(bufs, sizes, eps)
+		fn(bufs, eps)
 	}
 }
 
diff --git a/conn/bind_windows.go b/conn/bind_windows.go
@@ -417,20 +417,20 @@ retry:
 	return n, &ep, nil
 }
 
-func (bind *WinRingBind) receiveIPv4(bufs []*buffer.Buffer, sizes []int, eps []Endpoint) (int, error) {
+func (bind *WinRingBind) receiveIPv4(bufs []*buffer.Buffer, eps []Endpoint) (int, error) {
 	bind.mu.RLock()
 	defer bind.mu.RUnlock()
 	n, ep, err := bind.v4.Receive(bufs[0].Data(), &bind.isOpen)
-	sizes[0] = n
+	bufs[0].Size = n
 	eps[0] = ep
 	return 1, err
 }
 
-func (bind *WinRingBind) receiveIPv6(bufs []*buffer.Buffer, sizes []int, eps []Endpoint) (int, error) {
+func (bind *WinRingBind) receiveIPv6(bufs []*buffer.Buffer, eps []Endpoint) (int, error) {
 	bind.mu.RLock()
 	defer bind.mu.RUnlock()
 	n, ep, err := bind.v6.Receive(bufs[0].Data(), &bind.isOpen)
-	sizes[0] = n
+	bufs[0].Size = n
 	eps[0] = ep
 	return 1, err
 }
diff --git a/conn/bindtest/bindtest.go b/conn/bindtest/bindtest.go
@@ -95,7 +95,7 @@ func (c *ChannelBind) BatchSize() int { return 1 }
 func (c *ChannelBind) SetMark(mark uint32) error { return nil }
 
 func (c *ChannelBind) makeReceiveFunc(ch chan []byte) conn.ReceiveFunc {
-	return func(bufs []*buffer.Buffer, sizes []int, eps []conn.Endpoint) (n int, err error) {
+	return func(bufs []*buffer.Buffer, eps []conn.Endpoint) (n int, err error) {
 		select {
 		case <-c.closeSignal:
 			return 0, net.ErrClosed
@@ -104,7 +104,7 @@ func (c *ChannelBind) makeReceiveFunc(ch chan []byte) conn.ReceiveFunc {
 				bufs[0] = buffer.New(make([]byte, buffer.MaxMessageSize), nil)
 			}
 			copied := copy(bufs[0].Data(), rx)
-			sizes[0] = copied
+			bufs[0].Size = copied
 			eps[0] = c.target6
 			return 1, nil
 		}
diff --git a/conn/conn.go b/conn/conn.go
@@ -27,7 +27,7 @@ const (
 // sizes may be zero, and callers should ignore them. Callers must pass a sizes
 // and eps slice with a length greater than or equal to the length of packets.
 // These lengths must not exceed the length of the associated Bind.BatchSize().
-type ReceiveFunc func(bufs []*buffer.Buffer, sizes []int, eps []Endpoint) (n int, err error)
+type ReceiveFunc func(bufs []*buffer.Buffer, eps []Endpoint) (n int, err error)
 
 // A Bind listens on a port for both IPv6 and IPv4 UDP traffic.
 //
diff --git a/conn/conn_test.go b/conn/conn_test.go
@@ -13,7 +13,7 @@ import (
 
 func TestPrettyName(t *testing.T) {
 	var (
-		recvFunc ReceiveFunc = func(bufs []*buffer.Buffer, sizes []int, eps []Endpoint) (n int, err error) { return }
+		recvFunc ReceiveFunc = func(bufs []*buffer.Buffer, eps []Endpoint) (n int, err error) { return }
 	)
 
 	const want = "TestPrettyName"
diff --git a/device/device_test.go b/device/device_test.go
@@ -438,7 +438,7 @@ type fakeTUNDeviceSized struct {
 }
 
 func (t *fakeTUNDeviceSized) File() *os.File { return nil }
-func (t *fakeTUNDeviceSized) Read(bufs []*buffer.Buffer, sizes []int, offset int) (n int, err error) {
+func (t *fakeTUNDeviceSized) Read(bufs []*buffer.Buffer, offset int) (n int, err error) {
 	return 0, nil
 }
 func (t *fakeTUNDeviceSized) Write(bufs [][]byte, offset int) (int, error) { return 0, nil }
diff --git a/device/receive.go b/device/receive.go
@@ -86,7 +86,6 @@ func (device *Device) RoutineReceiveIncoming(maxBatchSize int, recv conn.Receive
 
 	var (
 		bufs        = make([]*buffer.Buffer, maxBatchSize) // nil entries; recv allocates
-		sizes       = make([]int, maxBatchSize)
 		err         error
 		count       int
 		endpoints   = make([]conn.Endpoint, maxBatchSize)
@@ -99,7 +98,7 @@ func (device *Device) RoutineReceiveIncoming(maxBatchSize int, recv conn.Receive
 	}()
 
 	for {
-		count, err = recv(bufs, sizes, endpoints)
+		count, err = recv(bufs, endpoints)
 		if err != nil {
 			if errors.Is(err, net.ErrClosed) {
 				return
@@ -118,14 +117,14 @@ func (device *Device) RoutineReceiveIncoming(maxBatchSize int, recv conn.Receive
 		deathSpiral = 0
 
 		// handle each packet in the batch
-		for i, size := range sizes[:count] {
-			if size < MinMessageSize {
+		for i := 0; i < count; i++ {
+			if bufs[i].Size < MinMessageSize {
 				continue
 			}
 
 			// check size of packet
 
-			packet := bufs[i].Data()[:size]
+			packet := bufs[i].Data()[:bufs[i].Size]
 			msgType := binary.LittleEndian.Uint32(packet[:4])
 
 			switch msgType {
diff --git a/device/send.go b/device/send.go
@@ -231,7 +231,6 @@ func (device *Device) RoutineReadFromTUN() {
 		bufs        = make([]*buffer.Buffer, batchSize)
 		elemsByPeer = make(map[*Peer]*QueueOutboundElementsContainer, batchSize)
 		count       = 0
-		sizes       = make([]int, batchSize)
 		offset      = MessageEncapsulatingTransportSize + MessageTransportHeaderSize
 	)
 
@@ -240,10 +239,10 @@ func (device *Device) RoutineReadFromTUN() {
 	}()
 
 	for {
-		count, readErr = device.tun.device.Read(bufs, sizes, offset)
+		count, readErr = device.tun.device.Read(bufs, offset)
 
 		for i := 0; i < count; i++ {
-			packet := bufs[i].Data()[offset : offset+sizes[i]]
+			packet := bufs[i].Data()[offset : offset+bufs[i].Size]
 			if len(packet) < 1 {
 				continue
 			}
diff --git a/tun/netstack/tun.go b/tun/netstack/tun.go
@@ -120,7 +120,7 @@ func (tun *netTun) Events() <-chan tun.Event {
 	return tun.events
 }
 
-func (tun *netTun) Read(bufs []*wgbuf.Buffer, sizes []int, offset int) (int, error) {
+func (tun *netTun) Read(bufs []*wgbuf.Buffer, offset int) (int, error) {
 	view, ok := <-tun.incomingPacket
 	if !ok {
 		return 0, os.ErrClosed
@@ -133,7 +133,7 @@ func (tun *netTun) Read(bufs []*wgbuf.Buffer, sizes []int, offset int) (int, err
 	if err != nil {
 		return 0, err
 	}
-	sizes[0] = n
+	bufs[0].Size = n
 	return 1, nil
 }
 
diff --git a/tun/offload.go b/tun/offload.go
@@ -107,7 +107,8 @@ func GSOSplit(in []byte, options GSOOptions, outBufs []*buffer.Buffer, pool buff
 			in[cSumAt], in[cSumAt+1] = 0, 0
 			binary.BigEndian.PutUint16(in[cSumAt:], ^Checksum(in[options.CsumStart:], initial))
 		}
-		copy(outBufs[0].Data()[outOffset:], in)
+		n := copy(outBufs[0].Data()[outOffset:], in)
+		outBufs[0].Size = n
 		return 1, nil
 	}
 
@@ -171,6 +172,7 @@ func GSOSplit(in []byte, options GSOOptions, outBufs []*buffer.Buffer, pool buff
 		if outBufs[i] == nil {
 			outBufs[i], _ = pool.Get(needed)
 		}
+		outBufs[i].Size = totalLen
 		out := outBufs[i].Data()[outOffset:]
 
 		copy(out, in[:iphLen])
diff --git a/tun/offload_linux_test.go b/tun/offload_linux_test.go
@@ -248,8 +248,8 @@ func Test_handleVirtioRead(t *testing.T) {
 				t.Fatalf("got %d packets, wanted %d", n, len(tt.wantLens))
 			}
 			for i := range tt.wantLens {
-				if tt.wantLens[i]+offset != out[i].Len() {
-					t.Fatalf("wantLens[%d]: %d != outSizes: %d", i, tt.wantLens[i]+offset, out[i].Len())
+				if tt.wantLens[i] != out[i].Size {
+					t.Fatalf("wantLens[%d]: %d != outSizes: %d", i, tt.wantLens[i], out[i].Size)
 				}
 			}
 		})
diff --git a/tun/tun.go b/tun/tun.go
@@ -29,7 +29,7 @@ type Device interface {
 	// Nil entries in bufs are allocated by the implementation.
 	// A nonzero offset can be used to instruct the Device on where to begin
 	// reading into each element of the bufs slice.
-	Read(bufs []*buffer.Buffer, sizes []int, offset int) (n int, err error)
+	Read(bufs []*buffer.Buffer, offset int) (n int, err error)
 
 	// Write one or more packets to the device (without any additional headers).
 	// On a successful write it returns the number of packets written. A nonzero
diff --git a/tun/tun_darwin.go b/tun/tun_darwin.go
@@ -218,7 +218,7 @@ func (tun *NativeTun) Events() <-chan Event {
 	return tun.events
 }
 
-func (tun *NativeTun) Read(bufs []*buffer.Buffer, sizes []int, offset int) (int, error) {
+func (tun *NativeTun) Read(bufs []*buffer.Buffer, offset int) (int, error) {
 	// TODO: the BSDs look very similar in Read() and Write(). They should be
 	// collapsed, with platform-specific files containing the varying parts of
 	// their implementations.
@@ -234,7 +234,7 @@ func (tun *NativeTun) Read(bufs []*buffer.Buffer, sizes []int, offset int) (int,
 		if n < 4 {
 			return 0, err
 		}
-		sizes[0] = n - 4
+		bufs[0].Size = n - 4
 		return 1, err
 	}
 }
diff --git a/tun/tun_freebsd.go b/tun/tun_freebsd.go
@@ -334,7 +334,7 @@ func (tun *NativeTun) Events() <-chan Event {
 	return tun.events
 }
 
-func (tun *NativeTun) Read(bufs []*buffer.Buffer, sizes []int, offset int) (int, error) {
+func (tun *NativeTun) Read(bufs []*buffer.Buffer, offset int) (int, error) {
 	select {
 	case err := <-tun.errors:
 		return 0, err
@@ -347,7 +347,7 @@ func (tun *NativeTun) Read(bufs []*buffer.Buffer, sizes []int, offset int) (int,
 		if n < 4 {
 			return 0, err
 		}
-		sizes[0] = n - 4
+		bufs[0].Size = n - 4
 		return 1, err
 	}
 }
diff --git a/tun/tun_linux.go b/tun/tun_linux.go
@@ -429,7 +429,7 @@ func handleVirtioRead(in []byte, bufs []*buffer.Buffer, pool buffer.Source, offs
 	return GSOSplit(in, options, bufs, pool, offset)
 }
 
-func (tun *NativeTun) Read(bufs []*buffer.Buffer, sizes []int, offset int) (int, error) {
+func (tun *NativeTun) Read(bufs []*buffer.Buffer, offset int) (int, error) {
 	tun.readOpMu.Lock()
 	defer tun.readOpMu.Unlock()
 	select {
@@ -453,16 +453,9 @@ func (tun *NativeTun) Read(bufs []*buffer.Buffer, sizes []int, offset int) (int,
 			return 0, err
 		}
 		if tun.vnetHdr {
-			count, err := handleVirtioRead(readInto[:n], bufs, tun.bufPool, offset)
-			if err != nil {
-				return 0, err
-			}
-			for i := 0; i < count; i++ {
-				sizes[i] = bufs[i].Len() - offset
-			}
-			return count, nil
+			return handleVirtioRead(readInto[:n], bufs, tun.bufPool, offset)
 		} else {
-			sizes[0] = n
+			bufs[0].Size = n
 			return 1, nil
 		}
 	}
diff --git a/tun/tun_openbsd.go b/tun/tun_openbsd.go
@@ -205,7 +205,7 @@ func (tun *NativeTun) Events() <-chan Event {
 	return tun.events
 }
 
-func (tun *NativeTun) Read(bufs []*buffer.Buffer, sizes []int, offset int) (int, error) {
+func (tun *NativeTun) Read(bufs []*buffer.Buffer, offset int) (int, error) {
 	select {
 	case err := <-tun.errors:
 		return 0, err
@@ -218,7 +218,7 @@ func (tun *NativeTun) Read(bufs []*buffer.Buffer, sizes []int, offset int) (int,
 		if n < 4 {
 			return 0, err
 		}
-		sizes[0] = n - 4
+		bufs[0].Size = n - 4
 		return 1, err
 	}
 }
diff --git a/tun/tun_plan9.go b/tun/tun_plan9.go
@@ -83,7 +83,7 @@ func (tun *NativeTun) Events() <-chan Event {
 	return tun.events
 }
 
-func (tun *NativeTun) Read(bufs []*buffer.Buffer, sizes []int, offset int) (int, error) {
+func (tun *NativeTun) Read(bufs []*buffer.Buffer, offset int) (int, error) {
 	select {
 	case err := <-tun.errors:
 		return 0, err
@@ -98,7 +98,7 @@ func (tun *NativeTun) Read(bufs []*buffer.Buffer, sizes []int, offset int) (int,
 			err = io.EOF
 			n = 0
 		}
-		sizes[0] = n
+		bufs[0].Size = n
 		return 1, err
 	}
 }
diff --git a/tun/tun_windows.go b/tun/tun_windows.go
@@ -145,7 +145,7 @@ func (tun *NativeTun) BatchSize() int {
 
 // Note: Read() and Write() assume the caller comes only from a single thread; there's no locking.
 
-func (tun *NativeTun) Read(bufs []*buffer.Buffer, sizes []int, offset int) (int, error) {
+func (tun *NativeTun) Read(bufs []*buffer.Buffer, offset int) (int, error) {
 	tun.running.Add(1)
 	defer tun.running.Done()
 retry:
@@ -166,7 +166,7 @@ retry:
 			}
 			packetSize := len(packet)
 			copy(bufs[0].Data()[offset:], packet)
-			sizes[0] = packetSize
+			bufs[0].Size = packetSize
 			tun.session.ReleaseReceivePacket(packet)
 			tun.rate.update(uint64(packetSize))
 			return 1, nil
diff --git a/tun/tuntest/tuntest.go b/tun/tuntest/tuntest.go
@@ -111,7 +111,7 @@ type chTun struct {
 
 func (t *chTun) File() *os.File { return nil }
 
-func (t *chTun) Read(bufs []*buffer.Buffer, sizes []int, offset int) (int, error) {
+func (t *chTun) Read(bufs []*buffer.Buffer, offset int) (int, error) {
 	select {
 	case <-t.c.closed:
 		return 0, os.ErrClosed
@@ -120,7 +120,7 @@ func (t *chTun) Read(bufs []*buffer.Buffer, sizes []int, offset int) (int, error
 			bufs[0] = buffer.New(make([]byte, buffer.MaxMessageSize), nil)
 		}
 		n := copy(bufs[0].Data()[offset:], msg)
-		sizes[0] = n
+		bufs[0].Size = n
 		return 1, nil
 	}
 }

Original file line number	Diff line number	Diff line change
`@@ -46,6 +46,8 @@ type Recycler interface {`
`46`	`46`	`type Buffer struct {`
`47`	`47`	`data []byte`
`48`	`48`	`recycler Recycler`
	`49`	`+`
	`50`	`+ Size int // size of the read, excluding offset`
`49`	`51`	`}`
`50`	`52`
`51`	`53`	`// New creates a standalone Buffer.`
Original file line number	Diff line number	Diff line change
`@@ -235,7 +235,6 @@ func (s *StdNetBind) receiveIP(`
`235`	`235`	`conn *net.UDPConn,`
`236`	`236`	`rxOffload bool,`
`237`	`237`	`bufs []*buffer.Buffer,`
`238`		`- sizes []int,`
`239`	`238`	`eps []Endpoint,`
`240`	`239`	`) (n int, err error) {`
`241`	`240`	`msgs := s.getMessages()`
`@@ -288,8 +287,8 @@ func (s *StdNetBind) receiveIP(`
`288`	`287`	`}`
`289`	`288`	`for i := 0; i < numMsgs; i++ {`
`290`	`289`	`msg := &(*msgs)[i]`
`291`		`- sizes[i] = msg.N`
`292`		`- if sizes[i] == 0 {`
	`290`	`+ bufs[i].Size = msg.N`
	`291`	`+ if bufs[i].Size == 0 {`
`293`	`292`	`continue`
`294`	`293`	`}`
`295`	`294`	`addrPort := msg.Addr.(*net.UDPAddr).AddrPort()`
`@@ -301,14 +300,14 @@ func (s *StdNetBind) receiveIP(`
`301`	`300`	`}`
`302`	`301`
`303`	`302`	`func (s StdNetBind) makeReceiveIPv4(pc ipv4.PacketConn, conn *net.UDPConn, rxOffload bool) ReceiveFunc {`
`304`		`- return func(bufs []*buffer.Buffer, sizes []int, eps []Endpoint) (n int, err error) {`
`305`		`- return s.receiveIP(pc, conn, rxOffload, bufs, sizes, eps)`
	`303`	`+ return func(bufs []*buffer.Buffer, eps []Endpoint) (n int, err error) {`
	`304`	`+ return s.receiveIP(pc, conn, rxOffload, bufs, eps)`
`306`	`305`	`}`
`307`	`306`	`}`
`308`	`307`
`309`	`308`	`func (s StdNetBind) makeReceiveIPv6(pc ipv6.PacketConn, conn *net.UDPConn, rxOffload bool) ReceiveFunc {`
`310`		`- return func(bufs []*buffer.Buffer, sizes []int, eps []Endpoint) (n int, err error) {`
`311`		`- return s.receiveIP(pc, conn, rxOffload, bufs, sizes, eps)`
	`309`	`+ return func(bufs []*buffer.Buffer, eps []Endpoint) (n int, err error) {`
	`310`	`+ return s.receiveIP(pc, conn, rxOffload, bufs, eps)`
`312`	`311`	`}`
`313`	`312`	`}`
`314`	`313`
Original file line number	Diff line number	Diff line change
`@@ -19,13 +19,12 @@ func TestStdNetBindReceiveFuncAfterClose(t *testing.T) {`
`19`	`19`	`bind.Close()`
`20`	`20`	`bufs := make([]*buffer.Buffer, IdealBatchSize)`
`21`	`21`	`bufs[0] = buffer.New(make([]byte, buffer.MaxMessageSize), nil)`
`22`		`- sizes := make([]int, IdealBatchSize)`
`23`	`22`	`eps := make([]Endpoint, IdealBatchSize)`
`24`	`23`	`for _, fn := range fns {`
`25`	`24`	`// The ReceiveFuncs must not access conn-related fields on StdNetBind`
`26`	`25`	`// unguarded. Close() nils the conn-related fields resulting in a panic`
`27`	`26`	`// if they violate the mutex.`
`28`		`- fn(bufs, sizes, eps)`
	`27`	`+ fn(bufs, eps)`
`29`	`28`	`}`
`30`	`29`	`}`
`31`	`30`
Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,7 @@ import (`
`13`	`13`
`14`	`14`	`func TestPrettyName(t *testing.T) {`
`15`	`15`	`var (`
`16`		`- recvFunc ReceiveFunc = func(bufs []*buffer.Buffer, sizes []int, eps []Endpoint) (n int, err error) { return }`
	`16`	`+ recvFunc ReceiveFunc = func(bufs []*buffer.Buffer, eps []Endpoint) (n int, err error) { return }`
`17`	`17`	`)`
`18`	`18`
`19`	`19`	`const want = "TestPrettyName"`
Original file line number	Diff line number	Diff line change
`@@ -438,7 +438,7 @@ type fakeTUNDeviceSized struct {`
`438`	`438`	`}`
`439`	`439`
`440`	`440`	`func (t fakeTUNDeviceSized) File() os.File { return nil }`
`441`		`-func (t fakeTUNDeviceSized) Read(bufs []buffer.Buffer, sizes []int, offset int) (n int, err error) {`
	`441`	`+func (t fakeTUNDeviceSized) Read(bufs []buffer.Buffer, offset int) (n int, err error) {`
`442`	`442`	`return 0, nil`
`443`	`443`	`}`
`444`	`444`	`func (t *fakeTUNDeviceSized) Write(bufs [][]byte, offset int) (int, error) { return 0, nil }`