From 73a7741c49524ad9c92d22c7d2f12921c1983c90 Mon Sep 17 00:00:00 2001
From: behzad nouri <behzadnouri@gmail.com>
Date: Sun, 2 Jan 2022 11:13:57 -0500
Subject: [PATCH] uses std::net::IpAddr type for Packet.Meta.addr

---
 core/benches/sigverify_stage.rs |  2 +-
 core/src/sigverify_stage.rs     |  2 +-
 sdk/src/packet.rs               | 59 ++++++++++++---------------------
 streamer/src/streamer.rs        |  6 ++--
 4 files changed, 27 insertions(+), 42 deletions(-)

diff --git a/core/benches/sigverify_stage.rs b/core/benches/sigverify_stage.rs
index 894c474ce..a0414fd76 100644
--- a/core/benches/sigverify_stage.rs
+++ b/core/benches/sigverify_stage.rs
@@ -37,7 +37,7 @@ fn bench_packet_discard(bencher: &mut Bencher) {
         .map(|_| {
             let mut addr = [0u16; 8];
             thread_rng().fill(&mut addr);
-            addr
+            std::net::IpAddr::from(addr)
         })
         .collect();
 
diff --git a/core/src/sigverify_stage.rs b/core/src/sigverify_stage.rs
index 9b63ba2b8..2297c362c 100644
--- a/core/src/sigverify_stage.rs
+++ b/core/src/sigverify_stage.rs
@@ -286,7 +286,7 @@ mod tests {
         solana_logger::setup();
         let mut batch = PacketBatch::default();
         batch.packets.resize(10, Packet::default());
-        batch.packets[3].meta.addr = [1u16; 8];
+        batch.packets[3].meta.addr = std::net::IpAddr::from([1u16; 8]);
         let mut batches = vec![batch];
         let max = 3;
         SigVerifyStage::discard_excess_packets(&mut batches, max);
diff --git a/sdk/src/packet.rs b/sdk/src/packet.rs
index 5bfc2404f..0c99eb039 100644
--- a/sdk/src/packet.rs
+++ b/sdk/src/packet.rs
@@ -3,7 +3,7 @@ use {
     serde::Serialize,
     std::{
         fmt, io,
-        net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
+        net::{IpAddr, Ipv4Addr, SocketAddr},
     },
 };
 
@@ -13,16 +13,15 @@ use {
 ///   8 bytes is the size of the fragment header
 pub const PACKET_DATA_SIZE: usize = 1280 - 40 - 8;
 
-#[derive(Clone, Default, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq)]
 #[repr(C)]
 pub struct Meta {
     pub size: usize,
     pub forwarded: bool,
     pub repair: bool,
     pub discard: bool,
-    pub addr: [u16; 8],
+    pub addr: IpAddr,
     pub port: u16,
-    pub v6: bool,
     pub is_tracer_tx: bool,
     pub is_simple_vote_tx: bool,
 }
@@ -92,40 +91,26 @@ impl PartialEq for Packet {
 
 impl Meta {
     pub fn addr(&self) -> SocketAddr {
-        if !self.v6 {
-            let addr = [
-                self.addr[0] as u8,
-                self.addr[1] as u8,
-                self.addr[2] as u8,
-                self.addr[3] as u8,
-            ];
-            let ipv4: Ipv4Addr = From::<[u8; 4]>::from(addr);
-            SocketAddr::new(IpAddr::V4(ipv4), self.port)
-        } else {
-            let ipv6: Ipv6Addr = From::<[u16; 8]>::from(self.addr);
-            SocketAddr::new(IpAddr::V6(ipv6), self.port)
-        }
+        SocketAddr::new(self.addr, self.port)
     }
 
-    pub fn set_addr(&mut self, a: &SocketAddr) {
-        match *a {
-            SocketAddr::V4(v4) => {
-                let ip = v4.ip().octets();
-                self.addr[0] = u16::from(ip[0]);
-                self.addr[1] = u16::from(ip[1]);
-                self.addr[2] = u16::from(ip[2]);
-                self.addr[3] = u16::from(ip[3]);
-                self.addr[4] = 0;
-                self.addr[5] = 0;
-                self.addr[6] = 0;
-                self.addr[7] = 0;
-                self.v6 = false;
-            }
-            SocketAddr::V6(v6) => {
-                self.addr = v6.ip().segments();
-                self.v6 = true;
-            }
-        }
-        self.port = a.port();
+    pub fn set_addr(&mut self, socket_addr: &SocketAddr) {
+        self.addr = socket_addr.ip();
+        self.port = socket_addr.port();
+    }
+}
+
+impl Default for Meta {
+    fn default() -> Self {
+        Self {
+            size: 0,
+            forwarded: false,
+            repair: false,
+            discard: false,
+            addr: IpAddr::V4(Ipv4Addr::UNSPECIFIED),
+            port: 0,
+            is_tracer_tx: false,
+            is_simple_vote_tx: false,
+        }
     }
 }
diff --git a/streamer/src/streamer.rs b/streamer/src/streamer.rs
index 58365f190..2328ff784 100644
--- a/streamer/src/streamer.rs
+++ b/streamer/src/streamer.rs
@@ -12,7 +12,7 @@ use {
     std::{
         cmp::Reverse,
         collections::HashMap,
-        net::UdpSocket,
+        net::{IpAddr, UdpSocket},
         sync::{
             atomic::{AtomicBool, Ordering},
             mpsc::{Receiver, RecvTimeoutError, SendError, Sender},
@@ -130,14 +130,14 @@ struct SendStats {
 
 #[derive(Default)]
 struct StreamerSendStats {
-    host_map: HashMap<[u16; 8], SendStats>,
+    host_map: HashMap<IpAddr, SendStats>,
     since: Option<Instant>,
 }
 
 impl StreamerSendStats {
     fn report_stats(
         name: &'static str,
-        host_map: HashMap<[u16; 8], SendStats>,
+        host_map: HashMap<IpAddr, SendStats>,
         sample_duration: Option<Duration>,
     ) {
         const MAX_REPORT_ENTRIES: usize = 5;