1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174

//! Official implementation of the IPGen Spec
//!
//! This library is the official reference implementation
//! of the [IPGen Spec] for generating unique and reproducible
//! IPv4 and IPv6 addresses.
//!
//! It exposes only two simple functions `ip` and `subnet`.
//!
//! [IPGen Spec]: https://github.com/ipgen/spec
mod error;

use blake2::digest::{Update, VariableOutput};
use blake2::Blake2bVar;
use ipnetwork::Ipv6Network;
use std::net::{IpAddr, Ipv6Addr};
use std::str::FromStr;

pub use error::Error;

pub type Result<T> = std::result::Result<T, Error>;

const IP4_PREFIX: u8 = 32;
const IP6_PREFIX: u8 = 128;

/// An IP network address typically in CIDR format
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct IpNetwork(ipnetwork::IpNetwork);

impl FromStr for IpNetwork {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self> {
        Ok(Self(s.parse()?))
    }
}

/// Generates an IP address
///
/// Takes any string and a IPv4 or IPv6 network local address
/// eg `fd52:f6b0:3162::/64` or `10.0.0.0/8` and computes a unique IP address.
pub fn ip(name: &str, net: IpNetwork) -> Result<IpAddr> {
    match net.0 {
        // handle IPv6 address
        ipnetwork::IpNetwork::V6(net6) => {
            if net6.prefix() == IP6_PREFIX {
                return Err(Error::PrefixTooBig(net));
            }
            ip6(name, net6).map(IpAddr::V6)
        }
        // handle IPv4 address
        ipnetwork::IpNetwork::V4(net4) => {
            if net4.prefix() == IP4_PREFIX {
                return Err(Error::PrefixTooBig(net));
            }
            let prefix = IP6_PREFIX - IP4_PREFIX + net4.prefix();
            let net6 = format!("::{}/{prefix}", net4.ip()).parse::<Ipv6Network>()?;
            let ipv6_addr = ip6(name, net6)?.to_string();
            let ip_addr = ipv6_addr
                .strip_prefix("::")
                // This error should never happen but I'm not a fan of panicking in libraries
                .ok_or_else(|| Error::InvalidIpNetwork(format!("[BUG] the generated IPv6 address `{ipv6_addr}` does not start with the expected prefix `::`")))?
                .parse()
                // This error should never happen but I'm not a fan of panicking in libraries
                .map_err(|_| Error::InvalidIpNetwork(format!("[BUG] failed to parse the generated IP address `{}` as IPv4", ipv6_addr.trim_start_matches(':'))))
                ?;
            Ok(IpAddr::V4(ip_addr))
        }
    }
}

// Generates an IPv6 address from an IPv6 network
fn ip6(name: &str, net: Ipv6Network) -> Result<Ipv6Addr> {
    // If we divide the prefix by 4 we will get the total number
    // of characters that we must never touch.
    let network_len = net.prefix() as usize / 4;
    let ip = net.ip().segments();
    // Uncompress the IP address and throw away the semi-colons
    // so we can easily extract the network part and later
    // join it to the address part that we will compute.
    let ip_parts: Vec<String> = ip.iter().map(|b| format!("{b:04x}")).collect();
    let ip_hash = ip_parts.join("");
    let network_hash = &ip_hash[..network_len];
    // The number of characters we need to generate
    //
    // * An IPv6 address has a total number of 32 (8*4) characters.
    // * Subtracting those characters from the total in an IP address
    //   gives us the number of characters we need to generate.
    let address_len = 32 - network_len;
    // Blake2b generates hashes in multiples of 2 so we need to divide
    // the total number of characters we need by 2. However, to fully
    // utilise the address space available to us, if this leaves a
    // remainder (which will always be 1) we add it back to output length
    // and then discard the last character of the resulting hash.
    let blake_len = (address_len / 2) + (address_len % 2);
    let address_hash = hash(name.as_bytes(), blake_len)?;
    let ip_hash = format!("{}{}", network_hash, address_hash);
    let ip_str = format!(
        "{}:{}:{}:{}:{}:{}:{}:{}",
        &ip_hash[..4],
        &ip_hash[4..8],
        &ip_hash[8..12],
        &ip_hash[12..16],
        &ip_hash[16..20],
        &ip_hash[20..24],
        &ip_hash[24..28],
        &ip_hash[28..32]
    );
    let ip_addr = ip_str.parse().map_err(|_| {
        // This error should never happen but I'm not a fan of panicking in libraries
        Error::InvalidIpNetwork(format!(
            "[BUG] failed to parse the generated IP string `{ip_str}` as IPv6",
        ))
    })?;
    Ok(ip_addr)
}

/// Computes a subnet ID for any identifier
pub fn subnet(name: &str) -> Result<String> {
    hash(name.as_bytes(), 2)
}

fn hash(name: &[u8], len: usize) -> Result<String> {
    let mut hasher = Blake2bVar::new(len)
        // This error should never happen but I'm not a fan of panicking in libraries
        .map_err(|_| {
            Error::InvalidIpNetwork(format!(
                "[BUG] output length of {len} resulted in an error in hash generation",
            ))
        })?;
    hasher.update(name);
    let mut buf = vec![0u8; len];
    hasher.finalize_variable(&mut buf).map_err(|_| {
        Error::InvalidIpNetwork(format!(
            "[BUG] buffer size of {len} resulted in an error in hash generation",
        ))
    })?;
    Ok(buf.iter().map(|v| format!("{:02x}", v)).collect())
}

#[cfg(test)]
mod tests {
    #[test]
    fn ip_generation() {
        // IPv6
        let ip = crate::ip("cassandra.1", "fd9d:bb35:94bf::/48".parse().unwrap())
            .unwrap()
            .to_string();
        assert_eq!(ip, "fd9d:bb35:94bf:c38a:ee1:c75d:8df3:c909");

        // IPv4
        let ip = crate::ip("postgresql.host1", "10.0.0.0/8".parse().unwrap())
            .unwrap()
            .to_string();
        assert_eq!(ip, "10.102.194.34");

        // an empty name
        let ip = crate::ip("", "fd9d:bb35:94bf::/48".parse().unwrap())
            .unwrap()
            .to_string();
        assert_eq!(ip, "fd9d:bb35:94bf:6fa1:d8fc:fd71:9046:d762");

        // an empty name
        let ip = crate::ip("", "fd9d:bb35:94bf::/48".parse().unwrap())
            .unwrap()
            .to_string();
        assert_eq!(ip, "fd9d:bb35:94bf:6fa1:d8fc:fd71:9046:d762");
    }

    #[test]
    fn subnet_generation() {
        let subnet = crate::subnet("consul").unwrap();
        assert_eq!(subnet, "1211");
    }
}