rework tun-side segmenentation checksums

overlay/tio/tun_linux_offload_test.go

@@ -353,9 +353,13 @@ func buildUSOv4(t *testing.T, payLen, gsoSize int) ([]byte, virtio.Hdr) {
 	copy(pkt[12:16], []byte{10, 0, 0, 1})
 	copy(pkt[16:20], []byte{10, 0, 0, 2})
 
-	// UDP header (length + checksum filled in per segment by segmentUDPYield)
-	binary.BigEndian.PutUint16(pkt[20:22], 12345) // sport
-	binary.BigEndian.PutUint16(pkt[22:24], 53)    // dport
+	// UDP header. The kernel hands us a USO superpacket whose length field
+	// covers the WHOLE superpacket; the segmenter overwrites it per segment.
+	// Populating it here matters: leaving it zero makes the base-checksum path
+	// that must exclude it untestable, since excluding zero is a no-op.
+	binary.BigEndian.PutUint16(pkt[20:22], 12345)                 // sport
+	binary.BigEndian.PutUint16(pkt[22:24], 53)                    // dport
+	binary.BigEndian.PutUint16(pkt[24:26], uint16(udpLen+payLen)) // superpacket length
 
 	for i := 0; i < payLen; i++ {
 		pkt[ipLen+udpLen+i] = byte(i & 0xff)
@@ -465,6 +469,8 @@ func TestSegmentUDPv6(t *testing.T) {
 
 	binary.BigEndian.PutUint16(pkt[40:42], 12345)
 	binary.BigEndian.PutUint16(pkt[42:44], 53)
+	// Superpacket-wide length, as the kernel supplies it; see buildUSOv4.
+	binary.BigEndian.PutUint16(pkt[44:46], uint16(udpLen+payLen))
 
 	for i := 0; i < payLen; i++ {
 		pkt[ipLen+udpLen+i] = byte(i)
@@ -669,6 +675,75 @@ func TestTunFileWriteVnetHdrNoAlloc(t *testing.T) {
 	}
 }
 
+// TestSegmentSuperpacketNoAlloc pins the segmenters' zero-allocation
+// contract. Both SegmentTCP and SegmentUDP derive their per-superpacket
+// constants into fixed-size arrays (tmp/ipTmp/savedHdr) that must stay on
+// the stack, and both take a yield closure that must not escape. Any of
+// those escaping turns one allocation into one-per-superpacket on the
+// hottest path in the reader, which BenchmarkSegmentSuperpacketAllocsTSO
+// reports but nothing fails on. This does.
+//
+// The yield closure here only touches captured scalars: appending segments
+// to a slice would allocate in the test itself and mask the measurement.
+func TestSegmentSuperpacketNoAlloc(t *testing.T) {
+	const mss = 1400
+	const numSeg = 8
+
+	cases := []struct {
+		name  string
+		build func() ([]byte, virtio.Hdr)
+	}{
+		{"tso-v4", func() ([]byte, virtio.Hdr) { return buildTSOv4(t, mss*numSeg, mss) }},
+		{"uso-v4", func() ([]byte, virtio.Hdr) { return buildUSOv4(t, mss*numSeg, mss) }},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			master, hdr := tc.build()
+			proto, err := protoFromGSOType(hdr.GSOType())
+			if err != nil {
+				t.Fatalf("protoFromGSOType: %v", err)
+			}
+			work := make([]byte, len(master))
+			p := Packet{Bytes: work, GSO: GSOInfo{
+				Size:      hdr.GSOSize,
+				HdrLen:    hdr.HdrLen,
+				CsumStart: hdr.CsumStart,
+				Proto:     proto,
+			}}
+
+			// Segmentation consumes its input destructively, so restore from
+			// the master copy each run; copy(2) into an existing slice does
+			// not allocate. seen/bytes keep the closure from being optimized
+			// away and double as a sanity check that work actually happened.
+			var seen, bytes int
+			run := func() {
+				copy(work, master)
+				seen, bytes = 0, 0
+				if err := SegmentSuperpacket(p, func(seg []byte) error {
+					seen++
+					bytes += len(seg)
+					return nil
+				}); err != nil {
+					t.Fatalf("SegmentSuperpacket: %v", err)
+				}
+			}
+
+			run() // warm up: absorb any one-time allocation elsewhere
+			if seen != numSeg {
+				t.Fatalf("yielded %d segments, want %d", seen, numSeg)
+			}
+
+			if allocs := testing.AllocsPerRun(200, run); allocs != 0 {
+				t.Fatalf("SegmentSuperpacket allocated %.1f times per call, want 0", allocs)
+			}
+			if seen != numSeg || bytes == 0 {
+				t.Fatalf("post-measure sanity: seen=%d bytes=%d", seen, bytes)
+			}
+		})
+	}
+}
+
 // TestWriteGSOSkipsEmptyPayloads is the defense-in-depth guard for the
 // zero-length UDP DoS: a payload fragment of length zero would make &p[0]
 // panic (index-out-of-range) when building the iovec array. WriteGSO must
@@ -958,3 +1033,108 @@ func TestWriteGSORejectsBadGeometry(t *testing.T) {
 		})
 	}
 }
