dynamo_runtime/utils/
ip_resolver.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! IP resolution utilities for getting local IP addresses with fallback support
5
6use crate::pipeline::network::tcp::server::{DefaultIpResolver, IpResolver};
7use local_ip_address::Error;
8use std::net::IpAddr;
9
10/// Get the local IP address for HTTP RPC host binding, using IpResolver with fallback to 127.0.0.1
11///
12/// This function attempts to resolve the local IP address using the provided resolver.
13/// If resolution fails, it falls back to 127.0.0.1 (localhost).
14///
15/// # Arguments
16/// * `resolver` - An implementation of IpResolver trait for getting local IP addresses
17///
18/// # Returns
19/// A string representation of the resolved IP address
20pub fn get_http_rpc_host_with_resolver<R: IpResolver>(resolver: R) -> String {
21    let resolved_ip = resolver.local_ip().or_else(|err| match err {
22        Error::LocalIpAddressNotFound => resolver.local_ipv6(),
23        _ => Err(err),
24    });
25
26    match resolved_ip {
27        Ok(addr) => addr,
28        Err(Error::LocalIpAddressNotFound) => IpAddr::from([127, 0, 0, 1]),
29        Err(_) => IpAddr::from([127, 0, 0, 1]), // Fallback for any other error
30    }
31    .to_string()
32}
33
34/// Get the local IP address for HTTP RPC host binding using the default resolver
35///
36/// This is a convenience function that uses the DefaultIpResolver.
37/// It follows the same logic as the TcpStreamServer for IP resolution.
38///
39/// # Returns
40/// A string representation of the resolved IP address, with fallback to "127.0.0.1"
41pub fn get_http_rpc_host() -> String {
42    get_http_rpc_host_with_resolver(DefaultIpResolver)
43}
44
45/// Get the HTTP RPC host from environment variable or resolve local IP as fallback
46///
47/// This function checks the DYN_HTTP_RPC_HOST environment variable first.
48/// If not set, it uses IP resolution to determine the local IP address.
49///
50/// # Returns
51/// A string representation of the HTTP RPC host address
52pub fn get_http_rpc_host_from_env() -> String {
53    std::env::var("DYN_HTTP_RPC_HOST").unwrap_or_else(|_| get_http_rpc_host())
54}
55
56/// Get the TCP RPC host from environment variable or resolve local IP as fallback
57///
58/// This function checks the DYN_TCP_RPC_HOST environment variable first.
59/// If not set, it uses IP resolution to determine the local IP address.
60///
61/// # Returns
62/// A string representation of the TCP RPC host address
63pub fn get_tcp_rpc_host_from_env() -> String {
64    std::env::var("DYN_TCP_RPC_HOST").unwrap_or_else(|_| get_http_rpc_host())
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use local_ip_address::Error;
71
72    // Mock resolver for testing
73    struct MockIpResolver {
74        ipv4_result: Result<IpAddr, Error>,
75        ipv6_result: Result<IpAddr, Error>,
76    }
77
78    impl IpResolver for MockIpResolver {
79        fn local_ip(&self) -> Result<IpAddr, Error> {
80            match &self.ipv4_result {
81                Ok(addr) => Ok(*addr),
82                Err(Error::LocalIpAddressNotFound) => Err(Error::LocalIpAddressNotFound),
83                Err(_) => Err(Error::LocalIpAddressNotFound), // Simplify for testing
84            }
85        }
86
87        fn local_ipv6(&self) -> Result<IpAddr, Error> {
88            match &self.ipv6_result {
89                Ok(addr) => Ok(*addr),
90                Err(Error::LocalIpAddressNotFound) => Err(Error::LocalIpAddressNotFound),
91                Err(_) => Err(Error::LocalIpAddressNotFound), // Simplify for testing
92            }
93        }
94    }
95
96    #[test]
97    fn test_get_http_rpc_host_with_successful_ipv4() {
98        let resolver = MockIpResolver {
99            ipv4_result: Ok(IpAddr::from([192, 168, 1, 100])),
100            ipv6_result: Ok(IpAddr::from([0, 0, 0, 0, 0, 0, 0, 1])),
101        };
102
103        let result = get_http_rpc_host_with_resolver(resolver);
104        assert_eq!(result, "192.168.1.100");
105    }
106
107    #[test]
108    fn test_get_http_rpc_host_with_ipv4_fail_ipv6_success() {
109        let resolver = MockIpResolver {
110            ipv4_result: Err(Error::LocalIpAddressNotFound),
111            ipv6_result: Ok(IpAddr::from([0x2001, 0xdb8, 0, 0, 0, 0, 0, 1])),
112        };
113
114        let result = get_http_rpc_host_with_resolver(resolver);
115        assert_eq!(result, "2001:db8::1");
116    }
117
118    #[test]
119    fn test_get_http_rpc_host_with_both_fail() {
120        let resolver = MockIpResolver {
121            ipv4_result: Err(Error::LocalIpAddressNotFound),
122            ipv6_result: Err(Error::LocalIpAddressNotFound),
123        };
124
125        let result = get_http_rpc_host_with_resolver(resolver);
126        assert_eq!(result, "127.0.0.1");
127    }
128
129    #[test]
130    fn test_get_http_rpc_host_from_env_with_env_var() {
131        // Set environment variable
132        unsafe {
133            std::env::set_var("DYN_HTTP_RPC_HOST", "10.0.0.1");
134        }
135
136        let result = get_http_rpc_host_from_env();
137        assert_eq!(result, "10.0.0.1");
138
139        // Clean up
140        unsafe {
141            std::env::remove_var("DYN_HTTP_RPC_HOST");
142        }
143    }
144
145    #[test]
146    fn test_get_http_rpc_host_from_env_without_env_var() {
147        // Note: We can't reliably unset environment variables in tests
148        // This test assumes DYN_HTTP_RPC_HOST is not set to a specific test value
149
150        let result = get_http_rpc_host_from_env();
151        // Should return some IP address (either resolved or fallback)
152        assert!(!result.is_empty());
153
154        // Should be a valid IP address
155        let _: IpAddr = result.parse().expect("Should be a valid IP address");
156    }
157}