diesel/pg/types/

network_address.rs

extern crate ipnetwork;
extern crate libc;

use self::ipnetwork::{IpNetwork, Ipv4Network, Ipv6Network};
use core::net::{Ipv4Addr, Ipv6Addr};
use std::io::prelude::*;

use crate::deserialize::{self, FromSql, FromSqlRow};
use crate::pg::{Pg, PgValue};
#[cfg(test)]
use crate::query_builder::bind_collector::ByteWrapper;
use crate::serialize::{self, IsNull, Output, ToSql};
use crate::sql_types::{Cidr, Inet};

#[cfg(windows)]
const AF_INET: u8 = 2;
// Maybe not used, but defining to follow Rust's libstd/net/sys
#[cfg(target_os = "redox")]
const AF_INET: u8 = 1;

#[allow(clippy::cast_possible_truncation)] // it's 2
#[cfg(not(any(windows, target_os = "redox")))]
const AF_INET: u8 = libc::AF_INET as u8;

const PGSQL_AF_INET: u8 = AF_INET;
const PGSQL_AF_INET6: u8 = AF_INET + 1;

#[allow(dead_code)]
mod foreign_derives {
    use super::*;
    use crate::expression::AsExpression;

    #[derive(AsExpression, FromSqlRow)]
    #[diesel(foreign_derive)]
    #[diesel(sql_type = Inet)]
    #[diesel(sql_type = Cidr)]
    struct IpNetworkProxy(IpNetwork);
}

macro_rules! err {
    () => {
        Err("invalid network address format".into())
    };
    ($msg:expr_2021) => {
        Err(format!("invalid network address format. {}", $msg).into())
    };
}

macro_rules! assert_or_error {
    ($cond:expr_2021) => {
        if !$cond {
            return err!();
        }
    };

    ($cond:expr_2021, $msg:expr_2021) => {
        if !$cond {
            return err!($msg);
        }
    };
}

macro_rules! impl_Sql {
    ($ty: ty, $net_type: expr_2021) => {
        #[cfg(all(feature = "postgres_backend", feature = "network-address"))]
        impl FromSql<$ty, Pg> for IpNetwork {
            fn from_sql(value: PgValue<'_>) -> deserialize::Result<Self> {
                // https://github.com/postgres/postgres/blob/55c3391d1e6a201b5b891781d21fe682a8c64fe6/src/include/utils/inet.h#L23-L28
                let bytes = value.as_bytes();
                assert_or_error!(4 <= bytes.len(), "input is too short.");
                let af = bytes[0];
                let prefix = bytes[1];
                let net_type = bytes[2];
                let len = bytes[3];
                assert_or_error!(
                    net_type == $net_type,
                    format!("returned type isn't a {}", stringify!($ty))
                );
                if af == PGSQL_AF_INET {
                    assert_or_error!(bytes.len() == 8);
                    assert_or_error!(len == 4, "the data isn't the size of ipv4");
                    let b = &bytes[4..];
                    let addr = Ipv4Addr::new(b[0], b[1], b[2], b[3]);
                    let inet = Ipv4Network::new(addr, prefix)?;
                    Ok(IpNetwork::V4(inet))
                } else if af == PGSQL_AF_INET6 {
                    assert_or_error!(bytes.len() == 20);
                    assert_or_error!(len == 16, "the data isn't the size of ipv6");
                    let b = &bytes[4..];
                    let addr = Ipv6Addr::from([
                        b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7], b[8], b[9], b[10], b[11],
                        b[12], b[13], b[14], b[15],
                    ]);
                    let inet = Ipv6Network::new(addr, prefix)?;
                    Ok(IpNetwork::V6(inet))
                } else {
                    err!()
                }
            }
        }

        #[cfg(all(feature = "postgres_backend", feature = "network-address"))]
        impl ToSql<$ty, Pg> for IpNetwork {
            fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> serialize::Result {
                use self::ipnetwork::IpNetwork::*;
                let net_type = $net_type;
                match *self {
                    V4(ref net) => {
                        let mut data = [0u8; 8];
                        let af = PGSQL_AF_INET;
                        let prefix = net.prefix();
                        let len: u8 = 4;
                        let addr = if net_type == 0 {
                            net.ip().octets()
                        } else {
                            network_v4(net).octets()
                        };
                        data[0] = af;
                        data[1] = prefix;
                        data[2] = net_type;
                        data[3] = len;
                        data[4..].copy_from_slice(&addr);
                        out.write_all(&data).map(|_| IsNull::No).map_err(Into::into)
                    }
                    V6(ref net) => {
                        let mut data = [0u8; 20];
                        let af = PGSQL_AF_INET6;
                        let prefix = net.prefix();
                        let len: u8 = 16;
                        let addr = if net_type == 0 {
                            net.ip().octets()
                        } else {
                            network_v6(net).octets()
                        };
                        data[0] = af;
                        data[1] = prefix;
                        data[2] = net_type;
                        data[3] = len;
                        data[4..].copy_from_slice(&addr);
                        out.write_all(&data).map(|_| IsNull::No).map_err(Into::into)
                    }
                }
            }
        }
    };
}

