1515// specific language governing permissions and limitations
1616// under the License.
1717
18+ #![ cfg_attr( not( target_os = "optee" ) , no_std) ]
1819#![ no_main]
1920
21+ cfg_block:: cfg_block! {
22+ // In Teaclave, if target_os = "optee", the codes is compiled with std.
23+ // Otherwise, no-std
24+ if #[ cfg( target_os = "optee" ) ] {
25+ use std:: io:: { Read , Write } ;
26+ } else {
27+ extern crate alloc;
28+ use optee_utee:: net:: { StdCompatConnect , StdCompatWrite , StdCompatRead } ;
29+ use alloc:: vec:: Vec ;
30+ use alloc:: string:: String ;
31+ }
32+ }
33+
2034use optee_utee:: net:: TcpStream ;
2135use optee_utee:: {
2236 ta_close_session, ta_create, ta_destroy, ta_invoke_command, ta_open_session, trace_println,
2337} ;
2438use optee_utee:: { Error , ErrorKind , Parameters , Result } ;
25- use proto:: Command ;
26- use std:: io:: Read ;
27- use std:: io:: Write ;
39+ use proto:: { Command , IpVersion } ;
2840
2941#[ ta_create]
3042fn create ( ) -> Result < ( ) > {
@@ -49,35 +61,59 @@ fn destroy() {
4961}
5062
5163#[ ta_invoke_command]
52- fn invoke_command ( cmd_id : u32 , _params : & mut Parameters ) -> Result < ( ) > {
64+ fn invoke_command ( cmd_id : u32 , params : & mut Parameters ) -> Result < ( ) > {
5365 trace_println ! ( "[+] TA invoke command" ) ;
5466 match Command :: from ( cmd_id) {
5567 Command :: Start => {
56- tcp_client ( ) ;
57- Ok ( ( ) )
68+ use core:: convert:: TryFrom ;
69+
70+ let mut param0 = unsafe { params. 0 . as_memref ( ) ? } ;
71+ let param1 = unsafe { params. 1 . as_value ( ) ? } ;
72+ let mut param2 = unsafe { params. 2 . as_memref ( ) ? } ;
73+
74+ let address = core:: str:: from_utf8 ( param0. buffer ( ) ) . unwrap ( ) ;
75+ let port = param1. a ( ) as u16 ;
76+ let ip_version =
77+ IpVersion :: try_from ( param1. b ( ) ) . map_err ( |_| ErrorKind :: BadParameters ) ?;
78+ let http_data = param2. buffer ( ) ;
79+
80+ tcp_client ( address, port, ip_version, http_data)
5881 }
5982 _ => Err ( Error :: new ( ErrorKind :: BadParameters ) ) ,
6083 }
6184}
6285
63- fn tcp_client ( ) {
64- let mut stream = TcpStream :: connect ( "teaclave.apache.org" , 80 ) . unwrap ( ) ;
65- stream
66- . write_all ( b"GET / HTTP/1.0\r \n Host: teaclave.apache.org\r \n \r \n " )
67- . unwrap ( ) ;
86+ fn tcp_client ( address : & str , port : u16 , ip_version : IpVersion , http_data : & [ u8 ] ) -> Result < ( ) > {
87+ let mut stream = match ip_version {
88+ IpVersion :: V4 => TcpStream :: connect_v4 ( address, port) ,
89+ IpVersion :: V6 => TcpStream :: connect_v6 ( address, port) ,
90+ }
91+ . map_err ( |err| {
92+ trace_println ! ( "failed to connect to {}:{} due to {:?}" , address, port, err) ;
93+ ErrorKind :: Generic
94+ } ) ?;
95+
96+ stream. set_send_timeout_in_milli ( 60 * 1000 ) ;
97+ stream. set_recv_timeout_in_milli ( 60 * 1000 ) ;
98+
99+ stream. write_all ( http_data) . map_err ( |err| {
100+ trace_println ! ( "failed to write_all due to {:?}" , err) ;
101+ ErrorKind :: Generic
102+ } ) ?;
68103 let mut response = Vec :: new ( ) ;
69104 let mut chunk = [ 0u8 ; 1024 ] ;
70105 loop {
71106 match stream. read ( & mut chunk) {
72107 Ok ( 0 ) => break ,
73108 Ok ( n) => response. extend_from_slice ( & chunk[ ..n] ) ,
74- Err ( _ ) => {
75- trace_println ! ( "Error" ) ;
76- panic ! ( ) ;
109+ Err ( err ) => {
110+ trace_println ! ( "failed to read due to {:?}" , err ) ;
111+ return Err ( ErrorKind :: Generic . into ( ) ) ;
77112 }
78113 }
79114 }
80115 trace_println ! ( "{}" , String :: from_utf8_lossy( & response) ) ;
116+ Ok ( ( ) )
81117}
82118
83119include ! ( concat!( env!( "OUT_DIR" ) , "/user_ta_header.rs" ) ) ;
0 commit comments