Changing port of ReponseBody::Pong to NonZeroU16 (#220)

ackintosh · web-flow · commit e64446b6415b · 2023-11-06T12:24:19.000+11:00
* Change `port` from u16 to NonZeroU16

* Fix tests

* Fix test: the PONG port can't be zero

* Fix clippy warnings
diff --git a/src/handler/tests.rs b/src/handler/tests.rs
@@ -7,7 +7,10 @@ use crate::{
     rpc::{Request, Response},
     ConfigBuilder, IpMode,
 };
-use std::net::{Ipv4Addr, Ipv6Addr};
+use std::{
+    convert::TryInto,
+    net::{Ipv4Addr, Ipv6Addr},
+};
 
 use crate::{
     handler::{session::build_dummy_session, HandlerOut::RequestFailed},
@@ -258,7 +261,7 @@ async fn multiple_messages() {
         body: ResponseBody::Pong {
             enr_seq: 1,
             ip: ip.into(),
-            port: sender_port,
+            port: sender_port.try_into().unwrap(),
         },
     };
 
diff --git a/src/ipmode.rs b/src/ipmode.rs
@@ -69,6 +69,17 @@ impl IpMode {
     }
 }
 
+/// Copied from the standard library. See <https://github.com/rust-lang/rust/issues/27709>
+/// The current code is behind the `ip` feature.
+pub const fn to_ipv4_mapped(ip: &std::net::Ipv6Addr) -> Option<std::net::Ipv4Addr> {
+    match ip.octets() {
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff, a, b, c, d] => {
+            Some(std::net::Ipv4Addr::new(a, b, c, d))
+        }
+        _ => None,
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -230,14 +241,3 @@ mod tests {
             .test();
     }
 }
-
-/// Copied from the standard library. See <https://github.com/rust-lang/rust/issues/27709>
-/// The current code is behind the `ip` feature.
-pub const fn to_ipv4_mapped(ip: &std::net::Ipv6Addr) -> Option<std::net::Ipv4Addr> {
-    match ip.octets() {
-        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff, a, b, c, d] => {
-            Some(std::net::Ipv4Addr::new(a, b, c, d))
-        }
-        _ => None,
-    }
-}
diff --git a/src/rpc.rs b/src/rpc.rs
@@ -1,6 +1,10 @@
 use enr::{CombinedKey, Enr};
 use rlp::{DecoderError, RlpStream};
-use std::net::{IpAddr, Ipv6Addr};
+use std::{
+    convert::TryInto,
+    net::{IpAddr, Ipv6Addr},
+    num::NonZeroU16,
+};
 use tracing::{debug, warn};
 
 /// Type to manage the request IDs.
@@ -89,7 +93,7 @@ pub enum ResponseBody {
         /// Our external IP address as observed by the responder.
         ip: IpAddr,
         /// Our external UDP port as observed by the responder.
-        port: u16,
+        port: NonZeroU16,
     },
     /// A NODES response.
     Nodes {
@@ -186,7 +190,7 @@ impl Response {
                     IpAddr::V4(addr) => s.append(&(&addr.octets() as &[u8])),
                     IpAddr::V6(addr) => s.append(&(&addr.octets() as &[u8])),
                 };
-                s.append(&port);
+                s.append(&port.get());
                 buf.extend_from_slice(&s.out());
                 buf
             }
@@ -378,15 +382,20 @@ impl Message {
                         return Err(DecoderError::RlpIncorrectListLen);
                     }
                 };
