Skip to content

Commit c1eeb8a

Browse files
rvolosatovsrvolosatovs
committed
feat(wasi-sockets): implement UDP
This is based on TCP implementation Signed-off-by: Roman Volosatovs <rvolosatovs@riseup.net>
1 parent 542c180 commit c1eeb8a

14 files changed

Lines changed: 763 additions & 5 deletions

File tree

‎crates/test-programs/tests/wasi-sockets.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,16 @@ async fn tcp_v6() {
7676
run("tcp_v6").await.unwrap();
7777
}
7878

79+
#[test_log::test(tokio::test(flavor = "multi_thread"))]
80+
async fn udp_v4() {
81+
run("udp_v4").await.unwrap();
82+
}
83+
84+
#[test_log::test(tokio::test(flavor = "multi_thread"))]
85+
async fn udp_v6() {
86+
run("udp_v6").await.unwrap();
87+
}
88+
7989
#[test_log::test(tokio::test(flavor = "multi_thread"))]
8090
async fn ip_name_lookup() {
8191
run("ip_name_lookup").await.unwrap();

‎crates/test-programs/wasi-sockets-tests/src/bin/tcp_v4.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,5 +24,5 @@ fn main() {
2424

2525
sock.finish_bind().unwrap();
2626

27-
example_body(net, sock, IpAddressFamily::Ipv4)
27+
example_body_tcp(net, sock, IpAddressFamily::Ipv4)
2828
}

‎crates/test-programs/wasi-sockets-tests/src/bin/tcp_v6.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
//! Like v4.rs, but with IPv6.
1+
//! Like tcp_v4.rs, but with IPv6.
22
33
use wasi::io::poll;
44
use wasi::sockets::network::{IpAddressFamily, IpSocketAddress, Ipv6SocketAddress};
@@ -26,5 +26,5 @@ fn main() {
2626

2727
sock.finish_bind().unwrap();
2828

29-
example_body(net, sock, IpAddressFamily::Ipv6)
29+
example_body_tcp(net, sock, IpAddressFamily::Ipv6)
3030
}

‎crates/test-programs/wasi-sockets-tests/src/bin/udp_v4.rs‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
//! A simple UDP testcase, using IPv4.
2+
3+
use wasi::io::poll;
4+
use wasi::sockets::network::{IpAddressFamily, IpSocketAddress, Ipv4SocketAddress};
5+
use wasi::sockets::{instance_network, udp_create_socket};
6+
use wasi_sockets_tests::*;
7+
8+
fn main() {
9+
let net = instance_network::instance_network();
10+
11+
let sock = udp_create_socket::create_udp_socket(IpAddressFamily::Ipv4).unwrap();
12+
13+
let addr = IpSocketAddress::Ipv4(Ipv4SocketAddress {
14+
port: 0, // use any free port
15+
address: (127, 0, 0, 1), // localhost
16+
});
17+
18+
let sub = sock.subscribe();
19+
20+
sock.start_bind(&net, addr).unwrap();
21+
22+
poll::poll_one(&sub);
23+
drop(sub);
24+
25+
sock.finish_bind().unwrap();
26+
27+
example_body_udp(net, sock, IpAddressFamily::Ipv4)
28+
}

‎crates/test-programs/wasi-sockets-tests/src/bin/udp_v6.rs‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
//! Like udp_v4.rs, but with IPv6.
2+
3+
use wasi::io::poll;
4+
use wasi::sockets::network::{IpAddressFamily, IpSocketAddress, Ipv6SocketAddress};
5+
use wasi::sockets::{instance_network, udp_create_socket};
6+
use wasi_sockets_tests::*;
7+
8+
fn main() {
9+
let net = instance_network::instance_network();
10+
11+
let sock = udp_create_socket::create_udp_socket(IpAddressFamily::Ipv6).unwrap();
12+
13+
let addr = IpSocketAddress::Ipv6(Ipv6SocketAddress {
14+
port: 0, // use any free port
15+
address: (0, 0, 0, 0, 0, 0, 0, 1), // localhost
16+
flow_info: 0,
17+
scope_id: 0,
18+
});
19+
20+
let sub = sock.subscribe();
21+
22+
sock.start_bind(&net, addr).unwrap();
23+
24+
poll::poll_one(&sub);
25+
drop(sub);
26+
27+
sock.finish_bind().unwrap();
28+
29+
example_body_udp(net, sock, IpAddressFamily::Ipv6)
30+
}

‎crates/test-programs/wasi-sockets-tests/src/lib.rs‎

Lines changed: 89 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ wit_bindgen::generate!("test-command-with-sockets" in "../../wasi/wit");
22

33
use wasi::io::poll;
44
use wasi::io::streams;
5-
use wasi::sockets::{network, tcp, tcp_create_socket};
5+
use wasi::sockets::{network, tcp, tcp_create_socket, udp, udp_create_socket};
66

77
pub fn write(output: &streams::OutputStream, mut bytes: &[u8]) -> Result<(), streams::StreamError> {
88
let pollable = output.subscribe();
@@ -24,7 +24,7 @@ pub fn write(output: &streams::OutputStream, mut bytes: &[u8]) -> Result<(), str
2424
Ok(())
2525
}
2626

27-
pub fn example_body(net: tcp::Network, sock: tcp::TcpSocket, family: network::IpAddressFamily) {
27+
pub fn example_body_tcp(net: tcp::Network, sock: tcp::TcpSocket, family: network::IpAddressFamily) {
2828
let first_message = b"Hello, world!";
2929
let second_message = b"Greetings, planet!";
3030

@@ -95,3 +95,90 @@ pub fn example_body(net: tcp::Network, sock: tcp::TcpSocket, family: network::Ip
9595
// Check that we sent and recieved our message!
9696
assert_eq!(data, second_message); // Not guaranteed to work but should work in practice.
9797
}
98+
99+
pub fn example_body_udp(net: udp::Network, sock: udp::UdpSocket, family: network::IpAddressFamily) {
100+
let first_message = b"Hello, world!";
101+
let second_message = b"Greetings, planet!";
102+
103+
let sub = sock.subscribe();
104+
105+
let addr = sock.local_address().unwrap();
106+
107+
let client = udp_create_socket::create_udp_socket(family).unwrap();
108+
let client_sub = client.subscribe();
109+
110+
client.start_connect(&net, addr).unwrap();
111+
poll::poll_one(&client_sub);
112+
client.finish_connect().unwrap();
113+
114+
let _client_addr = client.local_address().unwrap();
115+
116+
let n = client
117+
.send(&[
118+
udp::Datagram {
119+
data: vec![],
120+
remote_address: addr,
121+
},
122+
udp::Datagram {
123+
data: first_message.to_vec(),
124+
remote_address: addr,
125+
},
126+
])
127+
.unwrap();
128+
assert_eq!(n, 2);
129+
130+
drop(client_sub);
131+
drop(client);
132+
133+
poll::poll_one(&sub);
134+
let datagrams = sock.receive(2).unwrap();
135+
let mut datagrams = datagrams.into_iter();
136+
let (first, second) = match (datagrams.next(), datagrams.next(), datagrams.next()) {
137+
(Some(first), Some(second), None) => (first, second),
138+
(Some(_first), None, None) => panic!("only one datagram received"),
139+
(None, None, None) => panic!("no datagrams received"),
140+
_ => panic!("invalid datagram sequence received"),
141+
};
142+
143+
assert!(first.data.is_empty());
144+
145+
// TODO: Verify the `remote_address`
146+
//assert_eq!(first.remote_address, client_addr);
147+
148+
// Check that we sent and recieved our message!
149+
assert_eq!(second.data, first_message); // Not guaranteed to work but should work in practice.
150+
151+
// TODO: Verify the `remote_address`
152+
//assert_eq!(second.remote_address, client_addr);
153+
154+
// Another client
155+
let client = udp_create_socket::create_udp_socket(family).unwrap();
156+
let client_sub = client.subscribe();
157+
158+
client.start_connect(&net, addr).unwrap();
159+
poll::poll_one(&client_sub);
160+
client.finish_connect().unwrap();
161+
162+
let n = client
163+
.send(&[udp::Datagram {
164+
data: second_message.to_vec(),
165+
remote_address: addr,
166+
}])
167+
.unwrap();
168+
assert_eq!(n, 1);
169+
170+
drop(client_sub);
171+
drop(client);
172+
173+
poll::poll_one(&sub);
174+
let datagrams = sock.receive(2).unwrap();
175+
let mut datagrams = datagrams.into_iter();
176+
let first = match (datagrams.next(), datagrams.next()) {
177+
(Some(first), None) => first,
178+
(None, None) => panic!("no datagrams received"),
179+
_ => panic!("invalid datagram sequence received"),
180+
};
181+
182+
// Check that we sent and recieved our message!
183+
assert_eq!(first.data, second_message); // Not guaranteed to work but should work in practice.
184+
}

‎crates/wasi-http/wit/test.wit‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ world test-command-with-sockets {
3737
import wasi:cli/stderr
3838
import wasi:sockets/tcp
3939
import wasi:sockets/tcp-create-socket
40+
import wasi:sockets/udp
41+
import wasi:sockets/udp-create-socket
4042
import wasi:sockets/network
4143
import wasi:sockets/instance-network
4244
import wasi:sockets/ip-name-lookup

‎crates/wasi/src/preview2/command.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@ pub fn add_to_linker<T: WasiView>(l: &mut wasmtime::component::Linker<T>) -> any
4848
crate::preview2::bindings::cli::terminal_stderr::add_to_linker(l, |t| t)?;
4949
crate::preview2::bindings::sockets::tcp::add_to_linker(l, |t| t)?;
5050
crate::preview2::bindings::sockets::tcp_create_socket::add_to_linker(l, |t| t)?;
51+
crate::preview2::bindings::sockets::udp::add_to_linker(l, |t| t)?;
52+
crate::preview2::bindings::sockets::udp_create_socket::add_to_linker(l, |t| t)?;
5153
crate::preview2::bindings::sockets::instance_network::add_to_linker(l, |t| t)?;
5254
crate::preview2::bindings::sockets::network::add_to_linker(l, |t| t)?;
5355
crate::preview2::bindings::sockets::ip_name_lookup::add_to_linker(l, |t| t)?;
@@ -65,6 +67,7 @@ pub mod sync {
6567
"wasi:filesystem/types": crate::preview2::bindings::sync_io::filesystem::types,
6668
"wasi:filesystem/preopens": crate::preview2::bindings::filesystem::preopens,
6769
"wasi:sockets/tcp": crate::preview2::bindings::sockets::tcp,
70+
"wasi:sockets/udp": crate::preview2::bindings::sockets::udp,
6871
"wasi:clocks/monotonic_clock": crate::preview2::bindings::clocks::monotonic_clock,
6972
"wasi:io/poll": crate::preview2::bindings::sync_io::io::poll,
7073
"wasi:io/streams": crate::preview2::bindings::sync_io::io::streams,
@@ -107,6 +110,8 @@ pub mod sync {
107110
crate::preview2::bindings::cli::terminal_stderr::add_to_linker(l, |t| t)?;
108111
crate::preview2::bindings::sockets::tcp::add_to_linker(l, |t| t)?;
109112
crate::preview2::bindings::sockets::tcp_create_socket::add_to_linker(l, |t| t)?;
113+
crate::preview2::bindings::sockets::udp::add_to_linker(l, |t| t)?;
114+
crate::preview2::bindings::sockets::udp_create_socket::add_to_linker(l, |t| t)?;
110115
crate::preview2::bindings::sockets::instance_network::add_to_linker(l, |t| t)?;
111116
crate::preview2::bindings::sockets::network::add_to_linker(l, |t| t)?;
112117
crate::preview2::bindings::sockets::ip_name_lookup::add_to_linker(l, |t| t)?;

‎crates/wasi/src/preview2/host/mod.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,3 +8,5 @@ mod network;
88
mod random;
99
mod tcp;
1010
mod tcp_create_socket;
11+
mod udp;
12+
mod udp_create_socket;

0 commit comments

Comments
 (0)