This is an automated email from the ASF dual-hosted git repository.

yuanz pushed a commit to branch main
in repository 
https://gitbox.apache.org/repos/asf/incubator-teaclave-trustzone-sdk.git

commit 3fcbb3eda15a6ac71132ff0a69ed080429d3e9cd
Author: ivila <[email protected]>
AuthorDate: Mon Jan 20 15:14:14 2025 +0800

    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]>
    Reviewed-by: Sumit Garg <[email protected]>
---
 README.md                               |  4 +--
 ci/ci.sh                                |  4 +--
 examples/tcp_client-rs/host/Cargo.toml  |  1 +
 examples/tcp_client-rs/host/Makefile    |  5 ---
 examples/tcp_client-rs/host/src/main.rs | 62 ++++++++++++++++++++++++++++----
 examples/tcp_client-rs/proto/Cargo.toml |  1 +
 examples/tcp_client-rs/proto/src/lib.rs | 18 +++++-----
 examples/tcp_client-rs/ta/Cargo.toml    |  3 +-
 examples/tcp_client-rs/ta/Makefile      |  9 ++---
 examples/tcp_client-rs/ta/src/main.rs   | 64 +++++++++++++++++++++++++--------
 examples/udp_socket-rs/host/Makefile    |  5 ---
 examples/udp_socket-rs/host/src/main.rs | 52 +++++++++++++++++++--------
 examples/udp_socket-rs/proto/Cargo.toml |  1 +
 examples/udp_socket-rs/proto/src/lib.rs | 20 ++++++-----
 examples/udp_socket-rs/ta/Cargo.toml    |  4 +--
 examples/udp_socket-rs/ta/Makefile      |  9 ++---
 examples/udp_socket-rs/ta/src/main.rs   | 59 +++++++++++++++++++++++-------
 tests/optee-qemuv8.sh                   | 11 ++----
 18 files changed, 230 insertions(+), 102 deletions(-)

diff --git a/README.md b/README.md
index 7a13199..3b608be 100644
--- 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
index c58a03f..e10916f 100755
--- 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
index c8f3610..8c314d5 100644
--- 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
index 7675483..75a545f 100644
--- 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
index 75624db..6a0acdf 100644
--- a/examples/tcp_client-rs/host/src/main.rs
+++ b/examples/tcp_client-rs/host/src/main.rs
@@ -15,13 +15,33 @@
 // 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<()> {
+    println!("Test on: {}", address);
+
+    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)?;
+
+    println!("Success");
     Ok(())
 }
 
@@ -30,8 +50,36 @@ 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";
+    // Use the host directly to also check its domain name resolving 
capability.
+    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 child = 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();
+            break;
+        }
+    });
+    // Use the IP address directly to ensure we're actually trying an IPv6
+    // address.
+    tcp_client(&mut session, &ip, port, IpVersion::V6, &ip)?;
+    let _ = child.join();
 
-    println!("Success");
     Ok(())
 }
diff --git a/examples/tcp_client-rs/proto/Cargo.toml 
b/examples/tcp_client-rs/proto/Cargo.toml
index b76ceb2..4a173c8 100644
--- 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
index 7679b2d..9034643 100644
--- 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
index fce4cb4..5f71c19 100644
--- a/examples/tcp_client-rs/ta/Cargo.toml
+++ b/examples/tcp_client-rs/ta/Cargo.toml
@@ -28,6 +28,7 @@ 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" }
@@ -35,5 +36,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
index 42d9236..e69d885 100644
--- 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
index bc3a2b8..ddd9293 100644
--- 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,59 @@ 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.set_send_timeout_in_milli(60 * 1000);
+    stream.set_recv_timeout_in_milli(60 * 1000);
+
+    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
index 90b6d44..734ed5c 100644
--- 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
index 5e1db59..62076cc 100644
--- 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
index b76ceb2..4a173c8 100644
--- 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 }
diff --git a/examples/udp_socket-rs/proto/src/lib.rs 
b/examples/udp_socket-rs/proto/src/lib.rs
index 7679b2d..88fa36b 100644
--- a/examples/udp_socket-rs/proto/src/lib.rs
+++ b/examples/udp_socket-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, IntoPrimitive, TryFromPrimitive};
 
+#[derive(FromPrimitive, IntoPrimitive)]
+#[repr(u32)]
 pub enum Command {
-    Start,
+    Start = 1,
+    #[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/udp_socket-rs/ta/Cargo.toml 
b/examples/udp_socket-rs/ta/Cargo.toml
index ba7d6d4..5f71c19 100644
--- a/examples/udp_socket-rs/ta/Cargo.toml
+++ b/examples/udp_socket-rs/ta/Cargo.toml
@@ -25,10 +25,10 @@ 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" }
@@ -36,5 +36,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
index 42d9236..e69d885 100644
--- 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
index 867c9d8..c187797 100644
--- a/examples/udp_socket-rs/ta/src/main.rs
+++ b/examples/udp_socket-rs/ta/src/main.rs
@@ -15,16 +15,29 @@
 // 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 core::convert::TryFrom;
 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;
+use proto::{Command, IpVersion};
 
 #[ta_create]
 fn create() -> Result<()> {
@@ -49,20 +62,41 @@ 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 => {
-            udp_socket();
-            Ok(())
+            let mut param0 = unsafe { params.0.as_memref()? };
+            let param1 = unsafe { params.1.as_value()? };
+
+            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)?;
+
+            udp_socket(address, port, ip_version)
         }
         _ => 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(address: &str, port: u16, ip_version: IpVersion) -> Result<()> {
+    let mut stream = match ip_version {
+        IpVersion::V4 => UdpSocket::connect_v4(address, port),
+        IpVersion::V6 => UdpSocket::connect_v6(address, port),
+    }
+    .map_err(|err| {
+        trace_println!("failed to connect to {}:{} due to {:?}", address, 
port, err);
+        ErrorKind::Generic
+    })?;
+
+    stream.set_send_timeout_in_milli(60 * 1000);
+    stream.set_recv_timeout_in_milli(60 * 1000);
+
+    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 +108,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"));
diff --git a/tests/optee-qemuv8.sh b/tests/optee-qemuv8.sh
index 0b8cd44..2a44e21 100755
--- a/tests/optee-qemuv8.sh
+++ b/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
\ No newline at end of file
+    -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


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to