Fix UDP checksum calculation

A zero value for the UDP checksum indicates that the sender did not
include a checksum. This is allowed for IPv4 but not for IPv6. A
package may be specially crafted such that the checksum naturally has a
zero valued checksum. However, UDP requires that when this happens, the
sender must transmit the all ones value instead.

PiperOrigin-RevId: 429113399

pkg/tcpip/checker/checker.go

@@ -565,6 +565,17 @@ func DstPort(port uint16) TransportChecker {
 	}
 }
 
+// TransportChecksum creates a checker that checks the checksum value.
+func TransportChecksum(want uint16) TransportChecker {
+	return func(t *testing.T, transportHdr header.Transport) {
+		t.Helper()
+
+		if got := transportHdr.Checksum(); got != want {
+			t.Errorf("got transportHdr.Checksum() = %d, want = %d", got, want)
+		}
+	}
+}
+
 // NoChecksum creates a checker that checks if the checksum is zero.
 func NoChecksum(noChecksum bool) TransportChecker {
 	return func(t *testing.T, h header.Transport) {

pkg/tcpip/transport/udp/endpoint.go

@@ -17,6 +17,7 @@ package udp
 import (
 	"fmt"
 	"io"
+	"math"
 	"time"
 
 	"gvisor.dev/gvisor/pkg/sync"
@@ -471,10 +472,33 @@ func (e *endpoint) write(p tcpip.Payloader, opts tcpip.WriteOptions) (int64, tcp
 	// On IPv6, UDP checksum is not optional (RFC2460 Section 8.1).
 	if pktInfo.RequiresTXTransportChecksum &&
 		(!e.ops.GetNoChecksum() || pktInfo.NetProto == header.IPv6ProtocolNumber) {
-		udp.SetChecksum(^udp.CalculateChecksum(header.ChecksumCombine(
+		xsum := udp.CalculateChecksum(header.ChecksumCombine(
 			header.PseudoHeaderChecksum(ProtocolNumber, pktInfo.LocalAddress, pktInfo.RemoteAddress, length),
 			pkt.Data().AsRange().Checksum(),
-		)))
+		))
+		// As per RFC 768 page 2,
+		//
+		//   Checksum is the 16-bit one's complement of the one's complement sum of
+		//   a pseudo header of information from the IP header, the UDP header, and
+		//   the data, padded with zero octets at the end (if necessary) to make a
+		//   multiple of two octets.
+		//
+		//	 The pseudo header conceptually prefixed to the UDP header contains the
+		//   source address, the destination address, the protocol, and the UDP
+		//   length. This information gives protection against misrouted datagrams.
+		//   This checksum procedure is the same as is used in TCP.
+		//
+		//   If the computed checksum is zero, it is transmitted as all ones (the
+		//   equivalent in one's complement arithmetic). An all zero transmitted
+		//   checksum value means that the transmitter generated no checksum (for
+		//   debugging or for higher level protocols that don't care).
+		//
+		// To avoid the zero value, we only calculate the one's complement of the
+		// one's complement sum if the sum is not all ones.
+		if xsum != math.MaxUint16 {
+			xsum = ^xsum
+		}
+		udp.SetChecksum(xsum)
 	}
 	if err := udpInfo.ctx.WritePacket(pkt, false /* headerIncluded */); err != nil {
 		e.stack.Stats().UDP.PacketSendErrors.Increment()

pkg/tcpip/transport/udp/udp_test.go

@@ -16,6 +16,7 @@ package udp_test
 
 import (
 	"bytes"
+	"encoding/binary"
 	"fmt"
 	"io/ioutil"
 	"math"
@@ -1981,6 +1982,77 @@ func TestOutgoingSubnetBroadcast(t *testing.T) {
 	}
 }
 
+func TestChecksumWithZeroValueOnesComplementSum(t *testing.T) {
+	c := context.New(t, []stack.TransportProtocolFactory{udp.NewProtocol})
+	defer c.Cleanup()
+
+	c.CreateEndpoint(ipv6.ProtocolNumber, udp.ProtocolNumber)
+	var writeOpts tcpip.WriteOptions
+	h := context.UnicastV6.MakeHeader4Tuple(context.Outgoing)
+	writeDstAddr := context.UnicastV6.MapAddrIfApplicable(h.Dst.Addr)
+	writeOpts = tcpip.WriteOptions{
+		To: &tcpip.FullAddress{Addr: writeDstAddr, Port: h.Dst.Port},
+	}
+
+	// Write a packet to calculate what the checksum value will be with a zero
+	// value payload. We will then take that checksum value to construct another
+	// packet which would result in the ones complement of the packet to be zero.
+	var payload [2]byte
+	{
+		var r bytes.Reader
+		r.Reset(payload[:])
+		n, err := c.EP.Write(&r, writeOpts)
+		if err != nil {
+			t.Fatalf("Write failed: %s", err)
+		}
+		if want := int64(len(payload)); n != want {
+			t.Fatalf("got n = %d, want = %d", n, want)
+		}
+
+		pkt := c.LinkEP.Read()
+		if pkt == nil {
+			t.Fatal("Packet wasn't written out")
+		}
+
+		v := stack.PayloadSince(pkt.NetworkHeader())
+		checker.IPv6(t, v, checker.UDP())
+
+		// Simply replacing the payload with the checksum value is enough to make
+		// sure that we end up with an all ones value for the ones complement sum
+		// because the checksum value is held the ones complement of the ones
+		// complement sum.
+		//
+		// In ones complement arithmetic, adding a value A with a ones complement of
+		// another value B is the same as subtracting B from A.
+		//
+		// The resulting ones complement will be  C' = C - C so we know C' will be
+		// zero. The stack should never send a zero value though so we expect all
+		// ones below.
+		binary.BigEndian.PutUint16(payload[:], header.UDP(header.IPv6(v).Payload()).Checksum())
+	}
+
+	{
+		var r bytes.Reader
+		r.Reset(payload[:])
+		n, err := c.EP.Write(&r, writeOpts)
+		if err != nil {
+			t.Fatalf("Write failed: %s", err)
+		}
+		if want := int64(len(payload)); n != want {
+			t.Fatalf("got n = %d, want = %d", n, want)
+		}
+	}
+
+	{
+		pkt := c.LinkEP.Read()
+		if pkt == nil {
+			t.Fatal("Packet wasn't written out")
+		}
+
+		checker.IPv6(t, stack.PayloadSince(pkt.NetworkHeader()), checker.UDP(checker.TransportChecksum(math.MaxUint16)))
+	}
+}
+
 func TestMain(m *testing.M) {
 	refs.SetLeakMode(refs.LeaksPanic)
 	code := m.Run()