avalanche_types/packer/
ip.rs

1use std::{
2    convert::TryInto,
3    net::{IpAddr, Ipv4Addr, Ipv6Addr},
4};
5
6use crate::{
7    errors::Result,
8    packer::{self, Packer},
9};
10
11/// All IPs (either IPv4 or IPv6) are represented as a 16-byte (IPv6) array.
12/// ref. "go/net/IP"
13pub const IP_ADDR_LEN: usize = 16;
14
15/// number of bytes per IP + port
16pub const IP_LEN: usize = IP_ADDR_LEN + packer::U16_LEN;
17
18impl Packer {
19    /// Writes the "IP" value at the offset in 16-byte representation and increments the offset afterwards.
20    /// ref. "avalanchego/utils/wrappers.Packer.PackIP"
21    /// ref. <https://pkg.go.dev/github.com/ava-labs/avalanchego/utils/wrappers#Packer.PackIP>
22    /// ref. <https://doc.rust-lang.org/std/net/enum.IpAddr.html>
23    /// ref. <https://doc.rust-lang.org/std/net/struct.Ipv4Addr.html>
24    /// ref. <https://doc.rust-lang.org/std/net/struct.Ipv6Addr.html>
25    pub fn pack_ip(&self, ip_addr: IpAddr, port: u16) -> Result<()> {
26        let ip_bytes = match ip_addr {
27            IpAddr::V4(v) => {
28                // "avalanchego" encodes IPv4 address as it is
29                // (not compatible with IPv6, e.g., prepends 2 "0xFF"s as in Rust)
30                let octets = v.octets();
31                [
32                    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, octets[0], octets[1], octets[2], octets[3],
33                ]
34            }
35            IpAddr::V6(v) => v.octets(),
36        };
37        self.pack_bytes(&ip_bytes)?;
38        self.pack_u16(port)
39    }
40
41    /// Unpacks the "IP" in the "offset" position,
42    /// and advances the cursor and offset.
43    /// ref. "avalanchego/utils/wrappers.Packer.UnpackIP"
44    /// ref. <https://pkg.go.dev/github.com/ava-labs/avalanchego/utils/wrappers#Packer.UnpackIP>
45    pub fn unpack_ip(&self) -> Result<(IpAddr, u16)> {
46        let ip = self.unpack_bytes(IP_ADDR_LEN)?;
47        let ip_array: [u8; IP_ADDR_LEN] = fix_vector_size(ip);
48
49        let ip = {
50            // check if it were IPv4 or 6
51            if all_zeroes(&ip_array[..12]) && ip_array[12] > 0 {
52                IpAddr::V4(Ipv4Addr::new(
53                    ip_array[12],
54                    ip_array[13],
55                    ip_array[14],
56                    ip_array[15],
57                ))
58            } else {
59                let ip_u128 = u128::from_be_bytes(ip_array);
60                let ipv6 = Ipv6Addr::from(ip_u128);
61                IpAddr::V6(ipv6)
62            }
63        };
64        let port = self.unpack_u16()?;
65
66        Ok((ip, port))
67    }
68
69    /// Writes the list of "IP" values at the offset in 16-byte representation
70    /// and increments the offset afterwards.
71    /// ref. "avalanchego/utils/wrappers.Packer.PackIPs"
72    /// ref. <https://pkg.go.dev/github.com/ava-labs/avalanchego/utils/wrappers#Packer.PackIPs>
73    pub fn pack_ips(&self, ips: &[(IpAddr, u16)]) -> Result<()> {
74        let n = ips.len();
75        self.pack_u32(n as u32)?;
76        for ip in ips.iter() {
77            self.pack_ip(ip.0, ip.1)?;
78        }
79        Ok(())
80    }
81
82    /// Unpacks the list of "IP"s in the "offset" position,
83    /// and advances the cursor and offset.
84    /// ref. "avalanchego/utils/wrappers.Packer.UnpackIPs"
85    /// ref. <https://pkg.go.dev/github.com/ava-labs/avalanchego/utils/wrappers#Packer.UnpackIPs>
86    pub fn unpack_ips(&self) -> Result<Vec<(IpAddr, u16)>> {
87        let n = self.unpack_u32()?;
88        let mut rs: Vec<(IpAddr, u16)> = Vec::new();
89        for _ in 0..n {
90            let b = self.unpack_ip()?;
91            rs.push(b);
92        }
93        Ok(rs)
94    }
95}
96
97fn fix_vector_size<T, const N: usize>(v: Vec<T>) -> [T; N] {
98    v.try_into()
99        .unwrap_or_else(|v: Vec<T>| panic!("expected vec length {} but {}", N, v.len()))
100}
101
102/// ref. <https://doc.rust-lang.org/std/primitive.slice.html#method.align_to>
103fn all_zeroes(d: &[u8]) -> bool {
104    let (prefix, aligned, suffix) = unsafe { d.align_to::<u128>() };
105    prefix.iter().all(|&x| x == 0)
106        && suffix.iter().all(|&x| x == 0)
107        && aligned.iter().all(|&x| x == 0)
108}
109
110/// RUST_LOG=debug cargo test --package avalanche-types --lib -- packer::ip::test_pack_and_unpack --exact --show-output
111/// ref. "avalanchego/utils/wrappers.TestPackIPCert"
112#[test]
113fn test_pack_and_unpack() {
114    use bytes::BytesMut;
115    use std::cell::Cell;
116
117    let packer = Packer {
118        max_size: IP_LEN,
119        bytes: Cell::new(BytesMut::with_capacity(0)),
120        header: false,
121        offset: Cell::new(0),
122    };
123
124    packer
125        .pack_ip(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080)
126        .unwrap();
127    assert_eq!(packer.bytes_len(), IP_LEN);
128
129    // beyond max size
130    assert!(packer
131        .pack_ip(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080)
132        .is_err());
133
134    let b = packer.take_bytes();
135    let expected: Vec<u8> = vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 127, 0, 0, 1, 31, 144];
136    assert_eq!(&b[..], &expected[..]);
137
138    // test unpack
139    let b = BytesMut::from(&expected[..]);
140    let packer = Packer {
141        max_size: 0,
142        bytes: Cell::new(b),
143        header: false,
144        offset: Cell::new(0),
145    };
146    let b = packer.unpack_ip().unwrap();
147    assert_eq!(packer.get_offset(), IP_LEN);
148
149    assert_eq!(b.0, IpAddr::V4(Ipv4Addr::LOCALHOST));
150    assert_eq!(b.1, 8080);
151
152    assert!(packer.unpack_ip().is_err());
153
154    let packer = Packer::new_with_header(4 + IP_LEN, 0);
155    packer
156        .pack_ip(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080)
157        .unwrap();
158    let expected: Vec<u8> = vec![
159        0x00, 0x00, 0x00, 18, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 127, 0, 0, 1, 31, 144,
160    ];
161    assert_eq!(&packer.take_bytes()[..], &expected[..]);
162}
163
164/// RUST_LOG=debug cargo test --package avalanche-types --lib -- packer::ip::test_packs_and_unpacks --exact --show-output
165#[test]
166fn test_packs_and_unpacks() {
167    use bytes::BytesMut;
168    use std::cell::Cell;
169
170    let packer = Packer {
171        max_size: packer::U32_LEN + IP_LEN * 3,
172        bytes: Cell::new(BytesMut::with_capacity(0)),
173        header: false,
174        offset: Cell::new(0),
175    };
176
177    packer
178        .pack_ips(&[
179            (IpAddr::V4(Ipv4Addr::LOCALHOST), 8080),
180            (IpAddr::V6(Ipv6Addr::LOCALHOST), 8081),
181            (IpAddr::V6(Ipv6Addr::new(1, 1, 1, 1, 1, 1, 1, 1)), 80),
182        ])
183        .unwrap();
184
185    let b = packer.take_bytes();
186    let expected: Vec<u8> = vec![
187        0, 0, 0, 3, // length of IPs
188        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 127, 0, 0, 1, 31,
189        144, // IpAddr::V4(Ipv4Addr::LOCALHOST), 8080
190        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 31,
191        145, // IpAddr::V6(Ipv6Addr::LOCALHOST), 8081
192        0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0,
193        80, // IpAddr::V6(Ipv6Addr::new(1, 1, 1, 1, 1, 1, 1, 1)), 80
194    ];
195    assert_eq!(&b[..], &expected[..]);
196
197    // test unpack
198    let b = BytesMut::from(&expected[..]);
199    let packer = Packer {
200        max_size: 0,
201        bytes: Cell::new(b),
202        header: false,
203        offset: Cell::new(0),
204    };
205    let b = packer.unpack_ips().unwrap();
206    assert_eq!(packer.get_offset(), packer::U32_LEN + IP_LEN * 3);
207
208    assert_eq!(b[0].0, IpAddr::V4(Ipv4Addr::LOCALHOST));
209    assert_eq!(b[0].1, 8080);
210
211    assert_eq!(b[1].0, IpAddr::V6(Ipv6Addr::LOCALHOST));
212    assert_eq!(b[1].1, 8081);
213
214    assert_eq!(b[2].0, IpAddr::V6(Ipv6Addr::new(1, 1, 1, 1, 1, 1, 1, 1)));
215    assert_eq!(b[2].1, 80);
216
217    let packer = Packer::new_with_header(1024, 0);
218    packer
219        .pack_ips(&[
220            (IpAddr::V4(Ipv4Addr::LOCALHOST), 8080),
221            (IpAddr::V6(Ipv6Addr::LOCALHOST), 8081),
222            (IpAddr::V6(Ipv6Addr::new(1, 1, 1, 1, 1, 1, 1, 1)), 80),
223        ])
224        .unwrap();
225    let expected: Vec<u8> = vec![
226        0x00, 0x00, 0x00, 58, // length of message
227        0, 0, 0, 3, // length of IPs
228        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 127, 0, 0, 1, 31,
229        144, // IpAddr::V4(Ipv4Addr::LOCALHOST), 8080
230        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 31,
231        145, // IpAddr::V6(Ipv6Addr::LOCALHOST), 8081
232        0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0,
233        80, // IpAddr::V6(Ipv6Addr::new(1, 1, 1, 1, 1, 1, 1, 1)), 80
234    ];
235    assert_eq!(&packer.take_bytes()[..], &expected[..]);
236}
avalanche_types/packer/ip.rs

avalanche_types/packer/
ip.rs
