0g-storage-node/node/network/src/nat.rs
2024-01-04 18:05:32 +08:00

232 lines
9.1 KiB
Rust

//! This houses various NAT hole punching strategies.
//!
//! Currently supported strategies:
//! - UPnP
use crate::{NetworkConfig, NetworkMessage};
use if_addrs::get_if_addrs;
use std::net::{IpAddr, SocketAddr, SocketAddrV4};
use tokio::sync::mpsc;
/// Configuration required to construct the UPnP port mappings.
pub struct UPnPConfig {
/// The local tcp port.
tcp_port: u16,
/// The local udp port.
udp_port: u16,
/// Whether discovery is enabled or not.
disable_discovery: bool,
}
impl UPnPConfig {
pub fn from_config(config: &NetworkConfig) -> Option<Self> {
if config.enr_tcp_port.is_some()
|| config.enr_udp_port.is_some()
|| config.enr_address.is_some()
{
None
} else {
Some(UPnPConfig {
tcp_port: config.libp2p_port,
udp_port: config.discovery_port,
disable_discovery: config.disable_discovery,
})
}
}
}
/// Attempts to construct external port mappings with UPnP.
pub fn construct_upnp_mappings(
config: UPnPConfig,
network_send: mpsc::UnboundedSender<NetworkMessage>,
) {
info!("UPnP Attempting to initialise routes");
match igd::search_gateway(Default::default()) {
Err(e) => info!(error = %e, "UPnP not available"),
Ok(gateway) => {
// Need to find the local listening address matched with the router subnet
let interfaces = match get_if_addrs() {
Ok(v) => v,
Err(e) => {
info!(error = %e, "UPnP failed to get local interfaces");
return;
}
};
let local_ip = interfaces.iter().find_map(|interface| {
// Just use the first IP of the first interface that is not a loopback and not an
// ipv6 address.
if !interface.is_loopback() {
interface.ip().is_ipv4().then(|| interface.ip())
} else {
None
}
});
let local_ip = match local_ip {
None => {
info!("UPnP failed to find local IP address");
return;
}
Some(v) => v,
};
debug!(ip = %local_ip, "UPnP Local IP Discovered");
match local_ip {
IpAddr::V4(address) => {
let libp2p_socket = SocketAddrV4::new(address, config.tcp_port);
let external_ip = gateway.get_external_ip();
// We add specific port mappings rather than getting the router to arbitrary assign
// one.
// I've found this to be more reliable. If multiple users are behind a single
// router, they should ideally try to set different port numbers.
let tcp_socket = add_port_mapping(
&gateway,
igd::PortMappingProtocol::TCP,
libp2p_socket,
"tcp",
)
.and_then(|_| {
let external_socket = external_ip
.as_ref()
.map(|ip| SocketAddr::new((*ip).into(), config.tcp_port))
.map_err(|_| ());
info!(
external_socket = format!(
"{}:{}",
external_socket
.as_ref()
.map(|ip| ip.to_string())
.unwrap_or_else(|_| "".into()),
config.tcp_port
)
);
external_socket
})
.ok();
let udp_socket = if !config.disable_discovery {
let discovery_socket = SocketAddrV4::new(address, config.udp_port);
add_port_mapping(
&gateway,
igd::PortMappingProtocol::UDP,
discovery_socket,
"udp",
)
.and_then(|_| {
let external_socket = external_ip
.map(|ip| SocketAddr::new(ip.into(), config.udp_port))
.map_err(|_| ());
info!(
external_socket = format!(
"{}:{}",
external_socket
.as_ref()
.map(|ip| ip.to_string())
.unwrap_or_else(|_| "".into()),
config.udp_port
),
"UPnP UDP route established"
);
external_socket
})
.ok()
} else {
None
};
// report any updates to the network service.
network_send
.send(NetworkMessage::UPnPMappingEstablished {
tcp_socket,
udp_socket,
})
.unwrap_or_else(
|e| debug!(error = %e, "Could not send message to the network service"),
);
}
_ => debug!("UPnP no routes constructed. IPv6 not supported"),
}
}
};
}
/// Sets up a port mapping for a protocol returning the mapped port if successful.
fn add_port_mapping(
gateway: &igd::Gateway,
protocol: igd::PortMappingProtocol,
socket: SocketAddrV4,
protocol_string: &'static str,
) -> Result<(), ()> {
// We add specific port mappings rather than getting the router to arbitrary assign
// one.
// I've found this to be more reliable. If multiple users are behind a single
// router, they should ideally try to set different port numbers.
let mapping_string = &format!("lighthouse-{}", protocol_string);
for _ in 0..2 {
match gateway.add_port(protocol, socket.port(), socket, 0, mapping_string) {
Err(e) => {
match e {
igd::AddPortError::PortInUse => {
// Try and remove and re-create
debug!(
protocol = protocol_string,
port = socket.port(),
"UPnP port in use, attempting to remap"
);
match gateway.remove_port(protocol, socket.port()) {
Ok(()) => {
debug!(
protocol = protocol_string,
port = socket.port(),
"UPnP Removed port mapping"
)
}
Err(e) => {
debug!(protocol = protocol_string, port = socket.port(), error = %e, "UPnP Port remove failure");
return Err(());
}
}
}
e => {
info!(error = %e, "UPnP TCP route not set");
return Err(());
}
}
}
Ok(_) => {
return Ok(());
}
}
}
Err(())
}
/// Removes the specified TCP and UDP port mappings.
pub fn remove_mappings(tcp_port: Option<u16>, udp_port: Option<u16>) {
if tcp_port.is_some() || udp_port.is_some() {
debug!("Removing UPnP port mappings");
match igd::search_gateway(Default::default()) {
Ok(gateway) => {
if let Some(tcp_port) = tcp_port {
match gateway.remove_port(igd::PortMappingProtocol::TCP, tcp_port) {
Ok(()) => debug!(port = tcp_port, "UPnP Removed TCP port mapping"),
Err(e) => {
debug!(port = tcp_port, error = %e, "UPnP Failed to remove TCP port mapping")
}
}
}
if let Some(udp_port) = udp_port {
match gateway.remove_port(igd::PortMappingProtocol::UDP, udp_port) {
Ok(()) => debug!(port = tcp_port, "UPnP Removed UDP port mapping"),
Err(e) => {
debug!(port = tcp_port, error = %e, "UPnP Failed to remove UDP port mapping")
}
}
}
}
Err(e) => debug!(error = %e, "UPnP failed to remove mappings"),
}
}
}