impl_Sql!(Inet, 0);
impl_Sql!(Cidr, 1);

fn network_v4(net: &Ipv4Network) -> Ipv4Addr {
    let mask = u32::from(net.mask());
    let ip = u32::from(net.ip()) & mask;
    Ipv4Addr::from(ip)
}

fn network_v6(net: &Ipv6Network) -> Ipv6Addr {
    let mask = u128::from(net.mask());
    let network = u128::from(net.ip()) & mask;
    Ipv6Addr::from(network)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_helpers::format_error;

    #[diesel_test_helper::test]
    fn v4address_to_sql() {
        macro_rules! test_to_sql {
            ($ty:ty, $net_type:expr) => {
                let mut buffer = Vec::new();
                {
                    let mut bytes = Output::test(ByteWrapper(&mut buffer));
                    let test_address =
                        IpNetwork::V4(Ipv4Network::new(Ipv4Addr::new(127, 0, 0, 1), 32).unwrap());
                    ToSql::<$ty, Pg>::to_sql(&test_address, &mut bytes).unwrap();
                }
                assert_eq!(buffer, vec![PGSQL_AF_INET, 32, $net_type, 4, 127, 0, 0, 1]);
            };
        }

        test_to_sql!(Inet, 0);
        test_to_sql!(Cidr, 1);
    }

    #[diesel_test_helper::test]
    fn some_v4address_from_sql() {
        macro_rules! test_some_address_from_sql {
            ($ty:tt) => {
                let input_address =
                    IpNetwork::V4(Ipv4Network::new(Ipv4Addr::new(127, 0, 0, 1), 32).unwrap());
                let mut buffer = Vec::new();
                {
                    let mut bytes = Output::test(ByteWrapper(&mut buffer));
                    ToSql::<$ty, Pg>::to_sql(&input_address, &mut bytes).unwrap();
                }
                let output_address =
                    FromSql::<$ty, Pg>::from_sql(PgValue::for_test(&buffer)).unwrap();
                assert_eq!(input_address, output_address);
            };
        }

        test_some_address_from_sql!(Cidr);
        test_some_address_from_sql!(Inet);
    }

    #[diesel_test_helper::test]
    fn v6address_to_sql() {
        macro_rules! test_to_sql {
            ($ty:ty, $net_type:expr, $last:expr) => {
                let mut buffer = Vec::new();
                {
                    let mut bytes = Output::test(ByteWrapper(&mut buffer));
                    let test_address = IpNetwork::V6(
                        Ipv6Network::new(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1), 64).unwrap(),
                    );
                    ToSql::<$ty, Pg>::to_sql(&test_address, &mut bytes).unwrap();
                }
                assert_eq!(
                    buffer,
                    vec![
                        PGSQL_AF_INET6,
                        64,
                        $net_type,
                        16,
                        0,
                        0,
                        0,
                        0,
                        0,
                        0,
                        0,
                        0,
                        0,
                        0,
                        0,
                        0,
                        0,
                        0,
                        0,
                        $last,
                    ]
                );
            };
        }

        // `::1/64` keeps its host bit as an `Inet` and loses it as a `Cidr`
        test_to_sql!(Inet, 0, 1);
        test_to_sql!(Cidr, 1, 0);
    }

    #[diesel_test_helper::test]
    fn some_v6address_from_sql() {
        macro_rules! test_some_address_from_sql {
            ($ty:tt, $last:expr) => {
                let input_address = IpNetwork::V6(
                    Ipv6Network::new(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, $last), 64).unwrap(),
                );
                let mut buffer = Vec::new();
                {
                    let mut bytes = Output::test(ByteWrapper(&mut buffer));
                    ToSql::<$ty, Pg>::to_sql(&input_address, &mut bytes).unwrap();
                }
                let output_address =
                    FromSql::<$ty, Pg>::from_sql(PgValue::for_test(&buffer)).unwrap();
                assert_eq!(input_address, output_address);
            };
        }

        test_some_address_from_sql!(Inet, 1);
        test_some_address_from_sql!(Cidr, 0);
    }

    #[diesel_test_helper::test]
    fn bad_address_from_sql() {
        macro_rules! bad_address_from_sql {
            ($ty:tt) => {
                let address: Result<IpNetwork, _> =
                    FromSql::<$ty, Pg>::from_sql(PgValue::for_test(&[7, PGSQL_AF_INET, 0]));
                assert_eq!(
                    format_error(&*address.unwrap_err()),
                    "invalid network address format. input is too short."
                );
            };
        }

        bad_address_from_sql!(Inet);
        bad_address_from_sql!(Cidr);
    }

    #[diesel_test_helper::test]
    fn no_address_from_sql() {
        macro_rules! test_no_address_from_sql {
            ($ty:ty) => {
                let address: Result<IpNetwork, _> = FromSql::<$ty, Pg>::from_nullable_sql(None);
                assert_eq!(
                    format_error(&*address.unwrap_err()),
                    "Unexpected null for non-null column"
                );
            };
        }

        test_no_address_from_sql!(Inet);
        test_no_address_from_sql!(Cidr);
    }
}