feat(wasi-sockets): implement UDP

This is based on TCP implementation

Signed-off-by: Roman Volosatovs <rvolosatovs@riseup.net>

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

@@ -76,6 +76,16 @@ async fn tcp_v6() {
     run("tcp_v6").await.unwrap();
 }
 
+#[test_log::test(tokio::test(flavor = "multi_thread"))]
+async fn udp_v4() {
+    run("udp_v4").await.unwrap();
+}
+
+#[test_log::test(tokio::test(flavor = "multi_thread"))]
+async fn udp_v6() {
+    run("udp_v6").await.unwrap();
+}
+
 #[test_log::test(tokio::test(flavor = "multi_thread"))]
 async fn ip_name_lookup() {
     run("ip_name_lookup").await.unwrap();

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

@@ -24,5 +24,5 @@ fn main() {
 
     sock.finish_bind().unwrap();
 
-    example_body(net, sock, IpAddressFamily::Ipv4)
+    example_body_tcp(net, sock, IpAddressFamily::Ipv4)
 }

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

@@ -1,4 +1,4 @@
-//! Like v4.rs, but with IPv6.
+//! Like tcp_v4.rs, but with IPv6.
 
 use wasi::io::poll;
 use wasi::sockets::network::{IpAddressFamily, IpSocketAddress, Ipv6SocketAddress};
@@ -26,5 +26,5 @@ fn main() {
 
     sock.finish_bind().unwrap();
 
-    example_body(net, sock, IpAddressFamily::Ipv6)
+    example_body_tcp(net, sock, IpAddressFamily::Ipv6)
 }

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

@@ -0,0 +1,28 @@
+//! A simple UDP testcase, using IPv4.
+
+use wasi::io::poll;
+use wasi::sockets::network::{IpAddressFamily, IpSocketAddress, Ipv4SocketAddress};
+use wasi::sockets::{instance_network, udp_create_socket};
+use wasi_sockets_tests::*;
+
+fn main() {
+    let net = instance_network::instance_network();
+
+    let sock = udp_create_socket::create_udp_socket(IpAddressFamily::Ipv4).unwrap();
+
+    let addr = IpSocketAddress::Ipv4(Ipv4SocketAddress {
+        port: 0,                 // use any free port
+        address: (127, 0, 0, 1), // localhost
+    });
+
+    let sub = sock.subscribe();
+
+    sock.start_bind(&net, addr).unwrap();
+
+    poll::poll_one(&sub);
+    drop(sub);
+
+    sock.finish_bind().unwrap();
+
+    example_body_udp(net, sock, IpAddressFamily::Ipv4)
+}

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

@@ -0,0 +1,30 @@
+//! Like udp_v4.rs, but with IPv6.
+
+use wasi::io::poll;
+use wasi::sockets::network::{IpAddressFamily, IpSocketAddress, Ipv6SocketAddress};
+use wasi::sockets::{instance_network, udp_create_socket};
+use wasi_sockets_tests::*;
+
+fn main() {
+    let net = instance_network::instance_network();
+
+    let sock = udp_create_socket::create_udp_socket(IpAddressFamily::Ipv6).unwrap();
+
+    let addr = IpSocketAddress::Ipv6(Ipv6SocketAddress {
+        port: 0,                           // use any free port
+        address: (0, 0, 0, 0, 0, 0, 0, 1), // localhost
+        flow_info: 0,
+        scope_id: 0,
+    });
+
+    let sub = sock.subscribe();
+
+    sock.start_bind(&net, addr).unwrap();
+
+    poll::poll_one(&sub);
+    drop(sub);
+
+    sock.finish_bind().unwrap();
+
+    example_body_udp(net, sock, IpAddressFamily::Ipv6)
+}

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

@@ -2,7 +2,7 @@ wit_bindgen::generate!("test-command-with-sockets" in "../../wasi/wit");
 
 use wasi::io::poll;
 use wasi::io::streams;
-use wasi::sockets::{network, tcp, tcp_create_socket};
+use wasi::sockets::{network, tcp, tcp_create_socket, udp, udp_create_socket};
 
 pub fn write(output: &streams::OutputStream, mut bytes: &[u8]) -> Result<(), streams::StreamError> {
     let pollable = output.subscribe();
@@ -24,7 +24,7 @@ pub fn write(output: &streams::OutputStream, mut bytes: &[u8]) -> Result<(), str
     Ok(())
 }
 
-pub fn example_body(net: tcp::Network, sock: tcp::TcpSocket, family: network::IpAddressFamily) {
+pub fn example_body_tcp(net: tcp::Network, sock: tcp::TcpSocket, family: network::IpAddressFamily) {
     let first_message = b"Hello, world!";
     let second_message = b"Greetings, planet!";
 
@@ -95,3 +95,90 @@ pub fn example_body(net: tcp::Network, sock: tcp::TcpSocket, family: network::Ip
     // Check that we sent and recieved our message!
     assert_eq!(data, second_message); // Not guaranteed to work but should work in practice.
 }
+
+pub fn example_body_udp(net: udp::Network, sock: udp::UdpSocket, family: network::IpAddressFamily) {
+    let first_message = b"Hello, world!";
+    let second_message = b"Greetings, planet!";
+
+    let sub = sock.subscribe();
+
+    let addr = sock.local_address().unwrap();
+
+    let client = udp_create_socket::create_udp_socket(family).unwrap();
+    let client_sub = client.subscribe();
+
+    client.start_connect(&net, addr).unwrap();
+    poll::poll_one(&client_sub);
+    client.finish_connect().unwrap();
+
+    let _client_addr = client.local_address().unwrap();
+
+    let n = client
+        .send(&[
+            udp::Datagram {
+                data: vec![],
+                remote_address: addr,
+            },
+            udp::Datagram {
+                data: first_message.to_vec(),
+                remote_address: addr,
+            },
+        ])
+        .unwrap();
+    assert_eq!(n, 2);
+
+    drop(client_sub);
+    drop(client);
+
+    poll::poll_one(&sub);
+    let datagrams = sock.receive(2).unwrap();
+    let mut datagrams = datagrams.into_iter();
+    let (first, second) = match (datagrams.next(), datagrams.next(), datagrams.next()) {
+        (Some(first), Some(second), None) => (first, second),
+        (Some(_first), None, None) => panic!("only one datagram received"),
+        (None, None, None) => panic!("no datagrams received"),
+        _ => panic!("invalid datagram sequence received"),
+    };
+
+    assert!(first.data.is_empty());
+
+    // TODO: Verify the `remote_address`
+    //assert_eq!(first.remote_address, client_addr);
+
+    // Check that we sent and recieved our message!
+    assert_eq!(second.data, first_message); // Not guaranteed to work but should work in practice.
+
+    // TODO: Verify the `remote_address`
+    //assert_eq!(second.remote_address, client_addr);
+
+    // Another client
+    let client = udp_create_socket::create_udp_socket(family).unwrap();
+    let client_sub = client.subscribe();
+
+    client.start_connect(&net, addr).unwrap();
+    poll::poll_one(&client_sub);
+    client.finish_connect().unwrap();
+
+    let n = client
+        .send(&[udp::Datagram {
+            data: second_message.to_vec(),
+            remote_address: addr,
+        }])
+        .unwrap();
+    assert_eq!(n, 1);
+
+    drop(client_sub);
+    drop(client);
+
+    poll::poll_one(&sub);
+    let datagrams = sock.receive(2).unwrap();
+    let mut datagrams = datagrams.into_iter();
+    let first = match (datagrams.next(), datagrams.next()) {
+        (Some(first), None) => first,
+        (None, None) => panic!("no datagrams received"),
+        _ => panic!("invalid datagram sequence received"),
+    };
+
+    // Check that we sent and recieved our message!
+    assert_eq!(first.data, second_message); // Not guaranteed to work but should work in practice.
+}

crates/wasi/src/preview2/command.rs

@@ -48,6 +48,8 @@ pub fn add_to_linker<T: WasiView>(l: &mut wasmtime::component::Linker<T>) -> any
     crate::preview2::bindings::cli::terminal_stderr::add_to_linker(l, |t| t)?;
     crate::preview2::bindings::sockets::tcp::add_to_linker(l, |t| t)?;
     crate::preview2::bindings::sockets::tcp_create_socket::add_to_linker(l, |t| t)?;
+    crate::preview2::bindings::sockets::udp::add_to_linker(l, |t| t)?;
+    crate::preview2::bindings::sockets::udp_create_socket::add_to_linker(l, |t| t)?;
     crate::preview2::bindings::sockets::instance_network::add_to_linker(l, |t| t)?;
     crate::preview2::bindings::sockets::network::add_to_linker(l, |t| t)?;
     crate::preview2::bindings::sockets::ip_name_lookup::add_to_linker(l, |t| t)?;
@@ -65,6 +67,7 @@ pub mod sync {
            "wasi:filesystem/types": crate::preview2::bindings::sync_io::filesystem::types,
            "wasi:filesystem/preopens": crate::preview2::bindings::filesystem::preopens,
            "wasi:sockets/tcp": crate::preview2::bindings::sockets::tcp,
+           "wasi:sockets/udp": crate::preview2::bindings::sockets::udp,
            "wasi:clocks/monotonic_clock": crate::preview2::bindings::clocks::monotonic_clock,
            "wasi:io/poll": crate::preview2::bindings::sync_io::io::poll,
            "wasi:io/streams": crate::preview2::bindings::sync_io::io::streams,
@@ -107,6 +110,8 @@ pub mod sync {
         crate::preview2::bindings::cli::terminal_stderr::add_to_linker(l, |t| t)?;
         crate::preview2::bindings::sockets::tcp::add_to_linker(l, |t| t)?;
         crate::preview2::bindings::sockets::tcp_create_socket::add_to_linker(l, |t| t)?;
+        crate::preview2::bindings::sockets::udp::add_to_linker(l, |t| t)?;
+        crate::preview2::bindings::sockets::udp_create_socket::add_to_linker(l, |t| t)?;
         crate::preview2::bindings::sockets::instance_network::add_to_linker(l, |t| t)?;
         crate::preview2::bindings::sockets::network::add_to_linker(l, |t| t)?;
         crate::preview2::bindings::sockets::ip_name_lookup::add_to_linker(l, |t| t)?;

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

@@ -8,3 +8,5 @@ mod network;
 mod random;
 mod tcp;
 mod tcp_create_socket;
+mod udp;
+mod udp_create_socket;