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

Signed-off-by: ivila <[email protected]>
Reviewed-by: Yuan Zhuang <[email protected]>

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

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

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)

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,

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

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

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;
 
 #[ta_create]
 fn create() -> Result<()> {
@@ -52,32 +64,36 @@ 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();
+fn tcp_client() -> Result<()> {
+    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"));

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)

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,

examples/udp_socket-rs/ta/Cargo.toml

@@ -25,16 +25,16 @@ 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" }
+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

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

examples/udp_socket-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::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 +64,20 @@ 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();
+fn udp_socket() -> Result<()> {
+    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 +89,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"));

tests/optee-qemuv8.sh

@@ -17,13 +17,6 @@
 # specific language governing permissions and limitations
 # under the License.
 
-# std examples: tcp_client, udp_socket, tls_client, tls_server needs the external network
-if [ "$STD" ]; then
-    EXTERNAL_NETWORK_PARAMS=" \
-    -netdev user,id=vmnic,hostfwd=:127.0.0.1:54433-:4433 \
-    -device virtio-net-device,netdev=vmnic"
-fi
-
 cd $1 && ./qemu-system-aarch64 \
     -nodefaults \
     -nographic \
@@ -37,4 +30,6 @@ cd $1 && ./qemu-system-aarch64 \
     -append 'console=ttyAMA0,38400 keep_bootcon root=/dev/vda2' \
     -kernel Image -no-acpi \
     -fsdev local,id=fsdev0,path=$(pwd)/../shared,security_model=none \
-    -device virtio-9p-device,fsdev=fsdev0,mount_tag=host $EXTERNAL_NETWORK_PARAMS
+    -device virtio-9p-device,fsdev=fsdev0,mount_tag=host \
+    -netdev user,id=vmnic,hostfwd=:127.0.0.1:54433-:4433 \
+    -device virtio-net-device,netdev=vmnic