diff --git a/Cargo.lock b/Cargo.lock index 9968db8..d1dfe03 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1179,7 +1179,7 @@ dependencies = [ [[package]] name = "omnip" -version = "0.6.2" +version = "0.6.3" dependencies = [ "android_logger", "anyhow", diff --git a/Cargo.toml b/Cargo.toml index a69d0cd..26da2c3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "omnip" -version = "0.6.2" +version = "0.6.3" edition = "2021" [lib] diff --git a/src/server.rs b/src/server.rs index 40941df..7783a9f 100644 --- a/src/server.rs +++ b/src/server.rs @@ -976,32 +976,23 @@ impl Server { Ok(0) } Ok(n) => { + // this method is used as a select! branch, use while loop to write the data + // instead of using write_all (which is not cancellation safe) to ensure + // the code cancellation safe + let mut written_bytes = 0; + while written_bytes < n { + written_bytes += writer + .write(&buffer[written_bytes..n]) + .await + .map_err(|_| ProxyError::InternalError)?; + } *out_bytes += n as u64; - writer - .write_all(&buffer[..n]) - .await - .map_err(|_| ProxyError::InternalError)?; Ok(n) } Err(_) => Err(ProxyError::InternalError), // Connection mostly reset by peer } } - // async fn resolve_net_addr(&self, addr: &NetAddr) -> Result { - // if addr.is_ip() { - // return Ok(addr.to_socket_addr().unwrap()); - // } - // - // let resolver = if addr.is_internal_domain() { - // inner_state!(self, system_dns_resolver).clone() - // } else { - // inner_state!(self, dns_resolver).clone() - // }; - // - // let ip_arr = resolver.unwrap().lookup(addr.unwrap_domain()).await?; - // Ok(SocketAddr::new(*ip_arr.first().unwrap(), addr.port)) - // } - fn collect_and_report_server_stats(&self, mut stats_receiver: Receiver) { let inner_state = self.inner_state.clone(); tokio::spawn(async move {