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

apache · Jan 21, 2025 · b7c5ef7 · b7c5ef7
1 parent 5e9a0a2
commit b7c5ef7
Show file tree

Hide file tree

Showing 10 changed files with 130 additions and 49 deletions.
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/proto/src/lib.rs b/examples/tcp_client-rs/proto/src/lib.rs
@@ -14,6 +14,7 @@
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.
+#![no_std]
 
 pub enum Command {
     Start,

diff --git a/examples/tcp_client-rs/ta/Cargo.toml b/examples/tcp_client-rs/ta/Cargo.toml
@@ -35,5 +35,5 @@ 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,17 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#![cfg_attr(not(target_os = "optee"), no_std)]
 #![no_main]
 
+extern crate alloc;
+
 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;
 
 #[ta_create]
 fn create() -> Result<()> {
@@ -52,32 +53,78 @@ fn destroy() {
 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(())
-        }
+        Command::Start => tcp_client(),
         _ => Err(Error::new(ErrorKind::BadParameters)),
     }
 }
 
-fn tcp_client() {
-    let mut stream = TcpStream::connect("teaclave.apache.org", 80).unwrap();
+#[cfg(not(target_os = "optee"))]
+fn tcp_client() -> Result<()> {
+    use alloc::string::String;
+    use alloc::vec::Vec;
+    use optee_utee::net::Setup;
+
+    let setup = Setup::new_v4("teaclave.apache.org", 80)?;
+    let mut stream = TcpStream::open(setup).map_err(|err| {
+        trace_println!("failed to open due to {:?}", err);
+        ErrorKind::Generic
+    })?;
+
+    stream.set_send_timeout_in_milli(10 * 1000);
+    stream
+        .send(b"GET / HTTP/1.0\r\nHost: teaclave.apache.org\r\n\r\n")
+        .map_err(|err| {
+            trace_println!("failed to send due to {:?}", err);
+            ErrorKind::Generic
+        })?;
+
+    let mut response = Vec::new();
+    let mut chunk = [0u8; 1024];
+    stream.set_recv_timeout_in_milli(10 * 1000);
+    loop {
+        match stream.recv(&mut chunk) {
+            Ok(0) => break,
+            Ok(n) => response.extend_from_slice(&chunk[..n]),
+            Err(err) => {
+                trace_println!("failed to recv due to {:?}", err);
+                return Err(ErrorKind::Generic.into());
+            }
+        }
+    }
+    trace_println!("{}", String::from_utf8_lossy(&response));
+    Ok(())
+}
+
+#[cfg(target_os = "optee")]
+// For STD version, developers can also use APIs similar to std::net::TcpStream.
+fn tcp_client() -> Result<()> {
+    extern crate std;
+    use std::io::{Read, Write};
+
+    let mut stream = TcpStream::connect("teaclave.apache.org", 80).map_err(|err| {
+        trace_println!("failed to connect due to {:?}", err);
+        ErrorKind::Generic
+    })?;
     stream
         .write_all(b"GET / HTTP/1.0\r\nHost: teaclave.apache.org\r\n\r\n")
-        .unwrap();
+        .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/proto/src/lib.rs b/examples/udp_socket-rs/proto/src/lib.rs
@@ -14,6 +14,7 @@
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.
+#![no_std]
 
 pub enum Command {
     Start,

diff --git a/examples/udp_socket-rs/ta/Cargo.toml b/examples/udp_socket-rs/ta/Cargo.toml
@@ -25,7 +25,6 @@ description = "An example of Rust OP-TEE TrustZone SDK."
 edition = "2018"
 
 [dependencies]
-libc = { path = "../../../rust/libc" }
 proto = { path = "../proto" }
 optee-utee-sys = { path = "../../../optee-utee/optee-utee-sys" }
 optee-utee = { path = "../../../optee-utee" }
@@ -36,5 +35,5 @@ optee-utee-build = { path = "../../../optee-utee-build" }
 
 [profile.release]
 panic = "abort"
-lto = false
+lto = true
 opt-level = 1
diff --git a/examples/udp_socket-rs/ta/Makefile b/examples/udp_socket-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/udp_socket-rs/ta/src/main.rs b/examples/udp_socket-rs/ta/src/main.rs
@@ -15,16 +15,17 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#![cfg_attr(not(target_os = "optee"), no_std)]
 #![no_main]
 
+extern crate alloc;
+
 use optee_utee::net::UdpSocket;
 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;
 
 #[ta_create]
 fn create() -> Result<()> {
@@ -52,17 +53,64 @@ fn destroy() {
 fn invoke_command(cmd_id: u32, _params: &mut Parameters) -> Result<()> {
     trace_println!("[+] TA invoke command");
     match Command::from(cmd_id) {
-        Command::Start => {
-            udp_socket();
-            Ok(())
-        }
+        Command::Start => udp_socket(),
         _ => Err(Error::new(ErrorKind::BadParameters)),
     }
 }
 
-fn udp_socket() {
-    let mut stream = UdpSocket::connect("127.0.0.1", 34254).unwrap();
-    stream.write_all(b"[TA]: Hello, Teaclave!").unwrap();
+#[cfg(not(target_os = "optee"))]
+fn udp_socket() -> Result<()> {
+    use alloc::string::String;
+    use alloc::vec::Vec;
+    use optee_utee::net::Setup;
+
+    let setup = Setup::new_v4("127.0.0.1", 34254)?;
+    let mut stream = UdpSocket::open(setup).map_err(|err| {
+        trace_println!("failed to open due to {:?}", err);
+        ErrorKind::Generic
+    })?;
+
+    stream.set_send_timeout_in_milli(10 * 1000);
+    stream.send(b"[TA]: Hello, Teaclave!").map_err(|err| {
+        trace_println!("failed to send due to {:?}", err);
+        ErrorKind::Generic
+    })?;
+
+    let mut response = Vec::new();
+    let mut chunk = [0u8; 1024];
+    stream.set_recv_timeout_in_milli(10 * 1000);
+
+    // Loop until read something.
+    loop {
+        match stream.recv(&mut chunk) {
+            Ok(0) => continue,
+            Ok(n) => {
+                response.extend_from_slice(&chunk[..n]);
+                break;
+            }
+            Err(err) => {
+                trace_println!("failed to read due to {:?}", err);
+                return Err(ErrorKind::Generic.into());
+            }
+        }
+    }
+    trace_println!("{}", String::from_utf8_lossy(&response));
+    Ok(())
+}
+
+#[cfg(target_os = "optee")]
+// For STD version, developers can also use APIs similar to std::net::UdpSocket.
+fn udp_socket() -> Result<()> {
+    use std::io::{Read, Write};
+
+    let mut stream = UdpSocket::connect("127.0.0.1", 34254).map_err(|err| {
+        trace_println!("failed to connect due to {:?}", err);
+        ErrorKind::Generic
+    })?;
+    stream.write_all(b"[TA]: Hello, Teaclave!").map_err(|err| {
+        trace_println!("failed to write_all due to {:?}", err);
+        ErrorKind::Generic
+    })?;
     let mut response = Vec::new();
     let mut chunk = [0u8; 1024];
 
@@ -74,13 +122,14 @@ fn udp_socket() {
                 response.extend_from_slice(&chunk[..n]);
                 break;
             }
-            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"));