Refactor recv code into a function.

While here, check for TTL value to be 255.

1 files changed, 14 insertions, 9 deletions

diff --git a/src/main.rs b/src/main.rs
index bda3c70..17bb2a6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,7 +2,7 @@ use socket2::{Domain, Protocol, Socket, Type};
 use std::io::{Error, ErrorKind};
 use std::mem::MaybeUninit;
 use std::net::Ipv4Addr;
-use vrrpd::vrrpv2;
+use vrrpd::vrrpv2::{self, VRRPv2, VRRPv2Error};
 
 fn main() -> std::io::Result<()> {
     let sock = Socket::new(Domain::IPV4, Type::RAW, Some(Protocol::from(112)))?;
@@ -10,19 +10,24 @@ fn main() -> std::io::Result<()> {
         &Ipv4Addr::new(224, 0, 0, 18),
         &Ipv4Addr::UNSPECIFIED,
     )?;
+    let vrrp = recv_vrrp(sock).map_err(|err| {
+        Error::new(ErrorKind::InvalidData, format!("recv_vrrp(): {err}"))
+    })?;
+    eprintln!("vrrp: {vrrp:?}");
+    Ok(())
+}
+
+fn recv_vrrp(sock: Socket) -> Result<VRRPv2, VRRPv2Error> {
     let mut buf = [MaybeUninit::<u8>::uninit(); 512];
     let (nbytes, _) = sock.recv_from(&mut buf)?;
     let bytes = unsafe {
         std::slice::from_raw_parts(buf.as_ptr() as *const u8, nbytes)
     };
+    let ttl: u8 = bytes[8];
+    if ttl != 255 {
+        return Err(VRRPv2Error::InvalidTTL);
+    }
     let ihl = bytes[0] & 0xF;
     let data_offset = (ihl * 4) as usize;
-    let vrrp = vrrpv2::from_bytes(&bytes[data_offset..]).map_err(|err| {
-        Error::new(
-            ErrorKind::InvalidData,
-            format!("vrrpv2::from_bytes(): {err}"),
-        )
-    })?;
-    eprintln!("vrrp: {vrrp:?}");
-    Ok(())
+    vrrpv2::from_bytes(&bytes[data_offset..])
 }