PacketBuilder: add support for UDP

Note that IPv4 and UDP options are not supported.

Test: atest NetworkStaticLibTests
Change-Id: I15757da19311c5bae9b8b56674bef467afa7142a
diff --git a/staticlibs/device/com/android/net/module/util/PacketBuilder.java b/staticlibs/device/com/android/net/module/util/PacketBuilder.java
index e9ecb94..c908528 100644
--- a/staticlibs/device/com/android/net/module/util/PacketBuilder.java
+++ b/staticlibs/device/com/android/net/module/util/PacketBuilder.java
@@ -18,12 +18,16 @@
 
 import static android.system.OsConstants.IPPROTO_IP;
 import static android.system.OsConstants.IPPROTO_TCP;
+import static android.system.OsConstants.IPPROTO_UDP;
 
 import static com.android.net.module.util.IpUtils.ipChecksum;
 import static com.android.net.module.util.IpUtils.tcpChecksum;
+import static com.android.net.module.util.IpUtils.udpChecksum;
 import static com.android.net.module.util.NetworkStackConstants.IPV4_CHECKSUM_OFFSET;
 import static com.android.net.module.util.NetworkStackConstants.IPV4_LENGTH_OFFSET;
 import static com.android.net.module.util.NetworkStackConstants.TCP_CHECKSUM_OFFSET;
+import static com.android.net.module.util.NetworkStackConstants.UDP_CHECKSUM_OFFSET;
+import static com.android.net.module.util.NetworkStackConstants.UDP_LENGTH_OFFSET;
 
 import android.net.MacAddress;
 
@@ -32,6 +36,7 @@
 import com.android.net.module.util.structs.EthernetHeader;
 import com.android.net.module.util.structs.Ipv4Header;
 import com.android.net.module.util.structs.TcpHeader;
+import com.android.net.module.util.structs.UdpHeader;
 
 import java.io.IOException;
 import java.net.Inet4Address;
@@ -46,7 +51,7 @@
  * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
  * |                  Layer 3 header (Ipv4Header)                  |
  * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
- * |                  Layer 4 header (TcpHeader)                   |
+ * |           Layer 4 header (TcpHeader, UdpHeader)               |
  * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
  * |                           Payload                             | (optional)
  * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
@@ -73,6 +78,7 @@
 
     private int mIpv4HeaderOffset = -1;
     private int mTcpHeaderOffset = -1;
+    private int mUdpHeaderOffset = -1;
 
     public PacketBuilder(@NonNull ByteBuffer buffer) {
         mBuffer = buffer;
@@ -152,6 +158,26 @@
     }
 
     /**
+     * Write a UDP header.
+     * The UDP header length and checksum are calculated and written back in #finalizePacket.
+     *
+     * @param srcPort source port
+     * @param dstPort destination port
+     */
+    public void writeUdpHeader(short srcPort, short dstPort) throws IOException {
+        mUdpHeaderOffset = mBuffer.position();
+        final UdpHeader udpHeader = new UdpHeader(srcPort, dstPort,
+                (short) 0 /* length, calculate in #finalizePacket */,
+                (short) 0 /* checksum, calculate in #finalizePacket */);
+
+        try {
+            udpHeader.writeToByteBuffer(mBuffer);
+        } catch (IllegalArgumentException | BufferOverflowException e) {
+            throw new IOException("Error writing to buffer: ", e);
+        }
+    }
+
+    /**
      * Finalize the packet.
      *
      * Call after writing L4 header (no payload) or payload to the buffer used by the builder.
@@ -173,13 +199,21 @@
         mBuffer.putShort(mIpv4HeaderOffset + IPV4_CHECKSUM_OFFSET,
                 ipChecksum(mBuffer, mIpv4HeaderOffset /* headerOffset */));
 
