Incrementally update checksum when NAT-ing

...instead of calculating a fresh checksum to avoid re-calcalculating
a checksum on unchanged bytes.

Fixes #5340.

PiperOrigin-RevId: 381403888

1 files changed, 62 insertions, 0 deletions

diff --git a/pkg/tcpip/header/checksum.go b/pkg/tcpip/header/checksum.go
index 6aa9acfa8..e2c85e220 100644
--- a/pkg/tcpip/header/checksum.go
+++ b/pkg/tcpip/header/checksum.go
@@ -18,6 +18,7 @@ package header
 
 import (
 	"encoding/binary"
+	"fmt"
 
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/buffer"
@@ -234,3 +235,64 @@ func PseudoHeaderChecksum(protocol tcpip.TransportProtocolNumber, srcAddr tcpip.
 
 	return Checksum([]byte{0, uint8(protocol)}, xsum)
 }
+
+// checksumUpdate2ByteAlignedUint16 updates a uint16 value in a calculated
+// checksum.
+//
+// The value MUST begin at a 2-byte boundary in the original buffer.
+func checksumUpdate2ByteAlignedUint16(xsum, old, new uint16) uint16 {
+	// As per RFC 1071 page 4,
+	//	(4)  Incremental Update
+	//
+	//        ...
+	//
+	//        To update the checksum, simply add the differences of the
+	//        sixteen bit integers that have been changed.  To see why this
+	//        works, observe that every 16-bit integer has an additive inverse
+	//        and that addition is associative.  From this it follows that
+	//        given the original value m, the new value m', and the old
+	//        checksum C, the new checksum C' is:
+	//
+	//                C' = C + (-m) + m' = C + (m' - m)
+	return ChecksumCombine(xsum, ChecksumCombine(new, ^old))
+}
+
+// checksumUpdate2ByteAlignedAddress updates an address in a calculated
+// checksum.
+//
+// The addresses must have the same length and must contain an even number
+// of bytes. The address MUST begin at a 2-byte boundary in the original buffer.
+func checksumUpdate2ByteAlignedAddress(xsum uint16, old, new tcpip.Address) uint16 {
+	const uint16Bytes = 2
+
+	if len(old) != len(new) {
+		panic(fmt.Sprintf("buffer lengths are different; old = %d, new = %d", len(old), len(new)))
+	}
+
+	if len(old)%uint16Bytes != 0 {
+		panic(fmt.Sprintf("buffer has an odd number of bytes; got = %d", len(old)))
+	}
+
+	// As per RFC 1071 page 4,
+	//	(4)  Incremental Update
+	//
+	//        ...
+	//
+	//        To update the checksum, simply add the differences of the
+	//        sixteen bit integers that have been changed.  To see why this
+	//        works, observe that every 16-bit integer has an additive inverse
+	//        and that addition is associative.  From this it follows that
+	//        given the original value m, the new value m', and the old
+	//        checksum C, the new checksum C' is:
+	//
+	//                C' = C + (-m) + m' = C + (m' - m)
+	for len(old) != 0 {
+		// Convert the 2 byte sequences to uint16 values then apply the increment
+		// update.
+		xsum = checksumUpdate2ByteAlignedUint16(xsum, (uint16(old[0])<<8)+uint16(old[1]), (uint16(new[0])<<8)+uint16(new[1]))
+		old = old[uint16Bytes:]
+		new = new[uint16Bytes:]
+	}
+
+	return xsum
+}