Repetitively redefined different ARP requests were streamlined the co…

…de to be more generalized, shorter, and easier to follow, by refactoring the code. Signed-off-by: Andrew Yao <[email protected]>
firecracker-microvm · May 7, 2024 · 57fb051 · 57fb051
1 parent 81e8b68
commit 57fb051
Showing 1 changed file with 29 additions and 97 deletions.
diff --git a/src/vmm/src/dumbo/pdu/arp.rs b/src/vmm/src/dumbo/pdu/arp.rs
@@ -402,103 +402,35 @@ mod tests {
             EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap_err(),
             ArpError::Operation
         );
-
-        // TODO: The following test code is way more verbose than it should've been. Make it
-        // prettier at some point.
-
-        // Let's write a valid request.
-        EthIPv4ArpFrame::write_raw(
-            &mut a[..ETH_IPV4_FRAME_LEN],
-            HTYPE_ETHERNET,
-            ETHERTYPE_IPV4,
-            MAC_ADDR_LEN,
-            IPV4_ADDR_LEN,
-            OPER_REQUEST,
-            sha,
-            spa,
-            tha,
-            tpa,
-        )
-        .unwrap();
-        EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap();
-
-        // Now we start writing invalid requests. We've already tried with an invalid operation.
-
-        // Invalid htype.
-        EthIPv4ArpFrame::write_raw(
-            &mut a[..ETH_IPV4_FRAME_LEN],
-            HTYPE_ETHERNET + 1,
-            ETHERTYPE_IPV4,
-            MAC_ADDR_LEN,
-            IPV4_ADDR_LEN,
-            OPER_REQUEST,
-            sha,
-            spa,
-            tha,
-            tpa,
-        )
-        .unwrap();
-        assert_eq!(
-            EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap_err(),
-            ArpError::HType
-        );
-
-        // Invalid ptype.
-        EthIPv4ArpFrame::write_raw(
-            &mut a[..ETH_IPV4_FRAME_LEN],
-            HTYPE_ETHERNET,
-            ETHERTYPE_IPV4 + 1,
-            MAC_ADDR_LEN,
-            IPV4_ADDR_LEN,
-            OPER_REQUEST,
-            sha,
-            spa,
-            tha,
-            tpa,
-        )
-        .unwrap();
-        assert_eq!(
-            EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap_err(),
-            ArpError::PType
-        );
-
-        // Invalid hlen.
-        EthIPv4ArpFrame::write_raw(
-            &mut a[..ETH_IPV4_FRAME_LEN],
-            HTYPE_ETHERNET,
-            ETHERTYPE_IPV4,
-            MAC_ADDR_LEN + 1,
-            IPV4_ADDR_LEN,
-            OPER_REQUEST,
-            sha,
-            spa,
-            tha,
-            tpa,
-        )
-        .unwrap();
-        assert_eq!(
-            EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap_err(),
-            ArpError::HLen
-        );
-
-        // Invalid plen.
-        EthIPv4ArpFrame::write_raw(
-            &mut a[..ETH_IPV4_FRAME_LEN],
-            HTYPE_ETHERNET,
-            ETHERTYPE_IPV4,
-            MAC_ADDR_LEN,
-            IPV4_ADDR_LEN + 1,
-            OPER_REQUEST,
-            sha,
-            spa,
-            tha,
-            tpa,
-        )
-        .unwrap();
-        assert_eq!(
-            EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap_err(),
-            ArpError::PLen
-        );
+
+        // Various requests
+        let requests = [
+            (HTYPE_ETHERNET, ETHERTYPE_IPV4, MAC_ADDR_LEN, IPV4_ADDR_LEN, None),                            // Valid request
+            (HTYPE_ETHERNET + 1, ETHERTYPE_IPV4, MAC_ADDR_LEN, IPV4_ADDR_LEN, Some(ArpError::HType)),       // Invalid htype
+            (HTYPE_ETHERNET, ETHERTYPE_IPV4 + 1, MAC_ADDR_LEN, IPV4_ADDR_LEN, Some(ArpError::PType)),       // Invalid ptype
+            (HTYPE_ETHERNET, ETHERTYPE_IPV4, MAC_ADDR_LEN + 1, IPV4_ADDR_LEN, Some(ArpError::HLen)),        // Invalid hlen
+            (HTYPE_ETHERNET, ETHERTYPE_IPV4, MAC_ADDR_LEN, IPV4_ADDR_LEN + 1, Some(ArpError::PLen)),        // Invalid plen
+        ];
+
+        for (htype, ptype, hlen, plen, err) in requests.iter() {
+            EthIPv4ArpFrame::write_raw(
+                &mut a[..ETH_IPV4_FRAME_LEN],
+                *htype,
+                *ptype,
+                *hlen,
+                *plen,
+                OPER_REQUEST,
+                sha,
+                spa,
+                tha,
+                tpa,
+            )
+            .unwrap();
+            match err {
+                None => { EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap(); },
+                Some(arp_error) => assert_eq!(EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap_err(), *arp_error),
+            }
+        }
     }
 
     #[test]