5 files changed, 170 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
+}
diff --git a/pkg/tcpip/header/interfaces.go b/pkg/tcpip/header/interfaces.go
index 861cbbb70..3a41adfc4 100644
--- a/pkg/tcpip/header/interfaces.go
+++ b/pkg/tcpip/header/interfaces.go
@@ -53,6 +53,31 @@ type Transport interface {
 	Payload() []byte
 }
 
+// ChecksummableTransport is a Transport that supports checksumming.
+type ChecksummableTransport interface {
+	Transport
+
+	// SetSourcePortWithChecksumUpdate sets the source port and updates
+	// the checksum.
+	//
+	// The receiver's checksum must be a fully calculated checksum.
+	SetSourcePortWithChecksumUpdate(port uint16)
+
+	// SetDestinationPortWithChecksumUpdate sets the destination port and updates
+	// the checksum.
+	//
+	// The receiver's checksum must be a fully calculated checksum.
+	SetDestinationPortWithChecksumUpdate(port uint16)
+
+	// UpdateChecksumPseudoHeaderAddress updates the checksum to reflect an
+	// updated address in the pseudo header.
+	//
+	// If fullChecksum is true, the receiver's checksum field is assumed to hold a
+	// fully calculated checksum. Otherwise, it is assumed to hold a partially
+	// calculated checksum which only reflects the pseudo header.
+	UpdateChecksumPseudoHeaderAddress(old, new tcpip.Address, fullChecksum bool)
+}
+
 // Network offers generic methods to query and/or update the fields of the
 // header of a network protocol buffer.
 type Network interface {
@@ -90,3 +115,16 @@ type Network interface {
 	// SetTOS sets the values of the "type of service" and "flow label" fields.
 	SetTOS(t uint8, l uint32)
 }
+
+// ChecksummableNetwork is a Network that supports checksumming.
+type ChecksummableNetwork interface {
+	Network
+
+	// SetSourceAddressAndChecksum sets the source address and updates the
+	// checksum to reflect the new address.
+	SetSourceAddressWithChecksumUpdate(tcpip.Address)
+
+	// SetDestinationAddressAndChecksum sets the destination address and
+	// updates the checksum to reflect the new address.
+	SetDestinationAddressWithChecksumUpdate(tcpip.Address)
+}
diff --git a/pkg/tcpip/header/ipv4.go b/pkg/tcpip/header/ipv4.go
index e9abbb709..dcc549c7b 100644
--- a/pkg/tcpip/header/ipv4.go
+++ b/pkg/tcpip/header/ipv4.go
@@ -305,6 +305,18 @@ func (b IPv4) DestinationAddress() tcpip.Address {
 	return tcpip.Address(b[dstAddr : dstAddr+IPv4AddressSize])
 }
 
+// SetSourceAddressWithChecksumUpdate implements ChecksummableNetwork.
+func (b IPv4) SetSourceAddressWithChecksumUpdate(new tcpip.Address) {
+	b.SetChecksum(^checksumUpdate2ByteAlignedAddress(^b.Checksum(), b.SourceAddress(), new))
+	b.SetSourceAddress(new)
+}
+
+// SetDestinationAddressWithChecksumUpdate implements ChecksummableNetwork.
+func (b IPv4) SetDestinationAddressWithChecksumUpdate(new tcpip.Address) {
+	b.SetChecksum(^checksumUpdate2ByteAlignedAddress(^b.Checksum(), b.DestinationAddress(), new))
+	b.SetDestinationAddress(new)
+}
+
 // padIPv4OptionsLength returns the total length for IPv4 options of length l
 // after applying padding according to RFC 791:
 //    The internet header padding is used to ensure that the internet
diff --git a/pkg/tcpip/header/tcp.go b/pkg/tcpip/header/tcp.go
index 8dabe3354..a75e51a28 100644
--- a/pkg/tcpip/header/tcp.go
+++ b/pkg/tcpip/header/tcp.go
@@ -390,6 +390,35 @@ func (b TCP) EncodePartial(partialChecksum, length uint16, seqnum, acknum uint32
 	b.SetChecksum(^checksum)
 }
 
+// SetSourcePortWithChecksumUpdate implements ChecksummableTransport.
+func (b TCP) SetSourcePortWithChecksumUpdate(new uint16) {
+	old := b.SourcePort()
+	b.SetSourcePort(new)
+	b.SetChecksum(^checksumUpdate2ByteAlignedUint16(^b.Checksum(), old, new))
+}
+
+// SetDestinationPortWithChecksumUpdate implements ChecksummableTransport.
+func (b TCP) SetDestinationPortWithChecksumUpdate(new uint16) {
+	old := b.DestinationPort()
+	b.SetDestinationPort(new)
+	b.SetChecksum(^checksumUpdate2ByteAlignedUint16(^b.Checksum(), old, new))
+}
+
+// UpdateChecksumPseudoHeaderAddress implements ChecksummableTransport.
+func (b TCP) UpdateChecksumPseudoHeaderAddress(old, new tcpip.Address, fullChecksum bool) {
+	xsum := b.Checksum()
+	if fullChecksum {
+		xsum = ^xsum
+	}
+
+	xsum = checksumUpdate2ByteAlignedAddress(xsum, old, new)
+	if fullChecksum {
+		xsum = ^xsum
+	}
+
+	b.SetChecksum(xsum)
+}
+
 // ParseSynOptions parses the options received in a SYN segment and returns the
 // relevant ones. opts should point to the option part of the TCP header.
 func ParseSynOptions(opts []byte, isAck bool) TCPSynOptions {
diff --git a/pkg/tcpip/header/udp.go b/pkg/tcpip/header/udp.go
index ae9d167ff..f69d53314 100644
--- a/pkg/tcpip/header/udp.go
+++ b/pkg/tcpip/header/udp.go
@@ -130,3 +130,32 @@ func (b UDP) Encode(u *UDPFields) {
 	binary.BigEndian.PutUint16(b[udpLength:], u.Length)
 	binary.BigEndian.PutUint16(b[udpChecksum:], u.Checksum)
 }
+
+// SetSourcePortWithChecksumUpdate implements ChecksummableTransport.
+func (b UDP) SetSourcePortWithChecksumUpdate(new uint16) {
+	old := b.SourcePort()
+	b.SetSourcePort(new)
+	b.SetChecksum(^checksumUpdate2ByteAlignedUint16(^b.Checksum(), old, new))
+}
+
+// SetDestinationPortWithChecksumUpdate implements ChecksummableTransport.
+func (b UDP) SetDestinationPortWithChecksumUpdate(new uint16) {
+	old := b.DestinationPort()
+	b.SetDestinationPort(new)
+	b.SetChecksum(^checksumUpdate2ByteAlignedUint16(^b.Checksum(), old, new))
+}
+
+// UpdateChecksumPseudoHeaderAddress implements ChecksummableTransport.
+func (b UDP) UpdateChecksumPseudoHeaderAddress(old, new tcpip.Address, fullChecksum bool) {
+	xsum := b.Checksum()
+	if fullChecksum {
+		xsum = ^xsum
+	}
+
+	xsum = checksumUpdate2ByteAlignedAddress(xsum, old, new)
+	if fullChecksum {
+		xsum = ^xsum
+	}
+
+	b.SetChecksum(xsum)
+}