-        // Populate the TCP header checksum field.
         if (mTcpHeaderOffset > 0) {
+            // Populate the TCP header checksum field.
             mBuffer.putShort(mTcpHeaderOffset + TCP_CHECKSUM_OFFSET, tcpChecksum(mBuffer,
                     mIpv4HeaderOffset /* ipOffset */, mTcpHeaderOffset /* transportOffset */,
                     mBuffer.position() - mTcpHeaderOffset /* transportLen */));
-        } else {  // TODO: add support for UDP
-            throw new IOException("Packet is missing TCP header");
+        } else if (mUdpHeaderOffset > 0) {
+            // Populate the UDP header length field.
+            mBuffer.putShort(mUdpHeaderOffset + UDP_LENGTH_OFFSET,
+                    (short) (mBuffer.position() - mUdpHeaderOffset));
+
+            // Populate the UDP header checksum field.
+            mBuffer.putShort(mUdpHeaderOffset + UDP_CHECKSUM_OFFSET, udpChecksum(mBuffer,
+                    mIpv4HeaderOffset /* ipOffset */, mUdpHeaderOffset /* transportOffset */));
+        } else {
+            throw new IOException("Packet is missing neither TCP nor UDP header");
         }
 
         mBuffer.flip();
@@ -192,7 +226,8 @@
      * @param hasEther has ethernet header. Set this flag to indicate that the packet has an
      *        ethernet header.
      * @param l3proto the layer 3 protocol. Only {@code IPPROTO_IP} currently supported.
-     * @param l4proto the layer 4 protocol. Only {@code IPPROTO_TCP} currently supported.
+     * @param l4proto the layer 4 protocol. Only {@code IPPROTO_TCP} and {@code IPPROTO_UDP}
+     *        currently supported.
      * @param payloadLen length of the payload.
      */
     @NonNull
@@ -202,8 +237,7 @@
             throw new IllegalArgumentException("Unsupported layer 3 protocol " + l3proto);
         }
 
