fix resource leak

trojan-gfw · Mar 17, 2019 · c885117 · c885117
1 parent fe1f297
commit c885117
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 26 deletions.
diff --git a/src/clientsession.cpp b/src/clientsession.cpp
@@ -37,8 +37,13 @@ tcp::socket& ClientSession::accept_socket() {
 }
 
 void ClientSession::start() {
+    boost::system::error_code ec;
     start_time = time(NULL);
-    in_endpoint = in_socket.remote_endpoint();
+    in_endpoint = in_socket.remote_endpoint(ec);
+    if (ec) {
+        destroy();
+        return;
+    }
     auto ssl = out_socket.native_handle();
     if (config.ssl.sni != "") {
         SSL_set_tlsext_host_name(ssl, config.ssl.sni.c_str());
@@ -159,7 +164,12 @@ void ClientSession::in_recv(const string &data) {
             is_udp = req.command == TrojanRequest::UDP_ASSOCIATE;
             if (is_udp) {
                 udp::endpoint bindpoint(in_socket.local_endpoint().address(), 0);
-                udp_socket.open(bindpoint.protocol());
+                boost::system::error_code ec;
+                udp_socket.open(bindpoint.protocol(), ec);
+                if (ec) {
+                    destroy();
+                    return;
+                }
                 udp_socket.bind(bindpoint);
                 Log::log_with_endpoint(in_endpoint, "requested UDP associate to " + req.address.address + ':' + to_string(req.address.port) + ", open UDP socket " + udp_socket.local_endpoint().address().to_string() + ':' + to_string(udp_socket.local_endpoint().port()) + " for relay", Log::INFO);
                 in_async_write(string("\x05\x00\x00", 3) + SOCKS5Address::generate(udp_socket.local_endpoint()));
@@ -218,7 +228,12 @@ void ClientSession::in_sent() {
                     destroy();
                     return;
                 }
-                out_socket.lowest_layer().open(iterator->endpoint().protocol());
+                boost::system::error_code ec;
+                out_socket.lowest_layer().open(iterator->endpoint().protocol(), ec);
+                if (ec) {
+                    destroy();
+                    return;
+                }
                 if (config.tcp.no_delay) {
                     out_socket.lowest_layer().set_option(tcp::no_delay(true));
                 }
@@ -372,11 +387,11 @@ void ClientSession::destroy() {
     }
     if (out_socket.lowest_layer().is_open()) {
         out_socket.lowest_layer().cancel(ec);
-        auto self = shared_from_this();
-        out_socket.async_shutdown([this, self](const boost::system::error_code) {
-            boost::system::error_code ec;
-            out_socket.lowest_layer().shutdown(tcp::socket::shutdown_both, ec);
-            out_socket.lowest_layer().close(ec);
-        });
+        // only do unidirectional shutdown and don't wait for other side's close_notify
+        // a.k.a. call SSL_shutdown() once and discard its return value
+        ::SSL_set_shutdown(out_socket.native_handle(), SSL_RECEIVED_SHUTDOWN);
+        out_socket.shutdown(ec);
+        out_socket.lowest_layer().shutdown(tcp::socket::shutdown_both, ec);
+        out_socket.lowest_layer().close(ec);
     }
 }
diff --git a/src/forwardsession.cpp b/src/forwardsession.cpp
@@ -36,8 +36,13 @@ tcp::socket& ForwardSession::accept_socket() {
 }
 
 void ForwardSession::start() {
+    boost::system::error_code ec;
     start_time = time(NULL);
-    in_endpoint = in_socket.remote_endpoint();
+    in_endpoint = in_socket.remote_endpoint(ec);
+    if (ec) {
+        destroy();
+        return;
+    }
     auto ssl = out_socket.native_handle();
     if (config.ssl.sni != "") {
         SSL_set_tlsext_host_name(ssl, config.ssl.sni.c_str());
@@ -63,7 +68,12 @@ void ForwardSession::start() {
             destroy();
             return;
         }
-        out_socket.lowest_layer().open(iterator->endpoint().protocol());
+        boost::system::error_code ec;
+        out_socket.lowest_layer().open(iterator->endpoint().protocol(), ec);
+        if (ec) {
+            destroy();
+            return;
+        }
         if (config.tcp.no_delay) {
             out_socket.lowest_layer().set_option(tcp::no_delay(true));
         }
@@ -199,11 +209,11 @@ void ForwardSession::destroy() {
     }
     if (out_socket.lowest_layer().is_open()) {
         out_socket.lowest_layer().cancel(ec);
-        auto self = shared_from_this();
-        out_socket.async_shutdown([this, self](const boost::system::error_code) {
-            boost::system::error_code ec;
-            out_socket.lowest_layer().shutdown(tcp::socket::shutdown_both, ec);
-            out_socket.lowest_layer().close(ec);
-        });
+        // only do unidirectional shutdown and don't wait for other side's close_notify
+        // a.k.a. call SSL_shutdown() once and discard its return value
+        ::SSL_set_shutdown(out_socket.native_handle(), SSL_RECEIVED_SHUTDOWN);
+        out_socket.shutdown(ec);
+        out_socket.lowest_layer().shutdown(tcp::socket::shutdown_both, ec);
+        out_socket.lowest_layer().close(ec);
     }
 }
diff --git a/src/serversession.cpp b/src/serversession.cpp
@@ -38,8 +38,13 @@ tcp::socket& ServerSession::accept_socket() {
 }
 
 void ServerSession::start() {
+    boost::system::error_code ec;
     start_time = time(NULL);
-    in_endpoint = in_socket.lowest_layer().remote_endpoint();
+    in_endpoint = in_socket.lowest_layer().remote_endpoint(ec);
+    if (ec) {
+        destroy();
+        return;
+    }
     auto self = shared_from_this();
     in_socket.async_handshake(stream_base::server, [this, self](const boost::system::error_code error) {
         if (error) {
@@ -175,7 +180,12 @@ void ServerSession::in_recv(const string &data) {
                     }
                 }
             }
-            out_socket.open(iterator->endpoint().protocol());
+            boost::system::error_code ec;
+            out_socket.open(iterator->endpoint().protocol(), ec);
+            if (ec) {
+                destroy();
+                return;
+            }
             if (config.tcp.no_delay) {
                 out_socket.set_option(tcp::no_delay(true));
             }
@@ -278,7 +288,12 @@ void ServerSession::udp_sent() {
             }
             if (!udp_socket.is_open()) {
                 auto protocol = iterator->endpoint().protocol();
-                udp_socket.open(protocol);
+                boost::system::error_code ec;
+                udp_socket.open(protocol, ec);
+                if (ec) {
+                    destroy();
+                    return;
+                }
                 udp_socket.bind(udp::endpoint(protocol, 0));
                 udp_async_read();
             }
@@ -311,11 +326,11 @@ void ServerSession::destroy() {
     }
     if (in_socket.lowest_layer().is_open()) {
         in_socket.lowest_layer().cancel(ec);
-        auto self = shared_from_this();
-        in_socket.async_shutdown([this, self](const boost::system::error_code) {
-            boost::system::error_code ec;
-            in_socket.lowest_layer().shutdown(tcp::socket::shutdown_both, ec);
-            in_socket.lowest_layer().close(ec);
-        });
+        // only do unidirectional shutdown and don't wait for other side's close_notify
+        // a.k.a. call SSL_shutdown() once and discard its return value
+        ::SSL_set_shutdown(in_socket.native_handle(), SSL_RECEIVED_SHUTDOWN);
+        in_socket.shutdown(ec);
+        in_socket.lowest_layer().shutdown(tcp::socket::shutdown_both, ec);
+        in_socket.lowest_layer().close(ec);
     }
 }
diff --git a/src/service.cpp b/src/service.cpp
@@ -207,6 +207,8 @@ void Service::run() {
 }
 
 void Service::stop() {
+    boost::system::error_code ec;
+    socket_acceptor.cancel(ec);
     io_service.stop();
 }
 
@@ -220,6 +222,10 @@ void Service::async_accept() {
         session = make_shared<ClientSession>(config, io_service, ssl_context);
     }
     socket_acceptor.async_accept(session->accept_socket(), [this, session](const boost::system::error_code error) {
+        if (error == boost::asio::error::operation_aborted) {
+            // got cancel signal, stop calling myself
+            return;
+        }
         if (!error) {
             boost::system::error_code ec;
             auto endpoint = session->accept_socket().remote_endpoint(ec);