-                let port = rlp.val_at::<u16>(3)?;
-                Message::Response(Response {
-                    id,
-                    body: ResponseBody::Pong {
-                        enr_seq: rlp.val_at::<u64>(1)?,
-                        ip,
-                        port,
-                    },
-                })
+                let raw_port = rlp.val_at::<u16>(3)?;
+                if let Ok(port) = raw_port.try_into() {
+                    Message::Response(Response {
+                        id,
+                        body: ResponseBody::Pong {
+                            enr_seq: rlp.val_at::<u64>(1)?,
+                            ip,
+                            port,
+                        },
+                    })
+                } else {
+                    debug!("The port number should be non zero: {raw_port}");
+                    return Err(DecoderError::Custom("PONG response port number invalid"));
+                }
             }
             3 => {
                 // FindNodeRequest
@@ -530,7 +539,11 @@ mod tests {
         let port = 5000;
         let message = Message::Response(Response {
             id,
-            body: ResponseBody::Pong { enr_seq, ip, port },
+            body: ResponseBody::Pong {
+                enr_seq,
+                ip,
+                port: port.try_into().unwrap(),
+            },
         });
 
         // expected hex output
@@ -647,7 +660,7 @@ mod tests {
             body: ResponseBody::Pong {
                 enr_seq: 15,
                 ip: "127.0.0.1".parse().unwrap(),
-                port: 80,
+                port: 80.try_into().unwrap(),
             },
         });
 
@@ -665,7 +678,7 @@ mod tests {
             body: ResponseBody::Pong {
                 enr_seq: 15,
                 ip: IpAddr::V6(Ipv4Addr::new(192, 0, 2, 1).to_ipv6_mapped()),
-                port: 80,
+                port: 80.try_into().unwrap(),
             },
         });
 
@@ -676,7 +689,7 @@ mod tests {
             body: ResponseBody::Pong {
                 enr_seq: 15,
                 ip: IpAddr::V4(Ipv4Addr::new(192, 0, 2, 1)),
-                port: 80,
+                port: 80.try_into().unwrap(),
             },
         });
 
@@ -691,7 +704,7 @@ mod tests {
             body: ResponseBody::Pong {
                 enr_seq: 15,
                 ip: IpAddr::V6(Ipv6Addr::LOCALHOST),
-                port: 80,
+                port: 80.try_into().unwrap(),
             },
         });
 
diff --git a/src/service.rs b/src/service.rs
@@ -40,6 +40,7 @@ use parking_lot::RwLock;
 use rpc::*;
 use std::{
     collections::HashMap,
+    convert::TryInto,
     net::{IpAddr, SocketAddr},
     sync::Arc,
     task::Poll,
@@ -606,20 +607,24 @@ impl Service {
 
                 // build the PONG response
                 let src = node_address.socket_addr;
-                let response = Response {
-                    id,
-                    body: ResponseBody::Pong {
-                        enr_seq: self.local_enr.read().seq(),
-                        ip: src.ip(),
-                        port: src.port(),
-                    },
-                };
-                debug!("Sending PONG response to {}", node_address);
-                if let Err(e) = self
-                    .handler_send
-                    .send(HandlerIn::Response(node_address, Box::new(response)))
-                {
-                    warn!("Failed to send response {}", e)
+                if let Ok(port) = src.port().try_into() {
+                    let response = Response {
+                        id,
+                        body: ResponseBody::Pong {
+                            enr_seq: self.local_enr.read().seq(),
+                            ip: src.ip(),
+                            port,
+                        },
+                    };
+                    debug!("Sending PONG response to {}", node_address);
+                    if let Err(e) = self
+                        .handler_send
+                        .send(HandlerIn::Response(node_address, Box::new(response)))
+                    {
+                        warn!("Failed to send response {}", e);
+                    }
+                } else {
+                    warn!("The src port number should be non zero. {src}");
                 }
             }
             RequestBody::Talk { protocol, request } => {
@@ -778,12 +783,16 @@ impl Service {
                 ResponseBody::Pong { enr_seq, ip, port } => {
                     // Send the response to the user, if they are who asked
                     if let Some(CallbackResponse::Pong(callback)) = active_request.callback {
-                        let response = Pong { enr_seq, ip, port };
+                        let response = Pong {
+                            enr_seq,
+                            ip,
+                            port: port.get(),
+                        };
                         if let Err(e) = callback.send(Ok(response)) {
                             warn!("Failed to send callback response {:?}", e)
                         };
                     } else {
-                        let socket = SocketAddr::new(ip, port);
+                        let socket = SocketAddr::new(ip, port.get());
                         // perform ENR majority-based update if required.
 
                         // Only count votes that from peers we have contacted.
diff --git a/src/service/test.rs b/src/service/test.rs
@@ -144,7 +144,7 @@ async fn test_updating_connection_on_ping() {
         body: ResponseBody::Pong {
             enr_seq: 2,
             ip: ip2.into(),
-            port: DEFAULT_UDP_PORT,
+            port: 9000.try_into().unwrap(),
         },
     };
 
