prog, sys: add udp description and checksum

7 files changed, 174 insertions, 57 deletions

diff --git a/prog/checksum.go b/prog/checksum.go
index 94df4630a..fdc14a513 100644
--- a/prog/checksum.go
+++ b/prog/checksum.go
@@ -130,8 +130,8 @@ func getFieldByName(arg *Arg, name string) *Arg {
 	panic(fmt.Sprintf("failed to find %v field in %v", name, arg.Type.Name()))
 }
 
-func calcChecksumIPv4(arg *Arg, pid int) (*Arg, *Arg) {
-	csumField := getFieldByName(arg, "csum")
+func findCsumFieldIPv4(packet *Arg, pid int) *Arg {
+	csumField := getFieldByName(packet, "csum")
 	if typ, ok := csumField.Type.(*sys.CsumType); !ok {
 		panic(fmt.Sprintf("checksum field has bad type %v, arg: %+v", csumField.Type, csumField))
 	} else if typ.Kind != sys.CsumIPv4 {
@@ -140,11 +140,15 @@ func calcChecksumIPv4(arg *Arg, pid int) (*Arg, *Arg) {
 	if csumField.Value(pid) != 0 {
 		panic(fmt.Sprintf("checksum field has nonzero value %v, arg: %+v", csumField.Value(pid), csumField))
 	}
-	bytes := encodeArg(arg, pid)
+	return csumField
+}
+
+func calcChecksumIPv4(packet, csumField *Arg, pid int) *Arg {
+	bytes := encodeArg(packet, pid)
 	csum := ipChecksum(bytes)
 	newCsumField := *csumField
 	newCsumField.Val = uintptr(csum)
-	return csumField, &newCsumField
+	return &newCsumField
 }
 
 func extractHeaderParamsIPv4(arg *Arg) (*Arg, *Arg) {
@@ -171,29 +175,39 @@ func extractHeaderParamsIPv6(arg *Arg) (*Arg, *Arg) {
 	return srcAddr, dstAddr
 }
 
-func composeTCPPseudoHeaderIPv4(tcpPacket, srcAddr, dstAddr *Arg, pid int) []byte {
+func composePseudoHeaderIPv4(tcpPacket, srcAddr, dstAddr *Arg, protocol uint8, pid int) []byte {
 	header := []byte{}
 	header = append(header, encodeArg(srcAddr, pid)...)
 	header = append(header, encodeArg(dstAddr, pid)...)
-	header = append(header, []byte{0, 6}...) // IPPROTO_TCP == 6
+	header = append(header, []byte{0, protocol}...)
 	length := []byte{0, 0}
 	binary.BigEndian.PutUint16(length, uint16(tcpPacket.Size()))
 	header = append(header, length...)
 	return header
 }
 
-func composeTCPPseudoHeaderIPv6(tcpPacket, srcAddr, dstAddr *Arg, pid int) []byte {
+func composePseudoHeaderIPv6(tcpPacket, srcAddr, dstAddr *Arg, protocol uint8, pid int) []byte {
 	header := []byte{}
 	header = append(header, encodeArg(srcAddr, pid)...)
 	header = append(header, encodeArg(dstAddr, pid)...)
 	length := []byte{0, 0, 0, 0}
 	binary.BigEndian.PutUint32(length, uint32(tcpPacket.Size()))
 	header = append(header, length...)
-	header = append(header, []byte{0, 0, 0, 6}...) // IPPROTO_TCP == 6
+	header = append(header, []byte{0, 0, 0, protocol}...)
 	return header
 }
 
-func calcChecksumTCP(tcpPacket *Arg, pseudoHeader []byte, pid int) (*Arg, *Arg) {
+func findCsumFieldUDP(udpPacket *Arg) *Arg {
+	csumField := getFieldByName(udpPacket, "csum")
+	if typ, ok := csumField.Type.(*sys.CsumType); !ok {
+		panic(fmt.Sprintf("checksum field has bad type %v, arg: %+v", csumField.Type, csumField))
+	} else if typ.Kind != sys.CsumUDP {
+		panic(fmt.Sprintf("checksum field has bad kind %v, arg: %+v", typ.Kind, csumField))
+	}
+	return csumField
+}
+
+func findCsumFieldTCP(tcpPacket *Arg) *Arg {
 	tcpHeaderField := getFieldByName(tcpPacket, "header")
 	csumField := getFieldByName(tcpHeaderField, "csum")
 	if typ, ok := csumField.Type.(*sys.CsumType); !ok {
@@ -201,14 +215,16 @@ func calcChecksumTCP(tcpPacket *Arg, pseudoHeader []byte, pid int) (*Arg, *Arg)
 	} else if typ.Kind != sys.CsumTCP {
 		panic(fmt.Sprintf("checksum field has bad kind %v, arg: %+v", typ.Kind, csumField))
 	}
+	return csumField
+}
 
+func calcChecksumTCPUDP(packet, csumField *Arg, pseudoHeader []byte, pid int) *Arg {
 	var csum IPChecksum
 	csum.Update(pseudoHeader)
-	csum.Update(encodeArg(tcpPacket, pid))
-
+	csum.Update(encodeArg(packet, pid))
 	newCsumField := *csumField
 	newCsumField.Val = uintptr(csum.Digest())
-	return csumField, &newCsumField
+	return &newCsumField
 }
 
 func calcChecksumsCall(c *Call, pid int) map[*Arg]*Arg {
@@ -217,37 +233,48 @@ func calcChecksumsCall(c *Call, pid int) map[*Arg]*Arg {
 	ipv6HeaderParsed := false
 	var ipSrcAddr *Arg
 	var ipDstAddr *Arg
+	tcp := false
 	foreachArgArray(&c.Args, nil, func(arg, base *Arg, _ *[]*Arg) {
-		// syz_csum_ipv4_header struct is used in tests
-		if arg.Type.Name() == "ipv4_header" || arg.Type.Name() == "syz_csum_ipv4_header" {
+		// syz_csum_* structs are used in tests
+		switch arg.Type.Name() {
+		case "ipv4_header", "syz_csum_ipv4_header":
 			if csumMap == nil {
 				csumMap = make(map[*Arg]*Arg)
 			}
-			csumField, newCsumField := calcChecksumIPv4(arg, pid)
+			csumField := findCsumFieldIPv4(arg, pid)
+			newCsumField := calcChecksumIPv4(arg, csumField, pid)
 			csumMap[csumField] = newCsumField
 			ipSrcAddr, ipDstAddr = extractHeaderParamsIPv4(arg)
 			ipv4HeaderParsed = true
-		}
-		// syz_csum_ipv6_header struct is used in tests
-		if arg.Type.Name() == "ipv6_packet" || arg.Type.Name() == "syz_csum_ipv6_header" {
+		case "ipv6_packet", "syz_csum_ipv6_header":
 			ipSrcAddr, ipDstAddr = extractHeaderParamsIPv6(arg)
 			ipv6HeaderParsed = true
-		}
-		// syz_csum_tcp_packet struct is used in tests
-		if arg.Type.Name() == "tcp_packet" || arg.Type.Name() == "syz_csum_tcp_packet" {
-			if csumMap == nil {
-				csumMap = make(map[*Arg]*Arg)
-			}
+		case "tcp_packet", "syz_csum_tcp_packet":
+			tcp = true
+			fallthrough
+		case "udp_packet", "syz_csum_udp_packet":
 			if !ipv4HeaderParsed && !ipv6HeaderParsed {
-				panic("tcp packet is being parsed before ipv4 or ipv6 header")
+				panic(fmt.Sprintf("%s is being parsed before ipv4 or ipv6 header", arg.Type.Name()))
+			}
+			var csumField *Arg
+			var protocol uint8
+			if tcp {
+				csumField = findCsumFieldTCP(arg)
+				protocol = 6 // IPPROTO_TCP
+			} else {
+				csumField = findCsumFieldUDP(arg)
+				protocol = 17 // IPPROTO_UDP
 			}
 			var pseudoHeader []byte
 			if ipv4HeaderParsed {
-				pseudoHeader = composeTCPPseudoHeaderIPv4(arg, ipSrcAddr, ipDstAddr, pid)
+				pseudoHeader = composePseudoHeaderIPv4(arg, ipSrcAddr, ipDstAddr, protocol, pid)
 			} else {
-				pseudoHeader = composeTCPPseudoHeaderIPv6(arg, ipSrcAddr, ipDstAddr, pid)
+				pseudoHeader = composePseudoHeaderIPv6(arg, ipSrcAddr, ipDstAddr, protocol, pid)
+			}
+			if csumMap == nil {
+				csumMap = make(map[*Arg]*Arg)
 			}
-			csumField, newCsumField := calcChecksumTCP(arg, pseudoHeader, pid)
+			newCsumField := calcChecksumTCPUDP(arg, csumField, pseudoHeader, pid)
 			csumMap[csumField] = newCsumField
 		}
 	})
diff --git a/prog/checksum_test.go b/prog/checksum_test.go
index 7aa37a2ee..e5c2c0d32 100644
--- a/prog/checksum_test.go
+++ b/prog/checksum_test.go
@@ -60,9 +60,25 @@ func TestChecksumIP(t *testing.T) {
 			0x3250,
 		},
 		{
+			"\x00\x00\x12\x34\x56\x78",
+			0x9753,
+		},
+		{
+			"\x00\x00\x12\x34\x00\x00\x56\x78\x00\x06\x00\x04\xab\xcd",
+			0xeb7b,
+		},
+		{
 			"\x00\x11\x22\x33\x44\x55\x66\x77\x88\x99\xaa\xbb\xcc\xdd\xee\xff\xff\xee\xdd\xcc\xbb\xaa\x99\x88\x77\x66\x55\x44\x33\x22\x11\x00\x00\x00\x00\x04\x00\x00\x00\x06\x00\x00\xab\xcd",
 			0x5428,
 		},
+		{
+			"\x00\x00\x12\x34\x00\x00\x56\x78\x00\x11\x00\x04\xab\xcd",
+			0xeb70,
+		},
+		{
+			"\x00\x11\x22\x33\x44\x55\x66\x77\x88\x99\xaa\xbb\xcc\xdd\xee\xff\xff\xee\xdd\xcc\xbb\xaa\x99\x88\x77\x66\x55\x44\x33\x22\x11\x00\x00\x00\x00\x04\x00\x00\x00\x11\x00\x00\xab\xcd",
+			0x541d,
+		},
 	}
 
 	for _, test := range tests {
@@ -119,43 +135,37 @@ func TestChecksumEncode(t *testing.T) {
 	}
 }
 
-func TestChecksumIPv4Calc(t *testing.T) {
+func TestChecksumCalc(t *testing.T) {
 	tests := []struct {
 		prog string
+		kind sys.CsumKind
 		csum uint16
 	}{
 		{
-			"syz_test$csum_ipv4(&(0x7f0000000000)={0x0, {0x42, 0x43, [0x44, 0x45], 0xa, 0xb, \"aabbccdd\"}, 0x0, 0x0})",
-			0xe143,
+			"syz_test$csum_ipv4(&(0x7f0000000000)={0x0, 0x1234, 0x5678})",
+			sys.CsumIPv4,
+			0x9753,
 		},
-	}
-	for i, test := range tests {
-		p, err := Deserialize([]byte(test.prog))
-		if err != nil {
-			t.Fatalf("failed to deserialize prog %v: %v", test.prog, err)
-		}
-		_, csumField := calcChecksumIPv4(p.Calls[0].Args[0].Res, i%32)
-		// Can't compare serialized progs, since checksums are zerod on serialization.
-		csum := csumField.Value(i % 32)
-		if csum != uintptr(test.csum) {
-			t.Fatalf("failed to calc ipv4 checksum, got %x, want %x, prog: '%v'", csum, test.csum, test.prog)
-		}
-	}
-}
-
-func TestChecksumTCPCalc(t *testing.T) {
-	tests := []struct {
-		prog string
-		csum uint16
-	}{
 		{
-			"syz_test$csum_ipv4_tcp(&(0x7f0000000000)={{0x0, {0x42, 0x43, [0x44, 0x45], 0xa, 0xb, \"aabbccdd\"}, 0x0, 0x0}, {{0x0}, \"abcd\"}})",
-			0x5428,
+			"syz_test$csum_ipv4_tcp(&(0x7f0000000000)={{0x0, 0x1234, 0x5678}, {{0x0}, \"abcd\"}})",
+			sys.CsumTCP,
+			0xeb7b,
 		},
 		{
 			"syz_test$csum_ipv6_tcp(&(0x7f0000000000)={{\"00112233445566778899aabbccddeeff\", \"ffeeddccbbaa99887766554433221100\"}, {{0x0}, \"abcd\"}})",
+			sys.CsumTCP,
 			0x5428,
 		},
+		{
+			"syz_test$csum_ipv4_udp(&(0x7f0000000000)={{0x0, 0x1234, 0x5678}, {0x0, \"abcd\"}})",
+			sys.CsumUDP,
+			0xeb70,
+		},
+		{
+			"syz_test$csum_ipv6_udp(&(0x7f0000000000)={{\"00112233445566778899aabbccddeeff\", \"ffeeddccbbaa99887766554433221100\"}, {0x0, \"abcd\"}})",
+			sys.CsumUDP,
+			0x541d,
+		},
 	}
 	for i, test := range tests {
 		p, err := Deserialize([]byte(test.prog))
@@ -163,18 +173,23 @@ func TestChecksumTCPCalc(t *testing.T) {
 			t.Fatalf("failed to deserialize prog %v: %v", test.prog, err)
 		}
 		csumMap := calcChecksumsCall(p.Calls[0], i%32)
+		found := false
 		for oldField, newField := range csumMap {
 			if typ, ok := newField.Type.(*sys.CsumType); ok {
-				if typ.Kind == sys.CsumTCP {
+				if typ.Kind == test.kind {
+					found = true
 					csum := newField.Value(i % 32)
 					if csum != uintptr(test.csum) {
-						t.Fatalf("failed to calc tcp checksum, got %x, want %x, prog: '%v'", csum, test.csum, test.prog)
+						t.Fatalf("failed to calc checksum, got %x, want %x, kind %v, prog '%v'", csum, test.csum, test.kind, test.prog)
 					}
 				}
 			} else {
 				t.Fatalf("non csum key %+v in csum map %+v", oldField, csumMap)
 			}
 		}
+		if !found {
+			t.Fatalf("csum field not found, kind %v, prog '%v'", test.kind, test.prog);
+		}
 	}
 }
 
diff --git a/sys/decl.go b/sys/decl.go
index ba1806836..77fbfbc51 100644
--- a/sys/decl.go
+++ b/sys/decl.go
@@ -195,6 +195,7 @@ type CsumKind int
 const (
 	CsumIPv4 CsumKind = iota
 	CsumTCP
+	CsumUDP
 )
 
 type CsumType struct {
diff --git a/sys/socket.txt b/sys/socket.txt
index de9f44564..075f0536c 100644
--- a/sys/socket.txt
+++ b/sys/socket.txt
@@ -214,6 +214,26 @@ tcp_pair {
 	f1	sock_tcp
 }
 
+# AF_INET: UDP support
+
+resource sock_udp[sock]
+
+socket$udp(domain const[AF_INET], type const[SOCK_DGRAM], proto const[0]) sock_udp
+socketpair$udp(domain const[AF_INET], type const[SOCK_DGRAM], proto const[0], fds ptr[out, udp_pair])
+accept$udp(fd sock_udp, peer ptr[out, sockaddr_in, opt], peerlen ptr[inout, len[peer, int32]]) sock_udp
+accept4$udp(fd sock_udp, peer ptr[out, sockaddr_in, opt], peerlen ptr[inout, len[peer, int32]], flags flags[accept_flags]) sock_udp
+bind$udp(fd sock_udp, addr ptr[in, sockaddr_in], addrlen len[addr])
+connect$udp(fd sock_udp, addr ptr[in, sockaddr_in], addrlen len[addr])
+sendto$udp(fd sock_udp, buf buffer[in], len len[buf], f flags[send_flags], addr ptr[in, sockaddr_in, opt], addrlen len[addr])
+recvfrom$udp(fd sock_udp, buf buffer[out], len len[buf], f flags[recv_flags], addr ptr[in, sockaddr_in, opt], addrlen len[addr])
+getsockname$udp(fd sock_udp, addr ptr[out, sockaddr_in], addrlen ptr[inout, len[addr, int32]])
+getpeername$udp(fd sock_udp, peer ptr[out, sockaddr_in], peerlen ptr[inout, len[peer, int32]])
+
+udp_pair {
+	f0	sock_udp
+	f1	sock_udp
+}
+
 
 
 
@@ -237,6 +257,26 @@ tcp6_pair {
 	f1	sock_tcp6
 }
 
+# AF_INET6: UDP support
+
+resource sock_udp6[sock]
+
+socket$udp6(domain const[AF_INET6], type const[SOCK_DGRAM], proto const[0]) sock_udp6
+socketpair$udp6(domain const[AF_INET6], type const[SOCK_DGRAM], proto const[0], fds ptr[out, udp6_pair])
+accept$udp6(fd sock_udp6, peer ptr[out, sockaddr_in6, opt], peerlen ptr[inout, len[peer, int32]]) sock_udp6
+accept4$udp6(fd sock_udp6, peer ptr[out, sockaddr_in6, opt], peerlen ptr[inout, len[peer, int32]], flags flags[accept_flags]) sock_udp6
+bind$udp6(fd sock_udp6, addr ptr[in, sockaddr_in6], addrlen len[addr])
+connect$udp6(fd sock_udp6, addr ptr[in, sockaddr_in6], addrlen len[addr])
+sendto$udp6(fd sock_udp6, buf buffer[in], len len[buf], f flags[send_flags], addr ptr[in, sockaddr_in6, opt], addrlen len[addr])
+recvfrom$udp6(fd sock_udp6, buf buffer[out], len len[buf], f flags[recv_flags], addr ptr[in, sockaddr_in6, opt], addrlen len[addr])
+getsockname$udp6(fd sock_udp6, addr ptr[out, sockaddr_in6], addrlen ptr[inout, len[addr, int32]])
+getpeername$udp6(fd sock_udp6, peer ptr[out, sockaddr_in6], peerlen ptr[inout, len[peer, int32]])
+
+udp6_pair {
+	f0	sock_udp6
+	f1	sock_udp6
+}
+
 
 
 
diff --git a/sys/test.txt b/sys/test.txt
index 5d1d3266d..ae7b0b7f3 100644
--- a/sys/test.txt
+++ b/sys/test.txt
@@ -396,6 +396,8 @@ syz_test$csum_encode(a0 ptr[in, syz_csum_encode])
 syz_test$csum_ipv4(a0 ptr[in, syz_csum_ipv4_header])
 syz_test$csum_ipv4_tcp(a0 ptr[in, syz_csum_ipv4_tcp_packet])
 syz_test$csum_ipv6_tcp(a0 ptr[in, syz_csum_ipv6_tcp_packet])
+syz_test$csum_ipv4_udp(a0 ptr[in, syz_csum_ipv4_udp_packet])
+syz_test$csum_ipv6_udp(a0 ptr[in, syz_csum_ipv6_udp_packet])
 
 syz_csum_encode {
 	f0	int16
@@ -408,7 +410,6 @@ syz_csum_encode {
 
 syz_csum_ipv4_header {
 	csum		csum[ipv4, int16]
-	data		syz_csum_encode
 	src_ip		int32be
 	dst_ip		int32be
 } [packed]
@@ -436,3 +437,18 @@ syz_csum_ipv6_tcp_packet {
 	header		syz_csum_ipv6_header
 	payload		syz_csum_tcp_packet
 } [packed]
+
+syz_csum_udp_packet {
+	csum		csum[udp, int16]
+	payload		array[int8]
+} [packed]
+
+syz_csum_ipv4_udp_packet {
+	header		syz_csum_ipv4_header
+	payload		syz_csum_udp_packet
+} [packed]
+
+syz_csum_ipv6_udp_packet {
+	header		syz_csum_ipv6_header
+	payload		syz_csum_udp_packet
+} [packed]
diff --git a/sys/vnet.txt b/sys/vnet.txt
index 49582cae4..794103282 100644
--- a/sys/vnet.txt
+++ b/sys/vnet.txt
@@ -378,8 +378,24 @@ ipv6_packet_payload {
 ###################################### IP ######################################
 ################################################################################
 
-ip_payload {
+ip_payload [
 	tcp		tcp_packet
+	udp		udp_packet
+] [varlen]
+
+################################################################################
+###################################### UDP #####################################
+################################################################################
+
+# https://tools.ietf.org/html/rfc768
+# https://en.wikipedia.org/wiki/User_Datagram_Protocol#Packet_structure
+
+udp_packet {
+	src_port	proc[int16be, 20000, 4]
+	dst_port	proc[int16be, 20000, 4]
+	length		len[parent, int16be]
+	csum		csum[udp, int16be]
+	data		array[int8]
 } [packed]
 
 ################################################################################
diff --git a/sysgen/sysgen.go b/sysgen/sysgen.go
index 3da192472..918c6d847 100644
--- a/sysgen/sysgen.go
+++ b/sysgen/sysgen.go
@@ -511,6 +511,8 @@ func generateArg(
 			kind = "CsumIPv4"
 		case "tcp":
 			kind = "CsumTCP"
+		case "udp":
+			kind = "CsumUDP"
 		default:
 			failf("unknown checksum kind '%v'", a[0])
 		}