-        if (l4proto != IPPROTO_TCP) {
-            // TODO: add support for UDP
+        if (l4proto != IPPROTO_TCP && l4proto != IPPROTO_UDP) {
             throw new IllegalArgumentException("Unsupported layer 4 protocol " + l4proto);
         }
 
@@ -214,7 +248,8 @@
         int packetLen = 0;
         if (hasEther) packetLen += Struct.getSize(EthernetHeader.class);
         packetLen += Struct.getSize(Ipv4Header.class);
-        packetLen += Struct.getSize(TcpHeader.class);
+        packetLen += (l4proto == IPPROTO_TCP) ? Struct.getSize(TcpHeader.class)
+                : Struct.getSize(UdpHeader.class);
         packetLen += payloadLen;
 
         return ByteBuffer.allocate(packetLen);
diff --git a/staticlibs/framework/com/android/net/module/util/NetworkStackConstants.java b/staticlibs/framework/com/android/net/module/util/NetworkStackConstants.java
index 378e485..353fe69 100644
--- a/staticlibs/framework/com/android/net/module/util/NetworkStackConstants.java
+++ b/staticlibs/framework/com/android/net/module/util/NetworkStackConstants.java
@@ -187,7 +187,8 @@
      *     - https://tools.ietf.org/html/rfc768
      */
     public static final int UDP_HEADER_LEN = 8;
-
+    public static final int UDP_LENGTH_OFFSET = 4;
+    public static final int UDP_CHECKSUM_OFFSET = 6;
 
     /**
      * DHCP constants.
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/PacketBuilderTest.java b/staticlibs/tests/unit/src/com/android/net/module/util/PacketBuilderTest.java
index 275aa84..8f9a1f9 100644
--- a/staticlibs/tests/unit/src/com/android/net/module/util/PacketBuilderTest.java
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/PacketBuilderTest.java
@@ -18,15 +18,18 @@
 
 import static android.system.OsConstants.IPPROTO_IP;
 import static android.system.OsConstants.IPPROTO_TCP;
+import static android.system.OsConstants.IPPROTO_UDP;
 
 import static com.android.net.module.util.NetworkStackConstants.ETHER_TYPE_IPV4;
 import static com.android.net.module.util.NetworkStackConstants.IPV4_HEADER_MIN_LEN;
 import static com.android.net.module.util.NetworkStackConstants.TCPHDR_ACK;
 import static com.android.net.module.util.NetworkStackConstants.TCP_HEADER_MIN_LEN;
+import static com.android.net.module.util.NetworkStackConstants.UDP_HEADER_LEN;
 import static com.android.testutils.MiscAsserts.assertThrows;
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
 
 import android.net.InetAddresses;
 import android.net.MacAddress;
@@ -39,6 +42,7 @@
 import com.android.net.module.util.structs.EthernetHeader;
 import com.android.net.module.util.structs.Ipv4Header;
 import com.android.net.module.util.structs.TcpHeader;
+import com.android.net.module.util.structs.UdpHeader;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -168,8 +172,92 @@
                 (byte) 0xde, (byte) 0xad, (byte) 0xbe, (byte) 0xef
             };
 
+    private static final byte[] TEST_PACKET_ETHERHDR_IPV4HDR_UDPHDR =
+            new byte[] {
+                // packet = (scapy.Ether(src="11:22:33:44:55:66", dst="aa:bb:cc:dd:ee:ff",
+                //                 type='IPv4') /
+                //           scapy.IP(src="192.0.2.1", dst="198.51.100.1",
+                //                 tos=0, id=27149, flags='DF') /
+                //           scapy.UDP(sport=9876, dport=433))
+                // Ether header
+                (byte) 0xaa, (byte) 0xbb, (byte) 0xcc, (byte) 0xdd,
+                (byte) 0xee, (byte) 0xff, (byte) 0x11, (byte) 0x22,
+                (byte) 0x33, (byte) 0x44, (byte) 0x55, (byte) 0x66,
+                (byte) 0x08, (byte) 0x00,
+                // IP header
+                (byte) 0x45, (byte) 0x00, (byte) 0x00, (byte) 0x1c,
+                (byte) 0x6a, (byte) 0x0d, (byte) 0x40, (byte) 0x00,
+                (byte) 0x40, (byte) 0x11, (byte) 0xe4, (byte) 0x8d,
+                (byte) 0xc0, (byte) 0x00, (byte) 0x02, (byte) 0x01,
+                (byte) 0xc6, (byte) 0x33, (byte) 0x64, (byte) 0x01,
+                // UDP header
+                (byte) 0x26, (byte) 0x94, (byte) 0x01, (byte) 0xb1,
+                (byte) 0x00, (byte) 0x08, (byte) 0xeb, (byte) 0x62
+            };
+
+    private static final byte[] TEST_PACKET_ETHERHDR_IPV4HDR_UDPHDR_DATA =
+            new byte[] {
+                // packet = (scapy.Ether(src="11:22:33:44:55:66", dst="aa:bb:cc:dd:ee:ff",
+                //                 type='IPv4') /
+                //           scapy.IP(src="192.0.2.1", dst="198.51.100.1",
+                //                 tos=0, id=27149, flags='DF') /
+                //           scapy.UDP(sport=9876, dport=433) /
+                //           b'\xde\xad\xbe\xef')
+                // Ether header
+                (byte) 0xaa, (byte) 0xbb, (byte) 0xcc, (byte) 0xdd,
+                (byte) 0xee, (byte) 0xff, (byte) 0x11, (byte) 0x22,
+                (byte) 0x33, (byte) 0x44, (byte) 0x55, (byte) 0x66,
+                (byte) 0x08, (byte) 0x00,
+                // IP header
+                (byte) 0x45, (byte) 0x00, (byte) 0x00, (byte) 0x20,
+                (byte) 0x6a, (byte) 0x0d, (byte) 0x40, (byte) 0x00,
+                (byte) 0x40, (byte) 0x11, (byte) 0xe4, (byte) 0x89,
+                (byte) 0xc0, (byte) 0x00, (byte) 0x02, (byte) 0x01,
+                (byte) 0xc6, (byte) 0x33, (byte) 0x64, (byte) 0x01,
+                // UDP header
+                (byte) 0x26, (byte) 0x94, (byte) 0x01, (byte) 0xb1,
+                (byte) 0x00, (byte) 0x0c, (byte) 0x4d, (byte) 0xbd,
+                // Data
+                (byte) 0xde, (byte) 0xad, (byte) 0xbe, (byte) 0xef
+            };
+
+    private static final byte[] TEST_PACKET_IPV4HDR_UDPHDR =
+            new byte[] {
+                // packet = (scapy.IP(src="192.0.2.1", dst="198.51.100.1",
+                //                 tos=0, id=27149, flags='DF') /
+                //           scapy.UDP(sport=9876, dport=433))
+                // IP header
+                (byte) 0x45, (byte) 0x00, (byte) 0x00, (byte) 0x1c,
+                (byte) 0x6a, (byte) 0x0d, (byte) 0x40, (byte) 0x00,
+                (byte) 0x40, (byte) 0x11, (byte) 0xe4, (byte) 0x8d,
+                (byte) 0xc0, (byte) 0x00, (byte) 0x02, (byte) 0x01,
+                (byte) 0xc6, (byte) 0x33, (byte) 0x64, (byte) 0x01,
+                // UDP header
+                (byte) 0x26, (byte) 0x94, (byte) 0x01, (byte) 0xb1,
+                (byte) 0x00, (byte) 0x08, (byte) 0xeb, (byte) 0x62
+            };
+
+    private static final byte[] TEST_PACKET_IPV4HDR_UDPHDR_DATA =
+            new byte[] {
+                // packet = (scapy.IP(src="192.0.2.1", dst="198.51.100.1",
+                //                 tos=0, id=27149, flags='DF') /
+                //           scapy.UDP(sport=9876, dport=433) /
+                //           b'\xde\xad\xbe\xef')
+                // IP header
+                (byte) 0x45, (byte) 0x00, (byte) 0x00, (byte) 0x20,
+                (byte) 0x6a, (byte) 0x0d, (byte) 0x40, (byte) 0x00,
+                (byte) 0x40, (byte) 0x11, (byte) 0xe4, (byte) 0x89,
+                (byte) 0xc0, (byte) 0x00, (byte) 0x02, (byte) 0x01,
+                (byte) 0xc6, (byte) 0x33, (byte) 0x64, (byte) 0x01,
+                // UDP header
+                (byte) 0x26, (byte) 0x94, (byte) 0x01, (byte) 0xb1,
+                (byte) 0x00, (byte) 0x0c, (byte) 0x4d, (byte) 0xbd,
+                // Data
+                (byte) 0xde, (byte) 0xad, (byte) 0xbe, (byte) 0xef
+            };
+
     /**
-     * Build a TCPv4 packet which has ether header, IPv4 header, TCP header and data.
+     * Build an IPv4 packet which has ether header, IPv4 header, TCP/UDP header and data.
      * The ethernet header and data are optional. Note that both source mac address and
      * destination mac address are required for ethernet header.
      *
@@ -178,30 +266,40 @@
      * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
      * |                  Layer 3 header (Ipv4Header)                  |
      * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
-     * |                  Layer 4 header (TcpHeader)                   |
+     * |           Layer 4 header (TcpHeader, UdpHeader)               |
      * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
      * |                          Payload                              | (optional)
      * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
      *
      * @param srcMac source MAC address. used by L2 ether header.
      * @param dstMac destination MAC address. used by L2 ether header.
+     * @param l4proto the layer 4 protocol. support either IPPROTO_TCP or IPPROTO_UDP.
      * @param payload the payload.
      */
     @NonNull
-    private ByteBuffer buildTcpv4Packet(@Nullable final MacAddress srcMac,
-            @Nullable final MacAddress dstMac, @Nullable final ByteBuffer payload)
+    private ByteBuffer buildIpv4Packet(@Nullable final MacAddress srcMac,
+            @Nullable final MacAddress dstMac, final int l4proto,
+            @Nullable final ByteBuffer payload)
             throws Exception {
+        if (l4proto != IPPROTO_TCP && l4proto != IPPROTO_UDP) {
+            fail("Unsupported layer 4 protocol " + l4proto);
+        }
+
         final boolean hasEther = (srcMac != null && dstMac != null);
         final int payloadLen = (payload == null) ? 0 : payload.limit();
-        final ByteBuffer buffer = PacketBuilder.allocate(hasEther, IPPROTO_IP, IPPROTO_TCP,
+        final ByteBuffer buffer = PacketBuilder.allocate(hasEther, IPPROTO_IP, l4proto,
                 payloadLen);
         final PacketBuilder packetBuilder = new PacketBuilder(buffer);
 
         if (hasEther) packetBuilder.writeL2Header(srcMac, dstMac, (short) ETHER_TYPE_IPV4);
         packetBuilder.writeIpv4Header(TYPE_OF_SERVICE, ID, FLAGS_AND_FRAGMENT_OFFSET,
-                TIME_TO_LIVE, (byte) IPPROTO_TCP, IPV4_SRC_ADDR, IPV4_DST_ADDR);
-        packetBuilder.writeTcpHeader(SRC_PORT, DST_PORT, SEQ_NO, ACK_NO,
-                TCPHDR_ACK, WINDOW, URGENT_POINTER);
+                TIME_TO_LIVE, (byte) l4proto, IPV4_SRC_ADDR, IPV4_DST_ADDR);
+        if (l4proto == IPPROTO_TCP) {
+            packetBuilder.writeTcpHeader(SRC_PORT, DST_PORT, SEQ_NO, ACK_NO,
+                    TCPHDR_ACK, WINDOW, URGENT_POINTER);
+        } else if (l4proto == IPPROTO_UDP) {
+            packetBuilder.writeUdpHeader(SRC_PORT, DST_PORT);
+        }
         if (payload != null) {
             buffer.put(payload);
             // in case data might be reused by caller, restore the position and
@@ -212,6 +310,11 @@
         return packetBuilder.finalizePacket();
     }
 
+    /**
+     * Check ethernet header.
+     *
+     * @param actual the packet to check.
+     */
     private void checkEtherHeader(final ByteBuffer actual) {
         final EthernetHeader eth = Struct.parse(EthernetHeader.class, actual);
         assertEquals(SRC_MAC, eth.srcMac);
@@ -219,7 +322,19 @@
         assertEquals(ETHER_TYPE_IPV4, eth.etherType);
     }
 
-    private void checkIpv4Header(final boolean hasData, final ByteBuffer actual) {
+    /**
+     * Check IPv4 header.
+     *
+     * @param l4proto the layer 4 protocol. support either IPPROTO_TCP or IPPROTO_UDP.
+     * @param hasData true if the packet has data payload; false otherwise.
+     * @param actual the packet to check.
+     */
+    private void checkIpv4Header(final int l4proto, final boolean hasData,
+            final ByteBuffer actual) {
+        if (l4proto != IPPROTO_TCP && l4proto != IPPROTO_UDP) {
+            fail("Unsupported layer 4 protocol " + l4proto);
+        }
+
         final Ipv4Header ipv4Header = Struct.parse(Ipv4Header.class, actual);
         assertEquals(Ipv4Header.IPHDR_VERSION_IHL, ipv4Header.vi);
         assertEquals(TYPE_OF_SERVICE, ipv4Header.tos);
@@ -230,18 +345,32 @@
         assertEquals(IPV4_DST_ADDR, ipv4Header.dstIp);
 
         final int dataLength = hasData ? DATA.limit() : 0;
-        assertEquals(IPV4_HEADER_MIN_LEN + TCP_HEADER_MIN_LEN + dataLength,
-                ipv4Header.totalLength);
-        assertEquals((byte) IPPROTO_TCP, ipv4Header.protocol);
-        assertEquals(hasData ? (short) 0xe488 : (short) 0xe48c, ipv4Header.checksum);
+        if (l4proto == IPPROTO_TCP) {
+            assertEquals(IPV4_HEADER_MIN_LEN + TCP_HEADER_MIN_LEN + dataLength,
+                    ipv4Header.totalLength);
+            assertEquals((byte) IPPROTO_TCP, ipv4Header.protocol);
+            assertEquals(hasData ? (short) 0xe488 : (short) 0xe48c, ipv4Header.checksum);
+        } else if (l4proto == IPPROTO_UDP) {
+            assertEquals(IPV4_HEADER_MIN_LEN + UDP_HEADER_LEN + dataLength,
+                    ipv4Header.totalLength);
+            assertEquals((byte) IPPROTO_UDP, ipv4Header.protocol);
+            assertEquals(hasData ? (short) 0xe489 : (short) 0xe48d, ipv4Header.checksum);
+        }
     }
 
+    /**
+     * Check TCPv4 packet.
+     *
+     * @param hasEther true if the packet has ether header; false otherwise.
+     * @param hasData true if the packet has data payload; false otherwise.
+     * @param actual the packet to check.
+     */
     private void checkTcpv4Packet(final boolean hasEther, final boolean hasData,
             final ByteBuffer actual) {
         if (hasEther) {
             checkEtherHeader(actual);
         }
-        checkIpv4Header(hasData, actual);
+        checkIpv4Header(IPPROTO_TCP, hasData, actual);
 
         final TcpHeader tcpHeader = Struct.parse(TcpHeader.class, actual);
         assertEquals(SRC_PORT, tcpHeader.srcPort);
@@ -253,18 +382,48 @@
         assertEquals(WINDOW, tcpHeader.window);
         assertEquals(hasData ? (short) 0x4844 : (short) 0xe5e5, tcpHeader.checksum);
         assertEquals(URGENT_POINTER, tcpHeader.urgentPointer);
+
+        if (hasData) {
+            assertEquals(0xdeadbeef, actual.getInt());
+        }
+    }
+
+    /**
+     * Check UDPv4 packet.
+     *
+     * @param hasEther true if the packet has ether header; false otherwise.
+     * @param hasData true if the packet has data payload; false otherwise.
+     * @param actual the packet to check.
+     */
+    private void checkUdpv4Packet(final boolean hasEther, final boolean hasData,
+            final ByteBuffer actual) {
+        if (hasEther) {
+            checkEtherHeader(actual);
+        }
+        checkIpv4Header(IPPROTO_UDP, hasData, actual);
+
+        final UdpHeader udpHeader = Struct.parse(UdpHeader.class, actual);
+        assertEquals(SRC_PORT, udpHeader.srcPort);
+        assertEquals(DST_PORT, udpHeader.dstPort);
+        final int dataLength = hasData ? DATA.limit() : 0;
+        assertEquals(UDP_HEADER_LEN + dataLength, udpHeader.length);
+        assertEquals(hasData ? (short) 0x4dbd : (short) 0xeb62, udpHeader.checksum);
+
+        if (hasData) {
+            assertEquals(0xdeadbeef, actual.getInt());
+        }
     }
 
     @Test
     public void testBuildPacketEtherIPv4Tcp() throws Exception {
-        final ByteBuffer packet = buildTcpv4Packet(SRC_MAC, DST_MAC, null /* data */);
+        final ByteBuffer packet = buildIpv4Packet(SRC_MAC, DST_MAC, IPPROTO_TCP, null /* data */);
         checkTcpv4Packet(true /* hasEther */, false /* hasData */, packet);
         assertArrayEquals(TEST_PACKET_ETHERHDR_IPV4HDR_TCPHDR, packet.array());
     }
 
     @Test
     public void testBuildPacketEtherIPv4TcpData() throws Exception {
-        final ByteBuffer packet = buildTcpv4Packet(SRC_MAC, DST_MAC, DATA);
+        final ByteBuffer packet = buildIpv4Packet(SRC_MAC, DST_MAC, IPPROTO_TCP, DATA);
         checkTcpv4Packet(true /* hasEther */, true /* hasData */, packet);
         assertArrayEquals(TEST_PACKET_ETHERHDR_IPV4HDR_TCPHDR_DATA,
                 packet.array());
@@ -272,20 +431,51 @@
 
     @Test
     public void testBuildPacketIPv4Tcp() throws Exception {
-        final ByteBuffer packet = buildTcpv4Packet(null /* srcMac */, null /* dstMac */,
-                null /* data */);
+        final ByteBuffer packet = buildIpv4Packet(null /* srcMac */, null /* dstMac */,
+                IPPROTO_TCP, null /* data */);
         checkTcpv4Packet(false /* hasEther */, false /* hasData */, packet);
         assertArrayEquals(TEST_PACKET_IPV4HDR_TCPHDR, packet.array());
     }
 
     @Test
-    public void testBuildPacketIPv4TcpData() throws Exception  {
-        final ByteBuffer packet = buildTcpv4Packet(null /* srcMac */, null /* dstMac */, DATA);
+    public void testBuildPacketIPv4TcpData() throws Exception {
+        final ByteBuffer packet = buildIpv4Packet(null /* srcMac */, null /* dstMac */,
+                IPPROTO_TCP, DATA);
         checkTcpv4Packet(false /* hasEther */, true /* hasData */, packet);
         assertArrayEquals(TEST_PACKET_IPV4HDR_TCPHDR_DATA, packet.array());
     }
 
     @Test
+    public void testBuildPacketEtherIPv4Udp() throws Exception {
+        final ByteBuffer packet = buildIpv4Packet(SRC_MAC, DST_MAC, IPPROTO_UDP, null /* data */);
+        checkUdpv4Packet(true /* hasEther */, false /* hasData */, packet);
+        assertArrayEquals(TEST_PACKET_ETHERHDR_IPV4HDR_UDPHDR, packet.array());
+    }
+
+    @Test
+    public void testBuildPacketEtherIPv4UdpData() throws Exception {
+        final ByteBuffer packet = buildIpv4Packet(SRC_MAC, DST_MAC, IPPROTO_UDP, DATA);
+        checkUdpv4Packet(true /* hasEther */, true /* hasData */, packet);
+        assertArrayEquals(TEST_PACKET_ETHERHDR_IPV4HDR_UDPHDR_DATA, packet.array());
+    }
+
+    @Test
+    public void testBuildPacketIPv4Udp() throws Exception {
+        final ByteBuffer packet = buildIpv4Packet(null /* srcMac */, null /* dstMac */,
+                IPPROTO_UDP, null /*data*/);
+        checkUdpv4Packet(false /* hasEther */, false /* hasData */, packet);
+        assertArrayEquals(TEST_PACKET_IPV4HDR_UDPHDR, packet.array());
+    }
+
+    @Test
+    public void testBuildPacketIPv4UdpData() throws Exception {
+        final ByteBuffer packet = buildIpv4Packet(null /* srcMac */, null /* dstMac */,
+                IPPROTO_UDP, DATA);
+        checkUdpv4Packet(false /* hasEther */, true /* hasData */, packet);
+        assertArrayEquals(TEST_PACKET_IPV4HDR_UDPHDR_DATA, packet.array());
+    }
+
+    @Test
     public void testFinalizePacketWithoutIpv4Header() throws Exception {
         final ByteBuffer buffer = PacketBuilder.allocate(false /* hasEther */, IPPROTO_IP,
                 IPPROTO_TCP, 0 /* payloadLen */);
@@ -297,14 +487,14 @@
     }
 
     @Test
-    public void testFinalizePacketWithoutTcpHeader() throws Exception {
+    public void testFinalizePacketWithoutL4Header() throws Exception {
         final ByteBuffer buffer = PacketBuilder.allocate(false /* hasEther */, IPPROTO_IP,
                 IPPROTO_TCP, 0 /* payloadLen */);
         final PacketBuilder packetBuilder = new PacketBuilder(buffer);
         packetBuilder.writeIpv4Header(TYPE_OF_SERVICE, ID, FLAGS_AND_FRAGMENT_OFFSET,
                 TIME_TO_LIVE, (byte) IPPROTO_TCP, IPV4_SRC_ADDR, IPV4_DST_ADDR);
-        assertThrows("java.io.IOException: Packet is missing TCP headers", IOException.class,
-                () -> packetBuilder.finalizePacket());
+        assertThrows("java.io.IOException: Packet is missing neither TCP nor UDP header",
+                IOException.class, () -> packetBuilder.finalizePacket());
     }
 
     @Test
@@ -330,6 +520,12 @@
                         TCPHDR_ACK, WINDOW, URGENT_POINTER));
     }
 
+    @Test
+    public void testWriteUdpHeaderToInsufficientBuffer() throws Exception {
+        final PacketBuilder packetBuilder = new PacketBuilder(ByteBuffer.allocate(1));
+        assertThrows(IOException.class, () -> packetBuilder.writeUdpHeader(SRC_PORT, DST_PORT));
+    }
+
     private static Inet4Address addr(String addr) {
         return (Inet4Address) InetAddresses.parseNumericAddress(addr);
     }