Skip to content

Commit

Permalink
fix(igd): stop renewing port mapping after endpoint.close() is called
Browse files Browse the repository at this point in the history
  • Loading branch information
lionel-faber authored and dirvine committed Jul 1, 2021
1 parent 53de7a9 commit 949b372
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 3 deletions.
13 changes: 11 additions & 2 deletions src/endpoint.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ use super::{
use bytes::Bytes;
use log::{debug, error, info, trace, warn};
use std::{net::SocketAddr, time::Duration};
use tokio::sync::broadcast::{self, Receiver, Sender};
use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
use tokio::time::timeout;

Expand Down Expand Up @@ -81,6 +82,7 @@ pub struct Endpoint {
client_cfg: quinn::ClientConfig,
bootstrap_nodes: Vec<SocketAddr>,
qp2p_config: Config,
termination_tx: Sender<()>,
connection_pool: ConnectionPool,
connection_deduplicator: ConnectionDeduplicator,
}
Expand Down Expand Up @@ -118,6 +120,7 @@ impl Endpoint {
let (message_tx, message_rx) = mpsc::unbounded_channel();
let (connection_tx, connection_rx) = mpsc::unbounded_channel();
let (disconnection_tx, disconnection_rx) = mpsc::unbounded_channel();
let (termination_tx, termination_rx) = broadcast::channel(1);

let mut endpoint = Self {
local_addr,
Expand All @@ -128,6 +131,7 @@ impl Endpoint {
client_cfg,
bootstrap_nodes,
qp2p_config,
termination_tx,
connection_pool: connection_pool.clone(),
connection_deduplicator: ConnectionDeduplicator::new(),
};
Expand Down Expand Up @@ -182,7 +186,7 @@ impl Endpoint {
warn!("Public IP address not verified since bootstrap contacts are empty");
}
} else {
endpoint.public_addr = Some(endpoint.fetch_public_address().await?);
endpoint.public_addr = Some(endpoint.fetch_public_address(termination_rx).await?);
}

listen_for_incoming_connections(
Expand Down Expand Up @@ -219,7 +223,10 @@ impl Endpoint {
/// simply build our connection info by querying the underlying bound socket for our address.
/// Note that if such an obtained address is of unspecified category we will ignore that as
/// such an address cannot be reached and hence not useful.
async fn fetch_public_address(&mut self) -> Result<SocketAddr> {
async fn fetch_public_address(
&mut self,
#[allow(unused)] termination_rx: Receiver<()>,
) -> Result<SocketAddr> {
// Skip port forwarding
if self.local_addr.ip().is_loopback() || !self.qp2p_config.forward_port {
self.public_addr = Some(self.local_addr);
Expand Down Expand Up @@ -258,6 +265,7 @@ impl Endpoint {
self.qp2p_config
.upnp_lease_duration
.unwrap_or(DEFAULT_UPNP_LEASE_DURATION_SEC),
termination_rx,
),
)
.await
Expand Down Expand Up @@ -493,6 +501,7 @@ impl Endpoint {

/// Close all the connections of this endpoint immediately and stop accepting new connections.
pub fn close(&self) {
let _ = self.termination_tx.send(());
self.quic_endpoint.close(0_u32.into(), b"")
}

Expand Down
10 changes: 9 additions & 1 deletion src/igd.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,15 @@ use igd::SearchOptions;
use log::{debug, info, warn};
use std::net::SocketAddr;
use std::time::Duration;
use tokio::sync::broadcast::{error::TryRecvError, Receiver};
use tokio::time::{self, Instant};

/// Automatically forwards a port and setups a tokio task to renew it periodically.
pub async fn forward_port(local_addr: SocketAddr, lease_duration: u32) -> Result<u16> {
pub async fn forward_port(
local_addr: SocketAddr,
lease_duration: u32,
mut termination_rx: Receiver<()>,
) -> Result<u16> {
let igd_res = add_port(local_addr, lease_duration).await;

if let Ok(ext_port) = &igd_res {
Expand All @@ -27,6 +32,9 @@ pub async fn forward_port(local_addr: SocketAddr, lease_duration: u32) -> Result

loop {
let _ = timer.tick().await;
if let Err(TryRecvError::Empty) = termination_rx.try_recv() {
break;
}
debug!("Renewing IGD lease for port {}", local_addr);

let renew_res = renew_port(local_addr, ext_port, lease_duration).await;
Expand Down

0 comments on commit 949b372

Please sign in to comment.