examples: update tcp_client-rs and udp_socket-rs to support no_std

- add ipv6 call in tcp_client-rs and udp_socket-rs - enable no_std run of tcp_client-rs and udp_socket-rs in ci Signed-off-by: ivila <[email protected]> Reviewed-by: Yuan Zhuang <[email protected]>
apache · Jan 24, 2025 · a9e722b · a9e722b
1 parent e6d246e
commit a9e722b
Show file tree

Hide file tree

Showing 18 changed files with 214 additions and 101 deletions.
diff --git a/README.md b/README.md
@@ -81,8 +81,8 @@ branch (`main`), please refer to the
 - **Common**: See
   [Overview of OP-TEE Rust Examples](https://teaclave.apache.org/trustzone-sdk-docs/overview-of-optee-rust-examples/).
 
-- **`no-std`**: Excludes `test_serde`, `test_tcp_client`, `test_udp_socket`,
-  `test_message_passing_interface`, `test_tls_client`, `test_tls_server`.
+- **`no-std`**: Excludes `test_serde`, `test_message_passing_interface`,
+  `test_tls_client`, `test_tls_server`.
 
 
 ## Quick Start with the OP-TEE Repo for QEMUv8

diff --git a/ci/ci.sh b/ci/ci.sh
@@ -35,13 +35,13 @@ pushd ../tests
 ./test_signature_verification.sh
 ./test_supp_plugin.sh
 ./test_error_handling.sh
+./test_tcp_client.sh
+./test_udp_socket.sh
 
 # Run std only tests
 if [ "$STD" ]; then
     ./test_serde.sh
     ./test_message_passing_interface.sh
-    ./test_tcp_client.sh
-    ./test_udp_socket.sh
     ./test_tls_client.sh
     ./test_tls_server.sh
     ./test_eth_wallet.sh

diff --git a/examples/tcp_client-rs/host/Cargo.toml b/examples/tcp_client-rs/host/Cargo.toml
@@ -28,6 +28,7 @@ edition = "2018"
 libc = "0.2.48"
 proto = { path = "../proto" }
 optee-teec = { path = "../../../optee-teec" }
+tiny_http = "0.12.0"
 
 [profile.release]
 lto = true
diff --git a/examples/tcp_client-rs/host/Makefile b/examples/tcp_client-rs/host/Makefile
@@ -24,12 +24,7 @@ LINKER_CFG := target.$(TARGET).linker=\"$(CROSS_COMPILE)gcc\"
 
 OUT_DIR := $(CURDIR)/target/$(TARGET)/release
 
-ifeq ($(STD),)
-all:
-	@echo "Please \`export STD=y\` then rerun \`source environment\` to build the STD version"
-else
 all: host strip
-endif
 
 host:
 	@cargo build --target $(TARGET_HOST) --release --config $(LINKER_CFG)

diff --git a/examples/tcp_client-rs/host/src/main.rs b/examples/tcp_client-rs/host/src/main.rs
@@ -15,12 +15,28 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use optee_teec::ParamNone;
-use optee_teec::{Context, Operation, Session, Uuid};
-use proto::{Command, UUID};
+use std::net::{Ipv6Addr, SocketAddr, SocketAddrV6};
+use std::thread;
 
-fn tcp_client(session: &mut Session) -> optee_teec::Result<()> {
-    let mut operation = Operation::new(0, ParamNone, ParamNone, ParamNone, ParamNone);
+use optee_teec::{Context, Operation, ParamType, Session, Uuid};
+use optee_teec::{ParamNone, ParamTmpRef, ParamValue};
+use proto::{Command, IpVersion, UUID};
+
+fn tcp_client(
+    session: &mut Session,
+    address: &str,
+    port: u16,
+    ip_version: IpVersion,
+    host_name: &str,
+) -> optee_teec::Result<()> {
+    let http_data = format!("GET / HTTP/1.0\r\nHost: {}\r\n\r\n", host_name);
+    let mut operation = Operation::new(
+        0,
+        ParamTmpRef::new_input(address.as_bytes()),
+        ParamValue::new(port as u32, ip_version as u32, ParamType::ValueInput),
+        ParamTmpRef::new_input(http_data.as_bytes()),
+        ParamNone,
+    );
     session.invoke_command(Command::Start as u32, &mut operation)?;
     Ok(())
 }
@@ -30,7 +46,31 @@ fn main() -> optee_teec::Result<()> {
     let uuid = Uuid::parse_str(UUID).unwrap();
     let mut session = ctx.open_session(uuid)?;
 
-    tcp_client(&mut session)?;
+    // test ipv4
+    const IPV4_HOST: &str = "teaclave.apache.org";
+    tcp_client(&mut session, IPV4_HOST, 80, IpVersion::V4, IPV4_HOST)?;
+    // test ipv6
+    let addr = SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 0, 0, 0));
+    let server = tiny_http::Server::http(addr).unwrap();
+    let listen_addr = server.server_addr().to_ip().unwrap();
+    let ip = listen_addr.ip().to_string();
+    let port = listen_addr.port();
+
+    let _ = thread::spawn(move || {
+        for request in server.incoming_requests() {
+            println!(
+                "received request! method: {:?}, url: {:?}, headers: {:?}",
+                request.method(),
+                request.url(),
+                request.headers()
+            );
+
+            let response = tiny_http::Response::from_string("hello world");
+            request.respond(response).unwrap();
+        }
+    });
+
+    tcp_client(&mut session, &ip, port, IpVersion::V6, &ip)?;
 
     println!("Success");
     Ok(())

diff --git a/examples/tcp_client-rs/proto/Cargo.toml b/examples/tcp_client-rs/proto/Cargo.toml
@@ -25,6 +25,7 @@ description = "Data structures and functions shared by host and TA."
 edition = "2018"
 
 [dependencies]
+num_enum = { version = "0.7.3", default-features = false }
 
 [build-dependencies]
 uuid = { version = "1.8", default-features = false }
diff --git a/examples/tcp_client-rs/proto/src/lib.rs b/examples/tcp_client-rs/proto/src/lib.rs
@@ -14,20 +14,22 @@
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.
+#![no_std]
+use num_enum::{FromPrimitive, TryFromPrimitive, IntoPrimitive};
 
+#[derive(FromPrimitive, IntoPrimitive)]
+#[repr(u32)]
 pub enum Command {
     Start,
+    #[default]
     Unknown,
 }
 
-impl From<u32> for Command {
-    #[inline]
-    fn from(value: u32) -> Command {
-        match value {
-            0 => Command::Start,
-            _ => Command::Unknown,
-        }
-    }
+#[derive(TryFromPrimitive, IntoPrimitive)]
+#[repr(u32)]
+pub enum IpVersion {
+    V4 = 1,
+    V6 = 2,
 }
 
 pub const UUID: &str = &include_str!(concat!(env!("OUT_DIR"), "/uuid.txt"));
diff --git a/examples/tcp_client-rs/ta/Cargo.toml b/examples/tcp_client-rs/ta/Cargo.toml
@@ -28,12 +28,13 @@ edition = "2018"
 proto = { path = "../proto" }
 optee-utee-sys = { path = "../../../optee-utee/optee-utee-sys" }
 optee-utee = { path = "../../../optee-utee" }
+cfg_block = "0.2.0"
 
 [build-dependencies]
 proto = { path = "../proto" }
 optee-utee-build = { path = "../../../optee-utee-build" }
 
 [profile.release]
 panic = "abort"
-lto = false
+lto = true
 opt-level = 1
diff --git a/examples/tcp_client-rs/ta/Makefile b/examples/tcp_client-rs/ta/Makefile
@@ -29,15 +29,12 @@ TA_SIGN_KEY ?= $(TA_DEV_KIT_DIR)/keys/default_ta.pem
 SIGN := $(TA_DEV_KIT_DIR)/scripts/sign_encrypt.py
 OUT_DIR := $(CURDIR)/target/$(TARGET)/release
 
-ifeq ($(STD),)
-all:
-	@echo "Please \`export STD=y\` then rerun \`source environment\` to build the STD version"
-else
+BUILDER = $(if $(STD),xargo,cargo)
+
 all: ta strip sign
-endif
 
 ta:
-	@xargo build --target $(TARGET) --release --config $(LINKER_CFG)
+	@$(BUILDER) build --target $(TARGET) --release --config $(LINKER_CFG)
 
 strip: ta
 	@$(OBJCOPY) --strip-unneeded $(OUT_DIR)/ta $(OUT_DIR)/stripped_ta

diff --git a/examples/tcp_client-rs/ta/src/main.rs b/examples/tcp_client-rs/ta/src/main.rs
@@ -15,16 +15,28 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#![cfg_attr(not(target_os = "optee"), no_std)]
 #![no_main]
 
+cfg_block::cfg_block! {
+    // In Teaclave, if target_os = "optee", the codes is compiled with std.
+    // Otherwise, no-std
+    if #[cfg(target_os = "optee")] {
+        use std::io::{Read, Write};
+    } else {
+        extern crate alloc;
+        use optee_utee::net::{StdCompatConnect, StdCompatWrite, StdCompatRead};
+        use alloc::vec::Vec;
+        use alloc::string::String;
+    }
+}
+
 use optee_utee::net::TcpStream;
 use optee_utee::{
     ta_close_session, ta_create, ta_destroy, ta_invoke_command, ta_open_session, trace_println,
 };
 use optee_utee::{Error, ErrorKind, Parameters, Result};
-use proto::Command;
-use std::io::Read;
-use std::io::Write;
+use proto::{Command, IpVersion};
 
 #[ta_create]
 fn create() -> Result<()> {
@@ -49,35 +61,56 @@ fn destroy() {
 }
 
 #[ta_invoke_command]
-fn invoke_command(cmd_id: u32, _params: &mut Parameters) -> Result<()> {
+fn invoke_command(cmd_id: u32, params: &mut Parameters) -> Result<()> {
     trace_println!("[+] TA invoke command");
     match Command::from(cmd_id) {
         Command::Start => {
-            tcp_client();
-            Ok(())
+            use core::convert::TryFrom;
+
+            let mut param0 = unsafe { params.0.as_memref()? };
+            let param1 = unsafe { params.1.as_value()? };
+            let mut param2 = unsafe { params.2.as_memref()? };
+
+            let address = core::str::from_utf8(param0.buffer()).unwrap();
+            let port = param1.a() as u16;
+            let ip_version =
+                IpVersion::try_from(param1.b()).map_err(|_| ErrorKind::BadParameters)?;
+            let http_data = param2.buffer();
+
+            tcp_client(address, port, ip_version, http_data)
         }
         _ => Err(Error::new(ErrorKind::BadParameters)),
     }
 }
 
-fn tcp_client() {
-    let mut stream = TcpStream::connect("teaclave.apache.org", 80).unwrap();
-    stream
-        .write_all(b"GET / HTTP/1.0\r\nHost: teaclave.apache.org\r\n\r\n")
-        .unwrap();
+fn tcp_client(address: &str, port: u16, ip_version: IpVersion, http_data: &[u8]) -> Result<()> {
+    let mut stream = match ip_version {
+        IpVersion::V4 => TcpStream::connect_v4(address, port),
+        IpVersion::V6 => TcpStream::connect_v6(address, port),
+    }
+    .map_err(|err| {
+        trace_println!("failed to connect to {}:{} due to {:?}", address, port, err);
+        ErrorKind::Generic
+    })?;
+
+    stream.write_all(http_data).map_err(|err| {
+        trace_println!("failed to write_all due to {:?}", err);
+        ErrorKind::Generic
+    })?;
     let mut response = Vec::new();
     let mut chunk = [0u8; 1024];
     loop {
         match stream.read(&mut chunk) {
             Ok(0) => break,
             Ok(n) => response.extend_from_slice(&chunk[..n]),
-            Err(_) => {
-                trace_println!("Error");
-                panic!();
+            Err(err) => {
+                trace_println!("failed to read due to {:?}", err);
+                return Err(ErrorKind::Generic.into());
             }
         }
     }
     trace_println!("{}", String::from_utf8_lossy(&response));
+    Ok(())
 }
 
 include!(concat!(env!("OUT_DIR"), "/user_ta_header.rs"));
diff --git a/examples/udp_socket-rs/host/Makefile b/examples/udp_socket-rs/host/Makefile
@@ -24,12 +24,7 @@ LINKER_CFG := target.$(TARGET).linker=\"$(CROSS_COMPILE)gcc\"
 
 OUT_DIR := $(CURDIR)/target/$(TARGET)/release
 
-ifeq ($(STD),)
-all:
-	@echo "Please \`export STD=y\` then rerun \`source environment\` to build the STD version"
-else
 all: host strip
-endif
 
 host:
 	@cargo build --target $(TARGET_HOST) --release --config $(LINKER_CFG)

diff --git a/examples/udp_socket-rs/host/src/main.rs b/examples/udp_socket-rs/host/src/main.rs
@@ -14,28 +14,44 @@
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.
-
-use optee_teec::ParamNone;
-use optee_teec::{Context, Operation, Session, Uuid};
-use proto::{Command, UUID};
-use std::net::UdpSocket;
+use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, UdpSocket};
 use std::str;
 use std::thread;
 
-fn udp_socket(session: &mut Session) -> optee_teec::Result<()> {
-    let mut operation = Operation::new(0, ParamNone, ParamNone, ParamNone, ParamNone);
-    session.invoke_command(Command::Start as u32, &mut operation)?;
-    Ok(())
-}
+use optee_teec::{Context, Operation, Uuid};
+use optee_teec::{ParamNone, ParamTmpRef, ParamValue};
+use proto::{Command, IpVersion, UUID};
 
-fn main() -> optee_teec::Result<()> {
-    let socket = UdpSocket::bind("127.0.0.1:34254").unwrap();
+fn udp_socket(ip_version: IpVersion) {
+    let addr: SocketAddr = match ip_version {
+        IpVersion::V4 => SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)),
+        IpVersion::V6 => SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 0, 0, 0)),
+    };
+    let socket = UdpSocket::bind(addr).unwrap();
+    let local_addr = socket.local_addr().unwrap();
+    println!("Test on: {}", local_addr);
 
-    let mut ctx = Context::new()?;
-    let uuid = Uuid::parse_str(UUID).unwrap();
     let child = thread::spawn(move || {
+        let mut ctx = Context::new().unwrap();
+        let uuid = Uuid::parse_str(UUID).unwrap();
         let mut session = ctx.open_session(uuid).unwrap();
-        udp_socket(&mut session).unwrap();
+
+        let ip = local_addr.ip().to_string();
+        let port = local_addr.port();
+        let mut operation = Operation::new(
+            0,
+            ParamTmpRef::new_input(ip.as_bytes()),
+            ParamValue::new(
+                port as u32,
+                ip_version as u32,
+                optee_teec::ParamType::ValueInput,
+            ),
+            ParamNone,
+            ParamNone,
+        );
+        session
+            .invoke_command(Command::Start as u32, &mut operation)
+            .unwrap();
     });
 
     let mut buf = [0; 100];
@@ -47,5 +63,11 @@ fn main() -> optee_teec::Result<()> {
     let _ = child.join();
 
     println!("Success");
+}
+
+fn main() -> optee_teec::Result<()> {
+    udp_socket(IpVersion::V4);
+    udp_socket(IpVersion::V6);
+
     Ok(())
 }
diff --git a/examples/udp_socket-rs/proto/Cargo.toml b/examples/udp_socket-rs/proto/Cargo.toml
@@ -25,6 +25,7 @@ description = "Data structures and functions shared by host and TA."
 edition = "2018"
 
 [dependencies]
+num_enum = { version = "0.7.3", default-features = false }
 
 [build-dependencies]
 uuid = { version = "1.8", default-features = false }