Skip to content

Commit 51d6785

Browse files
committed
implement socket.inet_pton, socket.inet_ntop
1 parent c465662 commit 51d6785

File tree

1 file changed

+67
-1
lines changed

1 file changed

+67
-1
lines changed

vm/src/stdlib/socket.rs

Lines changed: 67 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use crate::common::cell::{PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard};
22
use std::io::{self, prelude::*};
3-
use std::net::{Ipv4Addr, Shutdown, SocketAddr, ToSocketAddrs};
3+
use std::net::{Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr, ToSocketAddrs};
44
use std::time::Duration;
55

66
use byteorder::{BigEndian, ByteOrder};
@@ -596,6 +596,56 @@ fn socket_gethostbyname(name: PyStringRef, vm: &VirtualMachine) -> PyResult<Stri
596596
}
597597
}
598598

599+
fn socket_inet_pton(af_inet: i32, ip_string: PyStringRef, vm: &VirtualMachine) -> PyResult {
600+
match af_inet {
601+
c::AF_INET => ip_string
602+
.borrow_value()
603+
.parse::<Ipv4Addr>()
604+
.map(|ip_addr| vm.ctx.new_bytes(ip_addr.octets().to_vec()))
605+
.map_err(|_| {
606+
vm.new_os_error("illegal IP address string passed to inet_pton".to_owned())
607+
}),
608+
c::AF_INET6 => ip_string
609+
.borrow_value()
610+
.parse::<Ipv6Addr>()
611+
.map(|ip_addr| vm.ctx.new_bytes(ip_addr.octets().to_vec()))
612+
.map_err(|_| {
613+
vm.new_os_error("illegal IP address string passed to inet_pton".to_owned())
614+
}),
615+
_ => Err(vm.new_os_error("Address family not supported by protocol".to_owned())),
616+
}
617+
}
618+
619+
fn socket_inet_ntop(af_inet: i32, packed_ip: PyBytesRef, vm: &VirtualMachine) -> PyResult {
620+
match af_inet {
621+
c::AF_INET => {
622+
if packed_ip.len() != 4 {
623+
return Err(
624+
vm.new_value_error("invalid length of packed IP address string".to_owned())
625+
);
626+
}
627+
let ip_num = BigEndian::read_u32(&packed_ip);
628+
Ok(vm.ctx.new_str(Ipv4Addr::from(ip_num).to_string()))
629+
}
630+
c::AF_INET6 => {
631+
if packed_ip.len() != 16 {
632+
return Err(
633+
vm.new_value_error("invalid length of packed IP address string".to_owned())
634+
);
635+
}
636+
let mut iter = packed_ip.iter();
637+
let mut ip_num: u128 = *iter.next().unwrap() as u128;
638+
for item in iter {
639+
ip_num <<= 8;
640+
ip_num |= *item as u128
641+
}
642+
let ipv6 = Ipv6Addr::from(ip_num);
643+
Ok(vm.ctx.new_str(get_ipv6_addr_str(ipv6)))
644+
}
645+
_ => Err(vm.new_value_error(format!("unknown address family {}", af_inet))),
646+
}
647+
}
648+
599649
fn get_addr<T, I>(vm: &VirtualMachine, addr: T) -> PyResult<socket2::SockAddr>
600650
where
601651
T: ToSocketAddrs<Iter = I>,
@@ -667,6 +717,20 @@ fn convert_sock_error(vm: &VirtualMachine, err: io::Error) -> PyBaseExceptionRef
667717
}
668718
}
669719

720+
fn get_ipv6_addr_str(ipv6: Ipv6Addr) -> String {
721+
match ipv6.to_ipv4() {
722+
Some(_) => {
723+
let segments = ipv6.segments();
724+
if segments[5] == 0 && segments[6] == 0 {
725+
format!("::{:x}", segments[7] as u32)
726+
} else {
727+
ipv6.to_string()
728+
}
729+
}
730+
None => ipv6.to_string(),
731+
}
732+
}
733+
670734
pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
671735
let ctx = &vm.ctx;
672736
let socket_timeout = ctx.new_class(
@@ -694,6 +758,8 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
694758
"ntohs" => ctx.new_function(u16::from_be),
695759
"getdefaulttimeout" => ctx.new_function(|vm: &VirtualMachine| vm.get_none()),
696760
"has_ipv6" => ctx.new_bool(false),
761+
"inet_pton" => ctx.new_function(socket_inet_pton),
762+
"inet_ntop" => ctx.new_function(socket_inet_ntop),
697763
// constants
698764
"AF_UNSPEC" => ctx.new_int(0),
699765
"AF_INET" => ctx.new_int(c::AF_INET),

0 commit comments

Comments
 (0)