Line data    Source code

       1              : use std::env::{VarError, var};
       2              : use std::error::Error;
       3              : use std::net::IpAddr;
       4              : use std::str::FromStr;
       5              : 
       6              : /// Name of the environment variable containing the reachable IP address of the node. If set, the IP address contained in this
       7              : /// environment variable is used as the reachable IP address of the pageserver or safekeeper node during node registration.
       8              : /// In a Kubernetes environment, this environment variable should be set by Kubernetes to the Pod IP (specified in the Pod
       9              : /// template).
      10              : pub const HADRON_NODE_IP_ADDRESS: &str = "HADRON_NODE_IP_ADDRESS";
      11              : 
      12              : /// Read the reachable IP address of this page server from env var HADRON_NODE_IP_ADDRESS.
      13              : /// In Kubernetes this environment variable is set to the Pod IP (specified in the Pod template).
      14            4 : pub fn read_node_ip_addr_from_env() -> Result<Option<IpAddr>, Box<dyn Error>> {
      15            4 :     match var(HADRON_NODE_IP_ADDRESS) {
      16            3 :         Ok(v) => {
      17            3 :             if let Ok(addr) = IpAddr::from_str(&v) {
      18            2 :                 Ok(Some(addr))
      19              :             } else {
      20            1 :                 Err(format!("Invalid IP address string: {v}. Cannot be parsed as either an IPv4 or an IPv6 address.").into())
      21              :             }
      22              :         }
      23            1 :         Err(VarError::NotPresent) => Ok(None),
      24            0 :         Err(e) => Err(e.into()),
      25              :     }
      26            4 : }
      27              : 
      28              : #[cfg(test)]
      29              : mod tests {
      30              :     use super::*;
      31              :     use std::env;
      32              :     use std::net::{Ipv4Addr, Ipv6Addr};
      33              : 
      34              :     #[test]
      35            1 :     fn test_read_node_ip_addr_from_env() {
      36              :         // SAFETY: test code
      37              :         unsafe {
      38              :             // Test with a valid IPv4 address
      39            1 :             env::set_var(HADRON_NODE_IP_ADDRESS, "192.168.1.1");
      40            1 :             let result = read_node_ip_addr_from_env().unwrap();
      41            1 :             assert_eq!(result, Some(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1))));
      42              : 
      43              :             // Test with a valid IPv6 address
      44            1 :             env::set_var(
      45              :                 HADRON_NODE_IP_ADDRESS,
      46              :                 "2001:0db8:85a3:0000:0000:8a2e:0370:7334",
      47              :             );
      48              :         }
      49            1 :         let result = read_node_ip_addr_from_env().unwrap();
      50            1 :         assert_eq!(
      51              :             result,
      52            1 :             Some(IpAddr::V6(
      53            1 :                 Ipv6Addr::from_str("2001:0db8:85a3:0000:0000:8a2e:0370:7334").unwrap()
      54            1 :             ))
      55              :         );
      56              : 
      57              :         // Test with an invalid IP address
      58              :         // SAFETY: test code
      59            1 :         unsafe {
      60            1 :             env::set_var(HADRON_NODE_IP_ADDRESS, "invalid_ip");
      61            1 :         }
      62            1 :         let result = read_node_ip_addr_from_env();
      63            1 :         assert!(result.is_err());
      64              : 
      65              :         // Test with no environment variable set
      66              :         // SAFETY: test code
      67            1 :         unsafe {
      68            1 :             env::remove_var(HADRON_NODE_IP_ADDRESS);
      69            1 :         }
      70            1 :         let result = read_node_ip_addr_from_env().unwrap();
      71            1 :         assert_eq!(result, None);
      72            1 :     }
      73              : }