net/packet: implement methods for rewriting v6 addresses

Implements the ability for the address-rewriting code to support rewriting IPv6 addresses.

Specifically, UpdateSrcAddr & UpdateDstAddr.

Signed-off-by: Tom DNetto <tom@tailscale.com>
Updates https://github.com/tailscale/corp/issues/11202

net/packet/packet.go

@@ -10,6 +10,8 @@ import (
 	"net/netip"
 	"strings"
 
+	"gvisor.dev/gvisor/pkg/tcpip"
+	"gvisor.dev/gvisor/pkg/tcpip/header"
 	"tailscale.com/net/netaddr"
 	"tailscale.com/types/ipproto"
 )
@@ -453,11 +455,14 @@ func (q *Parsed) IsEchoResponse() bool {
 }
 
 // UpdateSrcAddr updates the source address in the packet buffer (e.g. during
-// SNAT). It also updates the checksum. Currently (2022-12-10) only TCP/UDP/ICMP
-// over IPv4 is supported. It panics if called with IPv6 addr.
+// SNAT). It also updates the checksum. Currently (2023-09-22) only TCP/UDP/ICMP
+// is supported. It panics if provided with an address in a different
+// family to the parsed packet.
 func (q *Parsed) UpdateSrcAddr(src netip.Addr) {
-	if q.IPVersion != 4 || src.Is6() {
-		panic("UpdateSrcAddr: only IPv4 is supported")
+	if src.Is6() && q.IPVersion != 6 {
+		panic("UpdateSrcAddr: cannot write IPv6 address to v4 packet")
+	} else if src.Is4() && q.IPVersion != 4 {
+		panic("UpdateSrcAddr: cannot write IPv4 address to v6 packet")
 	}
 	q.CaptureMeta.DidSNAT = true
 	q.CaptureMeta.OriginalSrc = q.Src
@@ -466,19 +471,27 @@ func (q *Parsed) UpdateSrcAddr(src netip.Addr) {
 	q.Src = netip.AddrPortFrom(src, q.Src.Port())
 
 	b := q.Buffer()
-	v4 := src.As4()
-	copy(b[12:16], v4[:])
-	updateV4PacketChecksums(q, old, src)
+	if src.Is6() {
+		v6 := src.As16()
+		copy(b[8:24], v6[:])
+		updateV6PacketChecksums(q, old, src)
+	} else {
+		v4 := src.As4()
+		copy(b[12:16], v4[:])
+		updateV4PacketChecksums(q, old, src)
+	}
 }
 
-// UpdateDstAddr updates the source address in the packet buffer (e.g. during
+// UpdateDstAddr updates the destination address in the packet buffer (e.g. during
 // DNAT). It also updates the checksum. Currently (2022-12-10) only TCP/UDP/ICMP
-// over IPv4 is supported. It panics if called with IPv6 addr.
+// is supported. It panics if provided with an address in a different
+// family to the parsed packet.
 func (q *Parsed) UpdateDstAddr(dst netip.Addr) {
-	if q.IPVersion != 4 || dst.Is6() {
-		panic("UpdateDstAddr: only IPv4 is supported")
+	if dst.Is6() && q.IPVersion != 6 {
+		panic("UpdateDstAddr: cannot write IPv6 address to v4 packet")
+	} else if dst.Is4() && q.IPVersion != 4 {
+		panic("UpdateDstAddr: cannot write IPv4 address to v6 packet")
 	}
-
 	q.CaptureMeta.DidDNAT = true
 	q.CaptureMeta.OriginalDst = q.Dst
 
@@ -486,9 +499,15 @@ func (q *Parsed) UpdateDstAddr(dst netip.Addr) {
 	q.Dst = netip.AddrPortFrom(dst, q.Dst.Port())
 
 	b := q.Buffer()
-	v4 := dst.As4()
-	copy(b[16:20], v4[:])
-	updateV4PacketChecksums(q, old, dst)
+	if dst.Is6() {
+		v6 := dst.As16()
+		copy(b[24:36], v6[:])
+		updateV6PacketChecksums(q, old, dst)
+	} else {
+		v4 := dst.As4()
+		copy(b[16:20], v4[:])
+		updateV4PacketChecksums(q, old, dst)
+	}
 }
 
 // EchoIDSeq extracts the identifier/sequence bytes from an ICMP Echo response,
@@ -572,13 +591,13 @@ func updateV4PacketChecksums(p *Parsed, old, new netip.Addr) {
 	tr := p.Transport()
 	switch p.IPProto {
 	case ipproto.UDP, ipproto.DCCP:
-		if len(tr) < 8 {
+		if len(tr) < header.UDPMinimumSize {
 			// Not enough space for a UDP header.
 			return
 		}
 		updateV4Checksum(tr[6:8], o4[:], n4[:])
 	case ipproto.TCP:
-		if len(tr) < 18 {
+		if len(tr) < header.TCPMinimumSize {
 			// Not enough space for a TCP header.
 			return
 		}
@@ -596,6 +615,39 @@ func updateV4PacketChecksums(p *Parsed, old, new netip.Addr) {
 	}
 }
 
+// updateV6PacketChecksums updates the checksums in the packet buffer.
+// p is modified in place.
+// If p.IPProto is unknown, no checksums are updated.
+func updateV6PacketChecksums(p *Parsed, old, new netip.Addr) {
+	if len(p.Buffer()) < 40 {
+		// Not enough space for an IPv6 header.
+		return
+	}
+	o6, n6 := tcpip.AddrFrom16Slice(old.AsSlice()), tcpip.AddrFrom16Slice(new.AsSlice())
+
+	// Now update the transport layer checksums, where applicable.
+	tr := p.Transport()
+	switch p.IPProto {
+	case ipproto.ICMPv6:
+		if len(tr) < header.ICMPv6MinimumSize {
+			return
+		}
+		header.ICMPv6(tr).UpdateChecksumPseudoHeaderAddress(o6, n6)
+	case ipproto.UDP, ipproto.DCCP:
+		if len(tr) < header.UDPMinimumSize {
+			return
+		}
+		header.UDP(tr).UpdateChecksumPseudoHeaderAddress(o6, n6, true)
+	case ipproto.TCP:
+		if len(tr) < header.TCPMinimumSize {
+			return
+		}
+		header.TCP(tr).UpdateChecksumPseudoHeaderAddress(o6, n6, true)
+	case ipproto.SCTP:
+		// No transport layer update required.
+	}
+}
+
 // updateV4Checksum calculates and updates the checksum in the packet buffer for
 // a change between old and new. The oldSum must point to the 16-bit checksum
 // field in the packet buffer that holds the old checksum value, it will be