Merge pull request #153 from bradyjoestar/max_conn-in-mio

support max_conn in tlsserver of mio
diff --git a/samplecode/mio/README.md b/samplecode/mio/README.md
index e557d45..53c71d2 100644
--- a/samplecode/mio/README.md
+++ b/samplecode/mio/README.md
@@ -6,7 +6,7 @@
 cd server
 make
 cd bin
-./app
+./app (add '--maxconn 32' if you want to set the max_conn of tlsserver to 32)
 ```
 
 Start client 
diff --git a/samplecode/mio/server/app/build.rs b/samplecode/mio/server/app/build.rs
index de29f98..fb1218d 100644
--- a/samplecode/mio/server/app/build.rs
+++ b/samplecode/mio/server/app/build.rs
@@ -31,9 +31,9 @@
 fn main () {
 
     let sdk_dir = env::var("SGX_SDK")
-                    .unwrap_or_else(|_| "/opt/intel/sgxsdk".to_string());
+        .unwrap_or_else(|_| "/opt/intel/sgxsdk".to_string());
     let is_sim = env::var("SGX_MODE")
-                    .unwrap_or_else(|_| "HW".to_string());
+        .unwrap_or_else(|_| "HW".to_string());
 
     println!("cargo:rustc-link-search=native=../lib");
     println!("cargo:rustc-link-lib=static=Enclave_u");
@@ -44,4 +44,4 @@
         "HW" => println!("cargo:rustc-link-lib=dylib=sgx_urts"),
         _    => println!("cargo:rustc-link-lib=dylib=sgx_urts"), // Treat undefined as HW
     }
-}
+}
\ No newline at end of file
diff --git a/samplecode/mio/server/app/src/main.rs b/samplecode/mio/server/app/src/main.rs
index 78bec15..7bddda4 100644
--- a/samplecode/mio/server/app/src/main.rs
+++ b/samplecode/mio/server/app/src/main.rs
@@ -36,15 +36,17 @@
 use sgx_urts::SgxEnclave;
 
 use std::io::{Read, Write};
+use std::env;
 use std::fs;
 use std::path;
+use std::str::FromStr;
 
 
 static ENCLAVE_FILE: &'static str = "enclave.signed.so";
 static ENCLAVE_TOKEN: &'static str = "enclave.token";
 
 extern {
-    fn run_server(eid: sgx_enclave_id_t) -> sgx_status_t;
+    fn run_server(eid: sgx_enclave_id_t, max_conn: uint8_t) -> sgx_status_t;
 }
 
 fn init_enclave() -> SgxResult<SgxEnclave> {
@@ -116,6 +118,22 @@
 }
 
 fn main() {
+    let mut args: Vec<_> = env::args().collect();
+    //default max_conn is 30
+    let mut max_conn = 30;
+    args.remove(0);
+    while !args.is_empty() {
+        match args.remove(0).as_ref() {
+            "--maxconn" => {
+                max_conn =
+                    uint8_t::from_str(args.remove(0).as_ref()).expect("error parsing argument");
+                println!("max connections is: {}", max_conn);
+            }
+            _ => {
+                panic!("Only --maxconn is accepted");
+            }
+        }
+    }
 
     let enclave = match init_enclave() {
         Ok(r) => {
@@ -130,9 +148,7 @@
 
     println!("[+] Test server in enclave, start!");
 
-    let result = unsafe {
-        run_server(enclave.geteid())
-    };
+    let result = unsafe { run_server(enclave.geteid(), max_conn) };
 
     match result {
         sgx_status_t::SGX_SUCCESS => {},
diff --git a/samplecode/mio/server/enclave/Cargo.toml b/samplecode/mio/server/enclave/Cargo.toml
index e7ed160..54872f9 100644
--- a/samplecode/mio/server/enclave/Cargo.toml
+++ b/samplecode/mio/server/enclave/Cargo.toml
@@ -11,6 +11,7 @@
 default = []
 
 [target.'cfg(not(target_env = "sgx"))'.dependencies]
+sgx_types   = { path = "../../../../sgx_types" }
 sgx_tstd = { path = "../../../../sgx_tstd", features = ["net", "untrusted_fs"] }
 
 [dependencies]
diff --git a/samplecode/mio/server/enclave/Enclave.edl b/samplecode/mio/server/enclave/Enclave.edl
index 7873abe..597afed 100644
--- a/samplecode/mio/server/enclave/Enclave.edl
+++ b/samplecode/mio/server/enclave/Enclave.edl
@@ -40,6 +40,6 @@
 
     trusted {
         /* define ECALLs here. */
-        public void run_server();
+        public void run_server(uint8_t max_conn);
     };
 };
diff --git a/samplecode/mio/server/enclave/src/lib.rs b/samplecode/mio/server/enclave/src/lib.rs
index 99a8627..c634fa6 100644
--- a/samplecode/mio/server/enclave/src/lib.rs
+++ b/samplecode/mio/server/enclave/src/lib.rs
@@ -52,8 +52,11 @@
 extern crate webpki;
 extern crate rustls;
 extern crate mio;
+extern crate sgx_types;
+
 use rustls::{Session, NoClientAuth};
 use mio::net::{TcpListener, TcpStream};
+use sgx_types::uint8_t;
 
 // Token for our listening socket.
 const LISTENER: mio::Token = mio::Token(0);
@@ -208,7 +211,7 @@
             self.do_tls_write();
         }
 
-        if self.closing && !self.tls_session.wants_write() {
+        if self.closing {
             let _ = self.socket.shutdown(Shutdown::Both);
             self.close_back();
             self.closed = true;
@@ -438,7 +441,7 @@
 
 
 #[no_mangle]
-pub extern "C" fn run_server() {
+pub extern "C" fn run_server(max_conn: uint8_t) {
     let addr: net::SocketAddr = "0.0.0.0:8443".parse().unwrap();
     let cert = "end.fullchain";
     let key = "end.rsa";
@@ -466,6 +469,9 @@
         for event in events.iter() {
             match event.token() {
                 LISTENER => {
+                    if tlsserv.connections.len() as u8 == max_conn {
+                        continue;
+                    }
                     if !tlsserv.accept(&mut poll) {
                         break 'outer;
                     }