+
+// BenchmarkSegmentUDPv4 is the USO counterpart to BenchmarkSegmentTCPv4. The
+// yield is a no-op so the measurement is segmentation plus checksum work only.
+func BenchmarkSegmentUDPv4(b *testing.B) {
+	sizes := []struct {
+		name    string
+		payLen  int
+		gsoSize int
+	}{
+		{"64KiB_GSO1400", 64000, 1400},
+		{"16KiB_GSO1400", 16384, 1400},
+		{"4KiB_GSO1400", 4096, 1400},
+	}
+	for _, sz := range sizes {
+		b.Run(sz.name, func(b *testing.B) {
+			const ipLen = 20
+			const udpLen = 8
+			pkt := make([]byte, ipLen+udpLen+sz.payLen)
+			pkt[0] = 0x45
+			binary.BigEndian.PutUint16(pkt[2:4], uint16(ipLen+udpLen+sz.payLen))
+			binary.BigEndian.PutUint16(pkt[4:6], 0x4242)
+			pkt[8] = 64
+			pkt[9] = unix.IPPROTO_UDP
+			copy(pkt[12:16], []byte{10, 0, 0, 1})
+			copy(pkt[16:20], []byte{10, 0, 0, 2})
+			binary.BigEndian.PutUint16(pkt[20:22], 12345)
+			binary.BigEndian.PutUint16(pkt[22:24], 53)
+			binary.BigEndian.PutUint16(pkt[24:26], uint16(udpLen+sz.payLen))
+			for i := 0; i < sz.payLen; i++ {
+				pkt[ipLen+udpLen+i] = byte(i)
+			}
+
+			master := append([]byte(nil), pkt...)
+			work := make([]byte, len(pkt))
+			p := Packet{Bytes: work, GSO: GSOInfo{
+				Size:      uint16(sz.gsoSize),
+				HdrLen:    ipLen + udpLen,
+				CsumStart: ipLen,
+				Proto:     GSOProtoUDP,
+			}}
+
+			b.SetBytes(int64(len(pkt)))
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				copy(work, master)
+				if err := SegmentSuperpacket(p, func(seg []byte) error { return nil }); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}
+
+// BenchmarkSegmentUDPv6 mirrors BenchmarkSegmentUDPv4 for IPv6, where the
+// pseudo-header address sum is 32 bytes rather than 8.
+func BenchmarkSegmentUDPv6(b *testing.B) {
+	sizes := []struct {
+		name    string
+		payLen  int
+		gsoSize int
+	}{
+		{"64KiB_GSO1400", 64000, 1400},
+		{"16KiB_GSO1400", 16384, 1400},
+		{"4KiB_GSO1400", 4096, 1400},
+	}
+	for _, sz := range sizes {
+		b.Run(sz.name, func(b *testing.B) {
+			const ipLen = 40
+			const udpLen = 8
+			pkt := make([]byte, ipLen+udpLen+sz.payLen)
+			pkt[0] = 0x60
+			binary.BigEndian.PutUint16(pkt[4:6], uint16(udpLen+sz.payLen))
+			pkt[6] = unix.IPPROTO_UDP
+			pkt[7] = 64
+			pkt[8], pkt[9], pkt[23] = 0xfe, 0x80, 1
+			pkt[24], pkt[25], pkt[39] = 0xfe, 0x80, 2
+			binary.BigEndian.PutUint16(pkt[40:42], 12345)
+			binary.BigEndian.PutUint16(pkt[42:44], 53)
+			binary.BigEndian.PutUint16(pkt[44:46], uint16(udpLen+sz.payLen))
+			for i := 0; i < sz.payLen; i++ {
+				pkt[ipLen+udpLen+i] = byte(i)
+			}
+
+			master := append([]byte(nil), pkt...)
+			work := make([]byte, len(pkt))
+			p := Packet{Bytes: work, GSO: GSOInfo{
+				Size:      uint16(sz.gsoSize),
+				HdrLen:    ipLen + udpLen,
+				CsumStart: ipLen,
+				Proto:     GSOProtoUDP,
+			}}
+
+			b.SetBytes(int64(len(pkt)))
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				copy(work, master)
+				if err := SegmentSuperpacket(p, func(seg []byte) error { return nil }); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}

overlay/tio/virtio/segment_linux.go

@@ -27,11 +27,8 @@ const (
 	tcpHeaderMaxLen  = 60 // data-offset=15, max options
 )
 
-// maxSegHdrLen bounds the L3+L4 header we snapshot before stamping each
-// segment. The largest header the segmenter supports is IPv4 (max IHL 60)
-// plus TCP (max data-offset 60) = 120 bytes; the array is sized to that
-// worst case so the snapshot lives on the stack with no per-call heap
-// allocation.
+// maxSegHdrLen bounds the L3+L4 header we snapshot before stamping each segment.
+// The largest header the segmenter supports is IPv4 (max IHL 60) plus TCP (max data-offset 60) = 120 bytes
 const maxSegHdrLen = ipv4HeaderMaxLen + tcpHeaderMaxLen // 120
 
 // Byte offsets inside an IPv4 header.
@@ -68,9 +65,9 @@ const (
 // tcpFinPshMask is cleared on every segment except the last of a TSO burst.
 const tcpFinPshMask = 0x09 // FIN(0x01) | PSH(0x08)
 
-// tcpCwrFlag is cleared on every segment except the first. Per RFC 3168
-// §6.1.2 the CWR bit signals a one-shot transition (the sender just halved
-// its window) and must appear on the first segment of a TSO burst only.
+// tcpCwrFlag is cleared on every segment except the first.
+// Per RFC 3168 §6.1.2 the CWR bit signals a one-shot transition (the sender just halved its window)
+// and must appear on the first segment of a TSO burst only.
 const tcpCwrFlag = 0x80
 
 // CheckValid rejects packets whose virtio_net_hdr/IP combination would
@@ -112,15 +109,13 @@ func CheckValid(pkt []byte, hdr Hdr) error {
 	return nil
 }
 
-// CorrectHdrLen rewrites hdr.HdrLen based on the actual transport header
-// length read out of pkt. The kernel's hdr.HdrLen on the FORWARD path can
-// be the length of the entire first packet, so we don't trust it.
+// CorrectHdrLen rewrites hdr.HdrLen based on the actual transport header length read out of pkt.
+// The kernel's hdr.HdrLen on the FORWARD path can be the length of the entire first packet, so we don't trust it.
 func CorrectHdrLen(pkt []byte, hdr *Hdr) error {
 	// Thank you wireguard-go for documenting these edge-cases
 	// Don't trust hdr.hdrLen from the kernel as it can be equal to the length
-	// of the entire first packet when the kernel is handling it as part of a
-	// FORWARD path. Instead, parse the transport header length and add it onto
-	// csumStart, which is synonymous for IP header length.
+	// of the entire first packet when the kernel is handling it as part of a FORWARD path.
+	// Instead, parse the transport header length and add it onto csumStart, which is synonymous for IP header length.
 
 	if hdr.GSOType() == unix.VIRTIO_NET_HDR_GSO_UDP_L4 {
 		hdr.HdrLen = hdr.CsumStart + 8
@@ -130,8 +125,7 @@ func CorrectHdrLen(pkt []byte, hdr *Hdr) error {
 		}
 
 		tcpHLen := uint16(pkt[hdr.CsumStart+tcpDataOffOff] >> 4 * 4)
-		if tcpHLen < 20 || tcpHLen > 60 {
-			// A TCP header must be between 20 and 60 bytes in length.
+		if tcpHLen < tcpHeaderMinLen || tcpHLen > tcpHeaderMaxLen {
 			return fmt.Errorf("tcp header len is invalid: %d", tcpHLen)
 		}
 		hdr.HdrLen = hdr.CsumStart + tcpHLen
@@ -151,14 +145,63 @@ func CorrectHdrLen(pkt []byte, hdr *Hdr) error {
 	return nil
 }
 
+// segCount returns how many segments a payload of payLen bytes splits into at gsoSize,
+// with a floor of one so a header-only superpacket still yields a single segment.
+func segCount(payLen, gsoSize int) int {
+	n := (payLen + gsoSize - 1) / gsoSize
+	if n == 0 {
+		return 1
+	}
+	return n
+}
+
+// basePseudoSum folds the part of the L4 pseudo-header sum that is identical
+// for every segment: the source and destination addresses plus the protocol
+// number. The per-segment L4 length is added by the caller inside the loop.
+func basePseudoSum(pkt []byte, isV4 bool, proto uint32) uint32 {
+	if isV4 {
+		return uint32(checksum.Checksum(pkt[ipv4SrcOff:ipv4AddrsEnd], 0)) + proto
+	}
+	return uint32(checksum.Checksum(pkt[ipv6SrcOff:ipv6AddrsEnd], 0)) + proto
+}
+
+// baseIPv4HdrSum folds the IPv4 header checksum over the fields that stay constant across segments.
+// csumStart is the L3 header length, which bounds a valid IHL.
+func baseIPv4HdrSum(pkt []byte, csumStart int, zeroID bool) (uint32, error) {
+	ihl := int(pkt[0]&0x0f) * 4
+	if ihl < ipv4HeaderMinLen || ihl > csumStart {
+		return 0, fmt.Errorf("bad IPv4 IHL: %d", ihl)
+	}
+	// total_len and the checksum field itself are always excluded, since both are rewritten per segment.
+	sum := uint32(checksum.Checksum(pkt[:ihl], 0))
+	sum += uint32(^binary.BigEndian.Uint16(pkt[ipv4TotalLenOff : ipv4TotalLenOff+2]))
+	sum += uint32(^binary.BigEndian.Uint16(pkt[ipv4ChecksumOff : ipv4ChecksumOff+2]))
+	if zeroID { //only zero the ID if requested
+		sum += uint32(^binary.BigEndian.Uint16(pkt[ipv4IDOff : ipv4IDOff+2]))
+	}
+	sum = (sum & 0xffff) + (sum >> 16)
+	sum = (sum & 0xffff) + (sum >> 16)
+	return sum, nil
+}
+
+// baseTCPHdrSum folds the TCP header checksum over everything the segment loop does not rewrite
+func baseTCPHdrSum(pkt []byte, csumStart, headerLen int) uint32 {
+	seq := binary.BigEndian.Uint32(pkt[csumStart+tcpSeqOff : csumStart+tcpSeqOff+4])
+	flags := uint16(pkt[csumStart+tcpFlagsOff])
+
+	sum := uint32(checksum.Checksum(pkt[csumStart:headerLen], 0))
+	sum += uint32(^uint16(seq >> 16))
+	sum += uint32(^uint16(seq))
+	sum += uint32(^flags)
+	sum += uint32(^binary.BigEndian.Uint16(pkt[csumStart+tcpChecksumOff : csumStart+tcpChecksumOff+2]))
+	sum = (sum & 0xffff) + (sum >> 16)
+	sum = (sum & 0xffff) + (sum >> 16)
+	return sum
+}
+
 // SegmentTCP walks a TSO superpacket pkt, yielding each segment as a slice into pkt.
 // Per-segment plaintext is laid out by stamping a copy of the original L3+L4 header into pkt at offset i*gsoSize,
 // where it sits immediately before that segment's payload chunk in the original buffer.
-// The stamp is destructive: iter i's header write lands on pkt[i*G : i*G+hdrLen], which is the tail of seg_{i-1}'s payload (already
-// consumed) and ends exactly where seg_i's payload begins, so it never clobbers
-// live payload — this holds even when gsoSize < hdrLen.
-// The header bytes are sourced from a pristine snapshot taken before the loop (savedHdr), NOT from pkt[:hdrLen], because when gsoSize < hdrLen the stamps would otherwise
-// overwrite the leading header in place and every stamp after the first would copy corrupted bytes.
 // pkt is consumed by this call and must not be inspected by the caller after the final yield.
 func SegmentTCP(pkt []byte, hdrLenU, csumStartU, gsoSizeU uint16, yield func(seg []byte) error) error {
 	if gsoSizeU == 0 {
@@ -178,49 +221,28 @@ func SegmentTCP(pkt []byte, hdrLenU, csumStartU, gsoSizeU uint16, yield func(seg
 	tcpHdrLen := int(pkt[csumStart+tcpDataOffOff]>>4) * 4
 	payLen := len(pkt) - headerLen
 	gsoSize := int(gsoSizeU)
-	numSeg := (payLen + gsoSize - 1) / gsoSize
-	if numSeg == 0 {
-		numSeg = 1
-	}
+	numSeg := segCount(payLen, gsoSize)
 
 	origSeq := binary.BigEndian.Uint32(pkt[csumStart+tcpSeqOff : csumStart+tcpSeqOff+4])
 	origFlags := pkt[csumStart+tcpFlagsOff]
 
-	var tmp [tcpHeaderMaxLen]byte
-	copy(tmp[:tcpHdrLen], pkt[csumStart:headerLen])
-	tmp[tcpSeqOff], tmp[tcpSeqOff+1], tmp[tcpSeqOff+2], tmp[tcpSeqOff+3] = 0, 0, 0, 0
-	tmp[tcpFlagsOff] = 0
-	tmp[tcpChecksumOff], tmp[tcpChecksumOff+1] = 0, 0
-	baseTcpHdrSum := uint32(checksum.Checksum(tmp[:tcpHdrLen], 0))
-
-	var baseProtoSum uint32
-	if isV4 {
-		baseProtoSum = uint32(checksum.Checksum(pkt[ipv4SrcOff:ipv4AddrsEnd], 0))
-	} else {
-		baseProtoSum = uint32(checksum.Checksum(pkt[ipv6SrcOff:ipv6AddrsEnd], 0))
-	}
-	baseProtoSum += uint32(unix.IPPROTO_TCP)
+	baseProtoSum := basePseudoSum(pkt, isV4, unix.IPPROTO_TCP)
+	baseTcpHdrSum := baseTCPHdrSum(pkt, csumStart, headerLen)
 
 	var origIPID uint16
 	var baseIPHdrSum uint32
 	if isV4 {
 		origIPID = binary.BigEndian.Uint16(pkt[ipv4IDOff : ipv4IDOff+2])
-		ihl := int(pkt[0]&0x0f) * 4
-		if ihl < ipv4HeaderMinLen || ihl > csumStart {
-			return fmt.Errorf("bad IPv4 IHL: %d", ihl)
+		var err error
+		// TSO bumps the ID per segment, so it stays out of the base sum.
+		baseIPHdrSum, err = baseIPv4HdrSum(pkt, csumStart, true)
+		if err != nil {
+			return err
 		}
-		var ipTmp [ipv4HeaderMaxLen]byte
-		copy(ipTmp[:ihl], pkt[:ihl])
-		ipTmp[ipv4TotalLenOff], ipTmp[ipv4TotalLenOff+1] = 0, 0
-		ipTmp[ipv4IDOff], ipTmp[ipv4IDOff+1] = 0, 0
-		ipTmp[ipv4ChecksumOff], ipTmp[ipv4ChecksumOff+1] = 0, 0
-		baseIPHdrSum = uint32(checksum.Checksum(ipTmp[:ihl], 0))
 	}
 
-	// Snapshot the pristine L3+L4 header once. Every segment's header is
-	// stamped from this copy, so overlapping stamps (gsoSize < headerLen)
-	// can never corrupt the source. The variable fields (seq/flags/cksum/
-	// totalLen/id) captured here are stale but are overwritten per segment.
+	// Snapshot the pristine L3+L4 header once. '
+	// Every segment's header is stamped from this copy, so overlapping stamps (gsoSize < headerLen) can never corrupt the source.
 	var savedHdr [maxSegHdrLen]byte
 	copy(savedHdr[:headerLen], pkt[:headerLen])
 
@@ -234,12 +256,10 @@ func SegmentTCP(pkt []byte, hdrLenU, csumStartU, gsoSizeU uint16, yield func(seg
 		segLen := headerLen + segPayLen
 		headerOff := i * gsoSize
 
-		// Stamp the header into place immediately before this segment's
-		// payload, sourced from the pristine snapshot. Iter 0's header is
-		// already at pkt[:headerLen] (identical to savedHdr), so only i ≥ 1
-		// needs the stamp. The per-segment patches below overwrite the
-		// variable fields.
+		// Stamp the header into place immediately before this segment's payload, sourced from the snapshot.
+		// The per-segment patches below overwrite the variable fields. (seq/flags/cksum/totalLen/id)
 		if i > 0 {
+			// Iter 0's header is already at pkt[:headerLen] (identical to savedHdr), so only i >= 1 needs the stamp
 			copy(pkt[headerOff:headerOff+headerLen], savedHdr[:headerLen])
 		}
 		seg := pkt[headerOff : headerOff+segLen]
@@ -268,10 +288,9 @@ func SegmentTCP(pkt []byte, hdrLenU, csumStartU, gsoSizeU uint16, yield func(seg
 		seg[csumStart+tcpFlagsOff] = segFlags
 
 		tcpLen := tcpHdrLen + segPayLen
-		// Payload bytes still live at their original offset in pkt. The
-		// header slide above only writes into pkt[i*G : i*G+H], which is
-		// the tail of seg_{i-1}'s payload (already consumed) and never
-		// overlaps seg_i's own payload at pkt[H+i*G : H+(i+1)*G].
+		// Payload bytes still live at their original offset in pkt.
+		// The header slide above only writes into pkt[i*GSOSize : i*GSOSize+header], which is the tail of seg_{i-1}'s payload (already consumed)
+		// and never overlaps seg_i's own payload at pkt[header+i*GSOSize : header+(i+1)*GSOSize].
 		paySum := uint32(checksum.Checksum(pkt[headerLen+segStart:headerLen+segEnd], 0))
 		wide := uint64(baseTcpHdrSum) + uint64(paySum) + uint64(baseProtoSum)
 		wide += uint64(segSeq) + uint64(segFlags) + uint64(tcpLen)
@@ -287,17 +306,10 @@ func SegmentTCP(pkt []byte, hdrLenU, csumStartU, gsoSizeU uint16, yield func(seg
 	return nil
 }
 
-// SegmentUDP walks a USO superpacket, stamping a per-segment-patched copy of
-// the original L3+L4 header into pkt at offset i*gsoSize and yielding
-// pkt[i*G:i*G+segLen] to the caller. Per-segment patches are total_len +
-// IPv4 csum (or IPv6 payload_len) plus the UDP length and checksum. pkt is
-// consumed destructively; see SegmentTCP for the layout reasoning, including
-// why the header is stamped from a pristine snapshot rather than pkt[:hdrLen]
-// (correctness when gsoSize < hdrLen).
-//
-// UDP-GSO leaves the IPv4 ID identical across segments (the kernel does not
-// bump it), which is why the IP-level per-segment work is limited to
-// total_len + IPv4 header checksum (v4) or payload_len (v6).
+// SegmentUDP walks a USO superpacket, stamping a per-segment-patched copy of the original L3+L4 header
+// into pkt at offset i*GSOSize and yielding pkt[i*GSOSize:i*GSOSize+segLen] to the caller.
+// Per-segment patches are total_len + IPv4 csum (or IPv6 payload_len) plus the UDP length and checksum.
+// pkt is consumed destructively.
 func SegmentUDP(pkt []byte, hdrLenU, csumStartU, gsoSizeU uint16, yield func(seg []byte) error) error {
 	if gsoSizeU == 0 {
 		return fmt.Errorf("gso_size is zero")
@@ -318,41 +330,21 @@ func SegmentUDP(pkt []byte, hdrLenU, csumStartU, gsoSizeU uint16, yield func(seg
 
 	payLen := len(pkt) - headerLen
 	gsoSize := int(gsoSizeU)
-	numSeg := (payLen + gsoSize - 1) / gsoSize
-	if numSeg == 0 {
-		numSeg = 1
-	}
+	numSeg := segCount(payLen, gsoSize)
 
-	var udpTmp [udpHeaderLen]byte
-	copy(udpTmp[:], pkt[csumStart:headerLen])
-	udpTmp[udpLengthOff], udpTmp[udpLengthOff+1] = 0, 0
-	udpTmp[udpChecksumOff], udpTmp[udpChecksumOff+1] = 0, 0
-	baseUDPHdrSum := uint32(checksum.Checksum(udpTmp[:], 0))
-
-	var baseProtoSum uint32
-	if isV4 {
-		baseProtoSum = uint32(checksum.Checksum(pkt[ipv4SrcOff:ipv4AddrsEnd], 0))
-	} else {
-		baseProtoSum = uint32(checksum.Checksum(pkt[ipv6SrcOff:ipv6AddrsEnd], 0))
-	}
-	baseProtoSum += uint32(unix.IPPROTO_UDP)
+	baseProtoSum := basePseudoSum(pkt, isV4, unix.IPPROTO_UDP)
 
 	var baseIPHdrSum uint32
 	if isV4 {
-		ihl := int(pkt[0]&0x0f) * 4
-		if ihl < ipv4HeaderMinLen || ihl > csumStart {
-			return fmt.Errorf("bad IPv4 IHL: %d", ihl)
+		var err error
+		// UDP GSO holds the ID constant across the burst, so it stays in the base sum.
+		baseIPHdrSum, err = baseIPv4HdrSum(pkt, csumStart, false)
+		if err != nil {
+			return err
 		}
-		var ipTmp [ipv4HeaderMaxLen]byte
-		copy(ipTmp[:ihl], pkt[:ihl])
-		ipTmp[ipv4TotalLenOff], ipTmp[ipv4TotalLenOff+1] = 0, 0
-		ipTmp[ipv4ChecksumOff], ipTmp[ipv4ChecksumOff+1] = 0, 0
-		baseIPHdrSum = uint32(checksum.Checksum(ipTmp[:ihl], 0))
 	}
 
-	// Snapshot the pristine L3+L4 header once and stamp every segment from
-	// it; see SegmentTCP for why sourcing from pkt[:headerLen] corrupts
-	// segments when gsoSize < headerLen.
+	// Snapshot the pristine L3+L4 header once and stamp every segment from it
 	var savedHdr [maxSegHdrLen]byte
 	copy(savedHdr[:headerLen], pkt[:headerLen])
 
@@ -384,12 +376,13 @@ func SegmentUDP(pkt []byte, hdrLenU, csumStartU, gsoSizeU uint16, yield func(seg
 
 		binary.BigEndian.PutUint16(seg[csumStart+udpLengthOff:csumStart+udpLengthOff+2], uint16(udpLen))
 
-		paySum := uint32(checksum.Checksum(pkt[headerLen+segStart:headerLen+segEnd], 0))
-		wide := uint64(baseUDPHdrSum) + uint64(paySum) + uint64(baseProtoSum)
-		wide += uint64(udpLen) + uint64(udpLen)
-		wide = (wide & 0xffffffff) + (wide >> 32)
-		wide = (wide & 0xffffffff) + (wide >> 32)
-		csum := foldComplement(uint32(wide))
+		// Sum the UDP header (length just written, checksum zeroed) together with
+		// this segment's payload in one pass, seeded with the pseudo-header sum.
+		seg[csumStart+udpChecksumOff], seg[csumStart+udpChecksumOff+1] = 0, 0
+		pseudo := baseProtoSum + uint32(udpLen)
+		pseudo = (pseudo & 0xffff) + (pseudo >> 16)
+		pseudo = (pseudo & 0xffff) + (pseudo >> 16)
+		csum := ^checksum.Checksum(seg[csumStart:], uint16(pseudo))
 		if csum == 0 {
 			csum = 0xffff
 		}
@@ -403,10 +396,9 @@ func SegmentUDP(pkt []byte, hdrLenU, csumStartU, gsoSizeU uint16, yield func(seg
 	return nil
 }
 
-// FinishChecksum computes the L4 checksum for a non-GSO packet that the kernel
-// handed us with NEEDS_CSUM set. csum_start / csum_offset point at the 16-bit
-// checksum field; we zero it, fold a full sum (the field was pre-loaded with
-// the pseudo-header partial sum by the kernel), and store the result.
+// FinishChecksum computes the L4 checksum for a non-GSO packet that the kernel handed us with NEEDS_CSUM set.
+// CsumStart / CsumOffset point at the 16-bit checksum field.
+// We zero it, fold a full sum from the partial one that the kernel provided, and store the result.
 func FinishChecksum(seg []byte, hdr Hdr) error {
 	cs := int(hdr.CsumStart)
 	co := int(hdr.CsumOffset)

overlay/tio/virtio/segment_linux_test.go

@@ -494,3 +494,102 @@ func TestFoldComplementMatchesReference(t *testing.T) {
 		}
 	}
 }
+
+// referenceBaseIPv4HdrSum and referenceBaseTCPHdrSum are the straightforward
+// implementations that baseIPv4HdrSum/baseTCPHdrSum replaced: copy the header
+// into scratch, zero the fields the segment loop rewrites, sum. The production
+// versions instead sum in place and subtract those fields via one's-complement
+// arithmetic, which is faster but far less obvious — particularly for the TCP
+// flags byte, which is only half of a 16-bit word. These references exist so
+// that trade is checked rather than asserted.
+func referenceBaseIPv4HdrSum(pkt []byte, ihl int, zeroID bool) uint32 {
+	var ipTmp [ipv4HeaderMaxLen]byte
+	copy(ipTmp[:ihl], pkt[:ihl])
+	ipTmp[ipv4TotalLenOff], ipTmp[ipv4TotalLenOff+1] = 0, 0
+	ipTmp[ipv4ChecksumOff], ipTmp[ipv4ChecksumOff+1] = 0, 0
+	if zeroID {
+		ipTmp[ipv4IDOff], ipTmp[ipv4IDOff+1] = 0, 0
+	}
+	return uint32(checksum.Checksum(ipTmp[:ihl], 0))
+}
+
+func referenceBaseTCPHdrSum(pkt []byte, csumStart, headerLen int) uint32 {
+	tcpLen := headerLen - csumStart
+	var tmp [tcpHeaderMaxLen]byte
+	copy(tmp[:tcpLen], pkt[csumStart:headerLen])
+	tmp[tcpSeqOff], tmp[tcpSeqOff+1], tmp[tcpSeqOff+2], tmp[tcpSeqOff+3] = 0, 0, 0, 0
+	tmp[tcpFlagsOff] = 0
+	tmp[tcpChecksumOff], tmp[tcpChecksumOff+1] = 0, 0
+	return uint32(checksum.Checksum(tmp[:tcpLen], 0))
+}
+
+// randSeed is a tiny deterministic PRNG so this test needs no imports beyond
+// what the file already has and reproduces identically on every run.
+func randByte(state *uint32) byte {
+	*state = *state*1664525 + 1013904223
+	return byte(*state >> 24)
+}
+
+func TestBaseSumsMatchZeroingReference(t *testing.T) {
+	state := uint32(12345)
+
+	t.Run("ipv4", func(t *testing.T) {
+		for ihl := ipv4HeaderMinLen; ihl <= ipv4HeaderMaxLen; ihl += 4 {
+			for _, zeroID := range []bool{true, false} {
+				for iter := 0; iter < 5000; iter++ {
+					pkt := make([]byte, ihl)
+					for i := range pkt {
+						pkt[i] = randByte(&state)
+					}
+					pkt[0] = byte(0x40 | (ihl / 4))
+
+					want := referenceBaseIPv4HdrSum(pkt, ihl, zeroID)
+					got, err := baseIPv4HdrSum(pkt, ihl, zeroID)
+					if err != nil {
+						t.Fatalf("ihl=%d: %v", ihl, err)
+					}
+					// Compare the value that reaches the wire: the raw partial
+					// sums may legally differ by one's-complement -0 vs +0.
+					for _, tl := range []uint32{20, 1500, 65535} {
+						for _, id := range []uint32{0, 0x4242, 0xffff} {
+							if a, b := foldComplement(want+tl+id), foldComplement(got+tl+id); a != b {
+								t.Fatalf("ihl=%d zeroID=%v tl=%d id=%d: %#04x != %#04x", ihl, zeroID, tl, id, a, b)
+							}
+						}
+					}
+				}
+			}
+		}
+	})
+
+	t.Run("tcp", func(t *testing.T) {
+		const csumStart = 20
+		for dataOff := 5; dataOff <= 15; dataOff++ {
+			tcpLen := dataOff * 4
+			headerLen := csumStart + tcpLen
+			for iter := 0; iter < 5000; iter++ {
+				pkt := make([]byte, headerLen+64)
+				for i := range pkt {
+					pkt[i] = randByte(&state)
+				}
+				pkt[0] = 0x45
+				pkt[csumStart+tcpDataOffOff] = byte(dataOff << 4)
+
+				want := referenceBaseTCPHdrSum(pkt, csumStart, headerLen)
+				got := baseTCPHdrSum(pkt, csumStart, headerLen)
+				for _, seq := range []uint32{0, 1, 0x4242_4242, 0xffff_ffff} {
+					for _, fl := range []uint32{0x00, 0x10, 0x18, 0x19, 0xff} {
+						for _, l4 := range []uint32{20, 1460, 65535} {
+							a := foldComplement(want + seq + fl + l4)
+							b := foldComplement(got + seq + fl + l4)
+							if a != b {
+								t.Fatalf("dataOff=%d seq=%#x fl=%#x l4=%d: %#04x != %#04x",
+									dataOff, seq, fl, l4, a, b)
+							}
+						}
+					}
+				}
+			}
+		}
+	})
+}