fixup: remove all heap allocations

Author: tankyleo

lightning-net-tokio/src/lib.rs

@@ -513,7 +513,13 @@ async fn socks5_connect(
 	const SOCKS5_REQUEST_REPLY_MAX_LEN: usize = 1 /* VER */ + 1 /* CMD for request, REP for reply */ + 1 /* RSV */
 		+ 1 /* ATYP */ + 1 /* HOSTNAME len */ + 255 /* HOSTNAME */ + 2 /* PORT */;
 	const SOCKS5_REQUEST_REPLY_MIN_LEN: usize = 1 /* VER */ + 1 /* CMD for request, REP for reply */ + 1 /* RSV */
-		+ 1 /* ATYP */ + 4 /* IPV4 ADDR */ + 2 /* PORT */;
+		+ 1 /* ATYP */ + 1 /* HOSTNAME len */ + 1 /* HOSTNAME */ + 2 /* PORT */;
+	const SOCKS5_IPV4_REQUEST_LEN: usize = 1 /* VER */ + 1 /* CMD */ + 1 /* RSV */ + 1 /* ATYP */
+		+ 4 /* IPV4 DST.ADDR */ + 2 /* DST.PORT */;
+	const SOCKS5_IPV6_REQUEST_LEN: usize = 1 /* VER */ + 1 /* CMD */ + 1 /* RSV */ + 1 /* ATYP */
+		+ 16 /* IPV4 DST.ADDR */ + 2 /* DST.PORT */;
+	const SOCKS5_HOSTNAME_REQUEST_STATIC_FIELDS_LEN: usize = 1 /* VER */ + 1 /* CMD */ + 1 /* RSV */ + 1 /* ATYP */
+		+ 1 /* HOSTNAME len */ + 2 /* DST.PORT */;
 
 	let method_selection_message = [VERSION, NMETHODS, NO_AUTH];
 	let mut tcp_stream = TcpStream::connect(&socks5_proxy_addr).await.map_err(|_| ())?;
@@ -525,48 +531,73 @@ async fn socks5_connect(
 		return Err(());
 	}
 
-	let mut socks5_request: Vec<u8> = Vec::with_capacity(SOCKS5_REQUEST_REPLY_MAX_LEN);
+	let mut socks5_request = [0u8; SOCKS5_REQUEST_REPLY_MAX_LEN];
+	let request_size;
 
-	socks5_request.push(VERSION);
-	socks5_request.push(CMD_CONNECT);
-	socks5_request.push(RSV);
+	socks5_request[0..3].copy_from_slice(&[VERSION, CMD_CONNECT, RSV]);
 
 	match addr {
 		SocketAddress::TcpIpV4 { addr, port } => {
-			socks5_request.push(ATYP_IPV4);
-			socks5_request.extend_from_slice(&addr);
-			socks5_request.extend_from_slice(&port.to_be_bytes());
+			socks5_request[3] = ATYP_IPV4;
+			socks5_request[4..8].copy_from_slice(&addr);
+			socks5_request[8..10].copy_from_slice(&port.to_be_bytes());
+			request_size = SOCKS5_IPV4_REQUEST_LEN;
 		},
 		SocketAddress::TcpIpV6 { addr, port } => {
-			socks5_request.push(ATYP_IPV6);
-			socks5_request.extend_from_slice(&addr);
-			socks5_request.extend_from_slice(&port.to_be_bytes());
+			socks5_request[3] = ATYP_IPV6;
+			socks5_request[4..20].copy_from_slice(&addr);
+			socks5_request[20..22].copy_from_slice(&port.to_be_bytes());
+			request_size = SOCKS5_IPV6_REQUEST_LEN;
 		},
 		ref onion_v3 @ SocketAddress::OnionV3 { port, .. } => {
 			let onion_v3_url = onion_v3.to_string();
 			let hostname = onion_v3_url.split_once(':').ok_or(())?.0.as_bytes();
-			socks5_request.push(ATYP_DOMAINNAME);
-			socks5_request.extend_from_slice(&[hostname.len() as u8]);
-			socks5_request.extend_from_slice(hostname);
-			socks5_request.extend_from_slice(&port.to_be_bytes());
+			socks5_request[3] = ATYP_DOMAINNAME;
+			socks5_request[4] = hostname.len() as u8;
+			let port_index = 5 + hostname.len();
+			socks5_request[5..port_index].copy_from_slice(hostname);
+			socks5_request[port_index..port_index + 2].copy_from_slice(&port.to_be_bytes());
+			request_size = SOCKS5_HOSTNAME_REQUEST_STATIC_FIELDS_LEN + hostname.len();
 		},
 		SocketAddress::Hostname { hostname, port } => {
-			socks5_request.push(ATYP_DOMAINNAME);
-			socks5_request.extend_from_slice(&[hostname.len() as u8]);
-			socks5_request.extend_from_slice(hostname.as_bytes());
-			socks5_request.extend_from_slice(&port.to_be_bytes());
+			socks5_request[3] = ATYP_DOMAINNAME;
+			socks5_request[4] = hostname.len();
+			let port_index = 5 + hostname.len() as usize;
+			socks5_request[5..port_index].copy_from_slice(hostname.as_bytes());
+			socks5_request[port_index..port_index + 2].copy_from_slice(&port.to_be_bytes());
+			request_size = SOCKS5_HOSTNAME_REQUEST_STATIC_FIELDS_LEN + hostname.len() as usize;
 		},
 		SocketAddress::OnionV2 { .. } => return Err(()),
 	};
 
-	tcp_stream.write_all(&socks5_request).await.map_err(|_| ())?;
+	tcp_stream.write_all(&socks5_request[..request_size]).await.map_err(|_| ())?;
+
+	let mut buffer = [0u8; SOCKS5_REQUEST_REPLY_MAX_LEN];
+	// First pull some bytes into the buffer
+	let mut total_read = tcp_stream.read(&mut buffer).await.map_err(|_| ())?;
+	// Then make sure we've reached EOF, otherwise keep appending until the max size of the buffer
+	while total_read < SOCKS5_REQUEST_REPLY_MAX_LEN {
+		let n_read = match tcp_stream.try_read(&mut buffer[total_read..]) {
+			Ok(n) => n,
+			Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => 0,
+			Err(_e) => return Err(()),
+		};
+		if n_read == 0 {
+			break;
+		}
+		total_read += n_read;
+	}
+	// If we've filled the buffer, return `Err` if we receive more bytes than the max allowed by the RFC
+	if total_read == SOCKS5_REQUEST_REPLY_MAX_LEN {
+		match tcp_stream.try_read(&mut [0u8; 1]) {
+			Ok(0) => (),
+			Ok(1..) => return Err(()),
+			Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => (),
+			Err(_e) => return Err(()),
+		}
+	}
 
-	let mut buffer: Vec<u8> = Vec::with_capacity(SOCKS5_REQUEST_REPLY_MAX_LEN);
-	let n_read = tcp_stream.read_to_end(&mut buffer).await.map_err(|_| ())?;
-	if n_read < SOCKS5_REQUEST_REPLY_MIN_LEN
-		|| n_read > SOCKS5_REQUEST_REPLY_MAX_LEN
-		|| buffer[..3] != [VERSION, SUCCEEDED, RSV]
-	{
+	if total_read < SOCKS5_REQUEST_REPLY_MIN_LEN || buffer[..3] != [VERSION, SUCCEEDED, RSV] {
 		return Err(());
 	}