Skip to content

Commit

Permalink
Add getter for local_addr on TcpSocket (#1379)
Browse files Browse the repository at this point in the history
  • Loading branch information
zekisherif committed Oct 29, 2020
1 parent 183bbe4 commit b41a022
Show file tree
Hide file tree
Showing 5 changed files with 73 additions and 2 deletions.
7 changes: 7 additions & 0 deletions src/net/tcp/socket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,13 @@ impl TcpSocket {
pub fn set_linger(&self, dur: Option<Duration>) -> io::Result<()> {
sys::tcp::set_linger(self.sys, dur)
}

/// Returns the local address of this socket
///
/// Will return `Err` result in windows if called before calling `bind`
pub fn get_localaddr(&self) -> io::Result<SocketAddr> {
sys::tcp::get_localaddr(self.sys)
}
}

impl Drop for TcpSocket {
Expand Down
4 changes: 4 additions & 0 deletions src/sys/shell/tcp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -53,3 +53,7 @@ pub(crate) fn set_linger(_: TcpSocket, _: Option<Duration>) -> io::Result<()> {
pub fn accept(_: &net::TcpListener) -> io::Result<(net::TcpStream, SocketAddr)> {
os_required!();
}

pub(crate) fn get_localaddr(_: TcpSocket) -> io::Result<SocketAddr> {
os_required!();
}
13 changes: 13 additions & 0 deletions src/sys/unix/tcp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,19 @@ pub(crate) fn get_reuseport(socket: TcpSocket) -> io::Result<bool> {
Ok(optval != 0)
}

pub(crate) fn get_localaddr(socket: TcpSocket) -> io::Result<SocketAddr> {
let mut addr: libc::sockaddr_storage = unsafe { std::mem::zeroed() };
let mut length = size_of::<libc::sockaddr_storage>() as libc::socklen_t;

syscall!(getsockname(
socket,
&mut addr as *mut _ as *mut _,
&mut length
))?;

unsafe { to_socket_addr(&addr) }
}

pub(crate) fn set_linger(socket: TcpSocket, dur: Option<Duration>) -> io::Result<()> {
let val: libc::linger = libc::linger {
l_onoff: if dur.is_some() { 1 } else { 0 },
Expand Down
32 changes: 30 additions & 2 deletions src/sys/windows/tcp.rs
Original file line number Diff line number Diff line change
@@ -1,14 +1,17 @@
use std::io;
use std::mem::size_of;
use std::net::{self, SocketAddr};
use std::net::{self, SocketAddr, SocketAddrV4, SocketAddrV6};
use std::time::Duration;
use std::os::windows::io::FromRawSocket;
use std::os::windows::raw::SOCKET as StdSocket; // winapi uses usize, stdlib uses u32/u64.

use winapi::ctypes::{c_char, c_int, c_ushort};
use winapi::shared::ws2def::{SOCKADDR_STORAGE, AF_INET, SOCKADDR_IN};
use winapi::shared::ws2ipdef::SOCKADDR_IN6_LH;

use winapi::shared::minwindef::{BOOL, TRUE, FALSE};
use winapi::um::winsock2::{
self, closesocket, linger, setsockopt, getsockopt, PF_INET, PF_INET6, SOCKET, SOCKET_ERROR,
self, closesocket, linger, setsockopt, getsockopt, getsockname, PF_INET, PF_INET6, SOCKET, SOCKET_ERROR,
SOCK_STREAM, SOL_SOCKET, SO_LINGER, SO_REUSEADDR,
};

Expand Down Expand Up @@ -103,6 +106,31 @@ pub(crate) fn get_reuseaddr(socket: TcpSocket) -> io::Result<bool> {
}
}

pub(crate) fn get_localaddr(socket: TcpSocket) -> io::Result<SocketAddr> {
let mut addr: SOCKADDR_STORAGE = unsafe { std::mem::zeroed() };
let mut length = std::mem::size_of_val(&addr) as c_int;

match unsafe { getsockname(
socket,
&mut addr as *mut _ as *mut _,
&mut length
) } {
SOCKET_ERROR => Err(io::Error::last_os_error()),
_ => {
let storage: *const SOCKADDR_STORAGE = (&addr) as *const _;
if addr.ss_family as c_int == AF_INET {
let sock_addr : SocketAddrV4 = unsafe { *(storage as *const SOCKADDR_IN as *const _) };
Ok(sock_addr.into())
} else {
let sock_addr : SocketAddrV6 = unsafe { *(storage as *const SOCKADDR_IN6_LH as *const _) };
Ok(sock_addr.into())
}
},
}


}

pub(crate) fn set_linger(socket: TcpSocket, dur: Option<Duration>) -> io::Result<()> {
let val: linger = linger {
l_onoff: if dur.is_some() { 1 } else { 0 },
Expand Down
19 changes: 19 additions & 0 deletions tests/tcp_socket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,3 +37,22 @@ fn set_reuseport() {

let _ = socket.listen(128).unwrap();
}

#[test]
fn get_localaddr() {
let expected_addr = "127.0.0.1:0".parse().unwrap();
let socket = TcpSocket::new_v4().unwrap();

//Windows doesn't support calling getsockname before calling `bind`
#[cfg(not(windows))]
assert_eq!("0.0.0.0:0", socket.get_localaddr().unwrap().to_string());

socket.bind(expected_addr).unwrap();

let actual_addr = socket.get_localaddr().unwrap();

assert_eq!(expected_addr.ip(), actual_addr.ip());
assert!(actual_addr.port() > 0);

let _ = socket.listen(128).unwrap();
}

0 comments on commit b41a022

Please sign in to comment.