From 89471be3a1eec4699becbb381905e3f5365f8473 Mon Sep 17 00:00:00 2001
From: Andrew Yao <andr3wy@gmail.com>
Date: Mon, 6 May 2024 21:26:24 -0500
Subject: [PATCH] Improve efficiency of code in endpoint.rs

Signed-off-by: Andrew Yao <andrewyao@utexas.edu>
Co-authored-by: Himanshu Reddy <himanshureddy@utexas.edu>
Co-authored-by: Goldin Vo <goldinbaokimvo@utexas.edu>
Signed-off-by: Andrew Yao <andr3wy@gmail.com>
---
 src/vmm/src/dumbo/pdu/arp.rs | 126 +++++++++++++++++++++++++++--------
 1 file changed, 97 insertions(+), 29 deletions(-)

diff --git a/src/vmm/src/dumbo/pdu/arp.rs b/src/vmm/src/dumbo/pdu/arp.rs
index 1d53d9513254..749c64860da6 100644
--- a/src/vmm/src/dumbo/pdu/arp.rs
+++ b/src/vmm/src/dumbo/pdu/arp.rs
@@ -402,35 +402,103 @@ mod tests {
             EthIPv4ArpFrame::request_from_bytes(&a[..ETH_IPV4_FRAME_LEN]).unwrap_err(),
             ArpError::Operation
         );
-        
-        // 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),
-            }
-        }
+
+        // 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
+        );
     }
 
     #[test]