| // Copyright (C) 2017-2019 Baidu, Inc. All Rights Reserved. |
| // |
| // Redistribution and use in source and binary forms, with or without |
| // modification, are permitted provided that the following conditions |
| // are met: |
| // |
| // * Redistributions of source code must retain the above copyright |
| // notice, this list of conditions and the following disclaimer. |
| // * Redistributions in binary form must reproduce the above copyright |
| // notice, this list of conditions and the following disclaimer in |
| // the documentation and/or other materials provided with the |
| // distribution. |
| // * Neither the name of Baidu, Inc., nor the names of its |
| // contributors may be used to endorse or promote products derived |
| // from this software without specific prior written permission. |
| // |
| // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS |
| // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT |
| // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR |
| // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT |
| // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, |
| // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT |
| // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, |
| // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY |
| // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT |
| // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE |
| // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
| |
| #![allow(dead_code)] |
| #![allow(unused_assignments)] |
| |
| extern crate sgx_types; |
| extern crate sgx_urts; |
| extern crate dirs; |
| use sgx_types::*; |
| use sgx_urts::SgxEnclave; |
| |
| extern crate mio; |
| use mio::tcp::{TcpListener, TcpStream, Shutdown}; |
| |
| use std::os::unix::io::AsRawFd; |
| use std::ffi::CString; |
| use std::fs; |
| use std::path; |
| use std::net; |
| use std::str; |
| use std::io; |
| use std::io::{Read, Write}; |
| use std::collections::HashMap; |
| |
| const BUFFER_SIZE: usize = 1024; |
| |
| static ENCLAVE_FILE: &'static str = "enclave.signed.so"; |
| static ENCLAVE_TOKEN: &'static str = "enclave.token"; |
| |
| extern { |
| fn tls_server_new(eid: sgx_enclave_id_t, retval: *mut size_t, |
| fd: c_int, cert: *const c_char, key: *const c_char) -> sgx_status_t; |
| fn tls_server_read(eid: sgx_enclave_id_t, retval: *mut c_int, |
| session_id: size_t, buf: *mut c_void, cnt: c_int) -> sgx_status_t; |
| fn tls_server_write(eid: sgx_enclave_id_t, retval: *mut c_int, |
| session_id: size_t, buf: *const c_void, cnt: c_int) -> sgx_status_t; |
| fn tls_server_wants_read(eid: sgx_enclave_id_t, retval: *mut c_int, |
| session_id: size_t) -> sgx_status_t; |
| fn tls_server_wants_write(eid: sgx_enclave_id_t, retval: *mut c_int, |
| session_id: size_t) -> sgx_status_t; |
| fn tls_server_close(eid: sgx_enclave_id_t, |
| session_id: size_t) -> sgx_status_t; |
| fn tls_server_send_close(edi: sgx_enclave_id_t, |
| session_id: size_t) -> sgx_status_t; |
| } |
| |
| fn init_enclave() -> SgxResult<SgxEnclave> { |
| |
| let mut launch_token: sgx_launch_token_t = [0; 1024]; |
| let mut launch_token_updated: i32 = 0; |
| // Step 1: try to retrieve the launch token saved by last transaction |
| // if there is no token, then create a new one. |
| // |
| // try to get the token saved in $HOME */ |
| let mut home_dir = path::PathBuf::new(); |
| let use_token = match dirs::home_dir() { |
| Some(path) => { |
| println!("[+] Home dir is {}", path.display()); |
| home_dir = path; |
| true |
| }, |
| None => { |
| println!("[-] Cannot get home dir"); |
| false |
| } |
| }; |
| |
| let token_file: path::PathBuf = home_dir.join(ENCLAVE_TOKEN);; |
| if use_token == true { |
| match fs::File::open(&token_file) { |
| Err(_) => { |
| println!("[-] Open token file {} error! Will create one.", token_file.as_path().to_str().unwrap()); |
| }, |
| Ok(mut f) => { |
| println!("[+] Open token file success! "); |
| match f.read(&mut launch_token) { |
| Ok(1024) => { |
| println!("[+] Token file valid!"); |
| }, |
| _ => println!("[+] Token file invalid, will create new token file"), |
| } |
| } |
| } |
| } |
| |
| // Step 2: call sgx_create_enclave to initialize an enclave instance |
| // Debug Support: set 2nd parameter to 1 |
| let debug = 1; |
| let mut misc_attr = sgx_misc_attribute_t {secs_attr: sgx_attributes_t { flags:0, xfrm:0}, misc_select:0}; |
| let enclave = try!(SgxEnclave::create(ENCLAVE_FILE, |
| debug, |
| &mut launch_token, |
| &mut launch_token_updated, |
| &mut misc_attr)); |
| |
| // Step 3: save the launch token if it is updated |
| if use_token == true && launch_token_updated != 0 { |
| // reopen the file with write capablity |
| match fs::File::create(&token_file) { |
| Ok(mut f) => { |
| match f.write_all(&launch_token) { |
| Ok(()) => println!("[+] Saved updated launch token!"), |
| Err(_) => println!("[-] Failed to save updated launch token!"), |
| } |
| }, |
| Err(_) => { |
| println!("[-] Failed to save updated enclave token, but doesn't matter"); |
| }, |
| } |
| } |
| |
| Ok(enclave) |
| } |
| |
| // Token for our listening socket. |
| const LISTENER: mio::Token = mio::Token(0); |
| |
| // Which mode the server operates in. |
| #[derive(Clone)] |
| enum ServerMode { |
| /// Write back received bytes |
| Echo, |
| |
| /// Do one read, then write a bodged HTTP response and |
| /// cleanly close the connection. |
| Http, |
| |
| /// Forward traffic to/from given port on localhost. |
| Forward(u16), |
| } |
| |
| /// This binds together a TCP listening socket, some outstanding |
| /// connections, and a TLS server configuration. |
| struct TlsServer { |
| enclave_id: sgx_enclave_id_t, |
| server: TcpListener, |
| cert: CString, |
| key: CString, |
| mode: ServerMode, |
| connections: HashMap<mio::Token, Connection>, |
| next_id: usize, |
| } |
| |
| impl TlsServer { |
| fn new(enclave_id: sgx_enclave_id_t, server: TcpListener, mode: ServerMode, cert: CString, key: CString) -> TlsServer { |
| |
| println!("[+] TlsServer new {:?} {:?}", cert, key); |
| |
| TlsServer { |
| enclave_id: enclave_id, |
| server: server, |
| cert: cert, |
| key: key, |
| mode: mode, |
| connections: HashMap::new(), |
| next_id: 2 |
| } |
| } |
| |
| fn accept(&mut self, poll: &mut mio::Poll) -> bool { |
| match self.server.accept() { |
| Ok((socket, addr)) => { |
| |
| println!("Accepting new connection from {:?}", addr); |
| |
| let mut tlsserver_id: usize = 0xFFFF_FFFF_FFFF_FFFF; |
| let retval = unsafe { |
| tls_server_new(self.enclave_id, |
| &mut tlsserver_id, |
| socket.as_raw_fd(), |
| self.cert.as_bytes_with_nul().as_ptr() as * const c_char, |
| self.key.as_bytes_with_nul().as_ptr() as * const c_char) |
| }; |
| |
| if retval != sgx_status_t::SGX_SUCCESS { |
| println!("[-] ECALL Enclave [tls_server_new] Failed {}!", retval); |
| return false; |
| } |
| |
| if tlsserver_id == 0xFFFF_FFFF_FFFF_FFFF { |
| println!("[-] New enclave tlsserver error"); |
| return false; |
| } |
| |
| let mode = self.mode.clone(); |
| let token = mio::Token(self.next_id); |
| self.next_id += 1; |
| self.connections.insert(token, Connection::new(self.enclave_id, |
| socket, |
| token, |
| mode, |
| tlsserver_id)); |
| self.connections[&token].register(poll); |
| true |
| } |
| Err(e) => { |
| println!("encountered error while accepting connection; err={:?}", e); |
| false |
| } |
| } |
| } |
| |
| fn conn_event(&mut self, poll: &mut mio::Poll, event: &mio::Event) { |
| let token = event.token(); |
| |
| if self.connections.contains_key(&token) { |
| self.connections |
| .get_mut(&token) |
| .unwrap() |
| .ready(poll, event); |
| |
| if self.connections[&token].is_closed() { |
| self.connections.remove(&token); |
| } |
| } |
| } |
| } |
| |
| /// This is a connection which has been accepted by the server, |
| /// and is currently being served. |
| /// |
| /// It has a TCP-level stream, a TLS-level session, and some |
| /// other state/metadata. |
| struct Connection { |
| enclave_id: sgx_enclave_id_t, |
| socket: TcpStream, |
| token: mio::Token, |
| closing: bool, |
| mode: ServerMode, |
| tlsserver_id: usize, |
| back: Option<TcpStream>, |
| sent_http_response: bool, |
| } |
| |
| /// Open a plaintext TCP-level connection for forwarded connections. |
| fn open_back(mode: &ServerMode) -> Option<TcpStream> { |
| match *mode { |
| ServerMode::Forward(ref port) => { |
| let addr = net::SocketAddrV4::new(net::Ipv4Addr::new(127, 0, 0, 1), *port); |
| let conn = TcpStream::connect(&net::SocketAddr::V4(addr)).unwrap(); |
| Some(conn) |
| } |
| _ => None, |
| } |
| } |
| |
| /// This used to be conveniently exposed by mio: map EWOULDBLOCK |
| /// errors to something less-errory. |
| fn try_read(r: io::Result<usize>) -> io::Result<Option<usize>> { |
| match r { |
| Ok(len) => Ok(Some(len)), |
| Err(e) => { |
| if e.kind() == io::ErrorKind::WouldBlock { |
| Ok(None) |
| } else { |
| Err(e) |
| } |
| } |
| } |
| } |
| |
| impl Connection { |
| fn new(enclave_id: sgx_enclave_id_t, |
| socket: TcpStream, |
| token: mio::Token, |
| mode: ServerMode, |
| tlsserver_id: usize) |
| -> Connection { |
| let back = open_back(&mode); |
| Connection { |
| enclave_id: enclave_id, |
| socket: socket, |
| token: token, |
| closing: false, |
| mode: mode, |
| tlsserver_id: tlsserver_id, |
| back: back, |
| sent_http_response: false, |
| } |
| } |
| |
| fn read_tls(&self, buf: &mut [u8]) -> isize { |
| let mut retval = -1; |
| let result = unsafe { |
| tls_server_read(self.enclave_id, |
| &mut retval, |
| self.tlsserver_id, |
| buf.as_ptr() as * mut c_void, |
| buf.len() as c_int) |
| }; |
| match result { |
| sgx_status_t::SGX_SUCCESS => { retval as isize }, |
| _ => { |
| println!("[-] ECALL Enclave [tls_server_wants_read] Failed {}!", result); |
| return -1; |
| }, |
| } |
| } |
| |
| fn write_tls(&self, buf: &[u8]) -> isize { |
| let mut retval = -1; |
| let result = unsafe { |
| tls_server_write(self.enclave_id, |
| &mut retval, |
| self.tlsserver_id, |
| buf.as_ptr() as * const c_void, |
| buf.len() as c_int) |
| }; |
| match result { |
| sgx_status_t::SGX_SUCCESS => { retval as isize }, |
| _ => { |
| println!("[-] ECALL Enclave [tls_server_wants_read] Failed {}!", result); |
| return -1; |
| }, |
| } |
| } |
| |
| fn wants_read(&self) -> bool { |
| let mut retval = -1; |
| let result = unsafe { |
| tls_server_wants_read(self.enclave_id, |
| &mut retval, |
| self.tlsserver_id) |
| }; |
| match result { |
| sgx_status_t::SGX_SUCCESS => {}, |
| _ => { |
| println!("[-] ECALL Enclave [tls_server_wants_read] Failed {}!", result); |
| return false; |
| }, |
| } |
| |
| match retval { |
| 0 => false, |
| _ => true |
| } |
| } |
| |
| fn wants_write(&self) -> bool { |
| let mut retval = -1 ; |
| let result = unsafe { |
| tls_server_wants_write(self.enclave_id, |
| &mut retval, |
| self.tlsserver_id) |
| }; |
| |
| match result { |
| sgx_status_t::SGX_SUCCESS => {}, |
| _ => { |
| println!("[-] ECALL Enclave [tls_server_wants_write] Failed {}!", result); |
| return false; |
| }, |
| } |
| |
| match retval { |
| 0 => false, |
| _ => true |
| } |
| } |
| |
| fn tls_close(&self) { |
| unsafe { |
| tls_server_close(self.enclave_id, self.tlsserver_id) |
| }; |
| } |
| |
| fn send_close_notify(&self) { |
| unsafe { |
| tls_server_send_close(self.enclave_id, self.tlsserver_id); |
| } |
| } |
| |
| /// We're a connection, and we have something to do. |
| fn ready(&mut self, poll: &mut mio::Poll, ev: &mio::Event) { |
| // If we're readable: read some TLS. Then |
| // see if that yielded new plaintext. Then |
| // see if the backend is readable too. |
| if ev.readiness().is_readable() { |
| self.do_tls_read(); |
| self.try_plain_read(); |
| self.try_back_read(); |
| } |
| |
| if ev.readiness().is_writable() { |
| self.do_tls_write(); |
| } |
| |
| if self.closing { |
| self.tls_close(); |
| let _ = self.socket.shutdown(Shutdown::Both); |
| self.close_back(); |
| } else { |
| self.reregister(poll); |
| } |
| } |
| |
| /// Close the backend connection for forwarded sessions. |
| fn close_back(&mut self) { |
| if self.back.is_some() { |
| let back = self.back.as_mut().unwrap(); |
| back.shutdown(Shutdown::Both).unwrap(); |
| } |
| self.back = None; |
| } |
| |
| fn do_tls_read(&mut self) { |
| // Read some TLS data. |
| println!("Read some TLS data."); |
| let mut buf = Vec::new(); |
| let rc = self.read_tls(buf.as_mut_slice()); |
| if rc == -1 { |
| println!("read error {:?}", rc); |
| self.closing = true; |
| return; |
| } |
| } |
| |
| fn try_plain_read(&mut self) { |
| // Read and process all available plaintext. |
| let mut buf = vec![0; BUFFER_SIZE]; |
| let rc = self.read_tls(buf.as_mut_slice()); |
| if rc == -1 { |
| println!("plaintext read failed: {:?}", rc); |
| self.closing = true; |
| return; |
| } |
| buf.resize(rc as usize, 0); |
| if !buf.is_empty() { |
| println!("plaintext read {:?} {:? }", buf.len(), buf); |
| self.incoming_plaintext(&buf); |
| } |
| } |
| |
| fn try_back_read(&mut self) { |
| if self.back.is_none() { |
| return; |
| } |
| |
| println!("try back read"); |
| |
| // Try a non-blocking read. |
| let mut buf = [0u8; BUFFER_SIZE]; |
| let maybe_len = { |
| let back = self.back.as_mut().unwrap(); |
| let rc = try_read(back.read(&mut buf)); |
| if rc.is_err() { |
| println!("backend read failed: {:?}", rc); |
| self.closing = true; |
| return; |
| } |
| rc.unwrap() |
| }; |
| |
| // If we have a successful but empty read, that's an EOF. |
| // Otherwise, we shove the data into the TLS session. |
| match maybe_len { |
| Some(len) if len == 0 => { |
| println!("back eof"); |
| self.closing = true; |
| } |
| Some(len) => { |
| self.write_all(&buf[..len]).unwrap(); |
| } |
| None => {} |
| }; |
| } |
| |
| /// Process some amount of received plaintext. |
| fn incoming_plaintext(&mut self, buf: &[u8]) { |
| match self.mode { |
| ServerMode::Echo => { |
| self.write_all(buf).unwrap(); |
| } |
| ServerMode::Http => { |
| self.send_http_response_once(); |
| } |
| ServerMode::Forward(_) => { |
| self.back.as_mut().unwrap().write_all(buf).unwrap(); |
| } |
| } |
| } |
| |
| fn send_http_response_once(&mut self) { |
| |
| println!("send_http_response_once"); |
| |
| let response = b"HTTP/1.0 200 OK\r\nConnection: close\r\n\r\nHello world from rustls tlsserver\r\n"; |
| if !self.sent_http_response { |
| self.write_all(response) |
| .unwrap(); |
| self.sent_http_response = true; |
| self.send_close_notify(); |
| } |
| } |
| |
| fn do_tls_write(&mut self) { |
| let buf = Vec::new(); |
| let rc = self.write_tls(buf.as_slice()); |
| if rc == -1 { |
| println!("write failed {:?}", rc); |
| self.closing = true; |
| return; |
| } |
| } |
| |
| fn register(&self, poll: &mut mio::Poll) { |
| poll.register(&self.socket, |
| self.token, |
| self.event_set(), |
| mio::PollOpt::level() | mio::PollOpt::oneshot()) |
| .unwrap(); |
| |
| if self.back.is_some() { |
| poll.register(self.back.as_ref().unwrap(), |
| self.token, |
| mio::Ready::readable(), |
| mio::PollOpt::level() | mio::PollOpt::oneshot()) |
| .unwrap(); |
| } |
| } |
| |
| fn reregister(&self, poll: &mut mio::Poll) { |
| poll.reregister(&self.socket, |
| self.token, |
| self.event_set(), |
| mio::PollOpt::level() | mio::PollOpt::oneshot()) |
| .unwrap(); |
| |
| if self.back.is_some() { |
| poll.reregister(self.back.as_ref().unwrap(), |
| self.token, |
| mio::Ready::readable(), |
| mio::PollOpt::level() | mio::PollOpt::oneshot()) |
| .unwrap(); |
| } |
| } |
| |
| /// What IO events we're currently waiting for, |
| /// based on wants_read/wants_write. |
| fn event_set(&self) -> mio::Ready { |
| let rd = self.wants_read(); |
| let wr = self.wants_write(); |
| |
| if rd && wr { |
| mio::Ready::readable() | mio::Ready::writable() |
| } else if wr { |
| mio::Ready::writable() |
| } else { |
| mio::Ready::readable() |
| } |
| } |
| |
| fn is_closed(&self) -> bool { |
| self.closing |
| } |
| } |
| |
| /// We implement `io::Write` and pass through to the TLS session |
| impl io::Write for Connection { |
| fn write(&mut self, bytes: &[u8]) -> io::Result<usize> { |
| Ok(self.write_tls(bytes) as usize) |
| } |
| // unused |
| fn flush(&mut self) -> io::Result<()> { |
| Ok(()) |
| } |
| } |
| |
| impl io::Read for Connection { |
| fn read(&mut self, bytes: &mut [u8]) -> io::Result<usize> { |
| Ok(self.read_tls(bytes) as usize) |
| } |
| } |
| |
| fn main() { |
| |
| let enclave = match init_enclave() { |
| Ok(r) => { |
| println!("[+] Init Enclave Successful {}!", r.geteid()); |
| r |
| }, |
| Err(x) => { |
| println!("[-] Init Enclave Failed {}!", x.as_str()); |
| return; |
| }, |
| }; |
| |
| println!("[+] Test tlsclient in enclave, start!"); |
| |
| let addr: net::SocketAddr = "0.0.0.0:8443".parse().unwrap(); |
| let listener = TcpListener::bind(&addr).expect("cannot listen on port"); |
| |
| let cert = CString::new("end.fullchain").unwrap(); |
| let key = CString::new("end.rsa").unwrap(); |
| |
| let mut poll = mio::Poll::new().unwrap(); |
| poll.register(&listener, |
| LISTENER, |
| mio::Ready::readable(), |
| mio::PollOpt::level()).unwrap(); |
| |
| let mut tlsserv = TlsServer::new(enclave.geteid(), listener, ServerMode::Echo, cert, key); |
| |
| println!("[+] TlsServer new success!"); |
| |
| let mut events = mio::Events::with_capacity(256); |
| |
| 'outer: loop { |
| poll.poll(&mut events, None) |
| .unwrap(); |
| |
| for event in events.iter() { |
| match event.token() { |
| LISTENER => { |
| if !tlsserv.accept(&mut poll) { |
| break 'outer; |
| } |
| } |
| _ => tlsserv.conn_event(&mut poll, &event) |
| } |
| } |
| } |
| |
| println!("[+] Test tlsServer in enclave, done!"); |
| |
| enclave.destroy(); |
| } |