net/tstun: add inital support for NAT v4

This adds support in tstun to utitilize the SelfNodeV4MasqAddrForThisPeer and
perform the necessary modifications to the packet as it passes through tstun.

Currently this only handles ICMP, UDP and TCP traffic.
Subnet routers and Exit Nodes are also unsupported.

Updates tailscale/corp#8020

Co-authored-by: Melanie Warrick <warrick@tailscale.com>
Signed-off-by: Maisem Ali <maisem@tailscale.com>

net/packet/packet.go

@@ -440,6 +440,40 @@ 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.
+func (q *Parsed) UpdateSrcAddr(src netip.Addr) {
+	if q.IPVersion != 4 || src.Is6() {
+		panic("UpdateSrcAddr: only IPv4 is supported")
+	}
+
+	old := q.Src.Addr()
+	q.Src = netip.AddrPortFrom(src, q.Src.Port())
+
+	b := q.Buffer()
+	v4 := src.As4()
+	copy(b[12:16], v4[:])
+	updateV4PacketChecksums(q, old, src)
+}
+
+// UpdateDstAddr updates the source 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.
+func (q *Parsed) UpdateDstAddr(dst netip.Addr) {
+	if q.IPVersion != 4 || dst.Is6() {
+		panic("UpdateDstAddr: only IPv4 is supported")
+	}
+
+	old := q.Dst.Addr()
+	q.Dst = netip.AddrPortFrom(dst, q.Dst.Port())
+
+	b := q.Buffer()
+	v4 := dst.As4()
+	copy(b[16:20], v4[:])
+	updateV4PacketChecksums(q, old, dst)
+}
+
 // EchoIDSeq extracts the identifier/sequence bytes from an ICMP Echo response,
 // and returns them as a uint32, used to lookup internally routed ICMP echo
 // responses. This function is intentionally lightweight as it is called on
@@ -502,3 +536,69 @@ func withIP(ap netip.AddrPort, ip netip.Addr) netip.AddrPort {
 func withPort(ap netip.AddrPort, port uint16) netip.AddrPort {
 	return netip.AddrPortFrom(ap.Addr(), port)
 }
+
+// updateV4PacketChecksums updates the checksums in the packet buffer.
+// Currently (2023-03-01) only TCP/UDP/ICMP over IPv4 is supported.
+// p is modified in place.
+// If p.IPProto is unknown, only the IP header checksum is updated.
+// TODO(maisem): more protocols (sctp, gre, dccp)
+func updateV4PacketChecksums(p *Parsed, old, new netip.Addr) {
+	o4, n4 := old.As4(), new.As4()
+	updateV4Checksum(p.Buffer()[10:12], o4[:], n4[:]) // header
+	switch p.IPProto {
+	case ipproto.UDP:
+		updateV4Checksum(p.Transport()[6:8], o4[:], n4[:])
+	case ipproto.TCP:
+		updateV4Checksum(p.Transport()[16:18], o4[:], n4[:])
+	case ipproto.ICMPv4:
+		// Nothing to do.
+	}
+	// TODO(maisem): more protocols (sctp, gre, dccp)
+}
+
+// updateV4Checksum calculates and updates the checksum in the packet buffer
+// for a change between old and new. The checksum is updated in place.
+func updateV4Checksum(oldSum, old, new []byte) {
+	if len(old) != len(new) {
+		panic("old and new must be the same length")
+	}
+	if len(old)%2 != 0 {
+		panic("old and new must be even length")
+	}
+	/*
+		RFC 1624
+		Given the following notation:
+
+		    HC  - old checksum in header
+		    C   - one's complement sum of old header
+		    HC' - new checksum in header
+		    C'  - one's complement sum of new header
+		    m   - old value of a 16-bit field
+		    m'  - new value of a 16-bit field
+
+		    HC' = ~(C + (-m) + m')  --    [Eqn. 3]
+		    HC' = ~(~HC + ~m + m')
+
+		This can be simplified to:
+		    HC' = ~(C + ~m + m')    --    [Eqn. 3]
+		    HC' = ~C'
+		    C'  = C + ~m + m'
+	*/
+
+	c := uint32(^binary.BigEndian.Uint16(oldSum))
+
+	cPrime := c
+	for len(new) > 0 {
+		mNot := uint32(^binary.BigEndian.Uint16(old[:2]))
+		mPrime := uint32(binary.BigEndian.Uint16(new[:2]))
+		cPrime += mPrime + mNot
+		new, old = new[2:], old[2:]
+	}
+
+	// Account for overflows by adding the carry bits back into the sum.
+	for (cPrime >> 16) > 0 {
+		cPrime = cPrime&0xFFFF + cPrime>>16
+	}
+	hcPrime := ^uint16(cPrime)
+	binary.BigEndian.PutUint16(oldSum, hcPrime)
+}