diff --git a/src/header/netdb/lookup.rs b/src/header/netdb/lookup.rs index 0734eec6..ccb00b65 100644 --- a/src/header/netdb/lookup.rs +++ b/src/header/netdb/lookup.rs @@ -17,10 +17,11 @@ use crate::header::{ netinet_in::{IPPROTO_UDP, in_addr, sockaddr_in}, sys_socket::{ self, - constants::{AF_INET, SOCK_DGRAM}, + constants::{AF_INET, SOCK_DGRAM, SOL_SOCKET, SO_RCVTIMEO}, sockaddr, }, time, + sys_select::timeval, }; use super::{ @@ -89,11 +90,34 @@ pub fn lookup_host(host: &str) -> Result { drop(Box::from_raw(packet_data_ptr)); } - let i = 0 as socklen_t; let mut buf = vec![0u8; 65536]; let buf_ptr = buf.as_mut_ptr().cast::(); - let count = unsafe { sys_socket::recv(sock, buf_ptr, 65536, 0) }; + // Set 5s recv timeout (best-effort; if this fails, recv may block longer). + let tv = timeval { + tv_sec: 5, + tv_usec: 0, + }; + unsafe { + sys_socket::setsockopt( + sock, + SOL_SOCKET, + SO_RCVTIMEO, + &tv as *const timeval as *const c_void, + core::mem::size_of::() as socklen_t, + ); + } + + let mut count: isize = -1; + for _attempt in 0..2 { + count = unsafe { sys_socket::recv(sock, buf_ptr, 65536, 0) }; + if count >= 0 { + break; + } + if unsafe { sys_socket::send(sock, packet_data_ptr, packet_data_len, 0) } < 0 { + break; + } + } if count < 0 { return Err(EIO); } @@ -193,11 +217,34 @@ pub fn lookup_addr(addr: in_addr) -> Result>, c_int> { drop(Box::from_raw(packet_data_ptr)); } - let i = mem::size_of::() as socklen_t; let mut buf = [0u8; 65536]; let buf_ptr = buf.as_mut_ptr().cast::(); - let count = unsafe { sys_socket::recv(sock, buf_ptr, 65536, 0) }; + // Set 5s recv timeout (best-effort; if this fails, recv may block longer). + let tv = timeval { + tv_sec: 5, + tv_usec: 0, + }; + unsafe { + sys_socket::setsockopt( + sock, + SOL_SOCKET, + SO_RCVTIMEO, + &tv as *const timeval as *const c_void, + core::mem::size_of::() as socklen_t, + ); + } + + let mut count: isize = -1; + for _attempt in 0..2 { + count = unsafe { sys_socket::recv(sock, buf_ptr, 65536, 0) }; + if count >= 0 { + break; + } + if unsafe { sys_socket::send(sock, packet_data_ptr, packet_data_len, 0) } < 0 { + break; + } + } if count < 0 { return Err(EIO); }