Add --max-clients option to tls_proxy and tls_http_server

A clean exit is required in order to get coverage results

3 files changed, 121 insertions, 57 deletions

diff --git a/src/cli/tls_http_server.cpp b/src/cli/tls_http_server.cpp
index 98a51ce83..0ff1632af 100644
--- a/src/cli/tls_http_server.cpp
+++ b/src/cli/tls_http_server.cpp
@@ -15,6 +15,7 @@
 #include <string>
 #include <vector>
 #include <thread>
+#include <atomic>
 
 #define _GLIBCXX_HAVE_GTHR_DEFAULT
 #include <boost/asio.hpp>
@@ -58,6 +59,27 @@ inline void log_exception(const char* where, const std::exception& e)
    std::cout << where << ' ' << e.what() << std::endl;
    }
 
+class ServerStatus
+   {
+   public:
+      ServerStatus(size_t max_clients) : m_max_clients(max_clients), m_clients_serviced(0) {}
+
+      bool should_exit() const
+         {
+         if(m_max_clients == 0)
+            return false;
+
+         return clients_serviced() >= m_max_clients;
+         }
+
+      void client_serviced() { m_clients_serviced++; }
+
+      size_t clients_serviced() const { return m_clients_serviced.load(); }
+   private:
+      size_t m_max_clients;
+      std::atomic<size_t> m_clients_serviced;
+   };
+
 /*
 * This is an incomplete and highly buggy HTTP request parser. It is just
 * barely sufficient to handle a GET request sent by a browser.
@@ -198,8 +220,7 @@ class TLS_Asio_HTTP_Session final : public boost::enable_shared_from_this<TLS_As
          {
          if(error)
             {
-            stop();
-            return;
+            return stop();
             }
 
          try
@@ -209,8 +230,7 @@ class TLS_Asio_HTTP_Session final : public boost::enable_shared_from_this<TLS_As
          catch(Botan::Exception& e)
             {
             log_exception("TLS connection failed", e);
-            stop();
-            return;
+            return stop();
             }
 
          m_client_socket.async_read_some(
@@ -226,8 +246,7 @@ class TLS_Asio_HTTP_Session final : public boost::enable_shared_from_this<TLS_As
          {
          if(error)
             {
-            stop();
-            return;
+            return stop();
             }
 
          m_s2c.clear();
@@ -272,11 +291,7 @@ class TLS_Asio_HTTP_Session final : public boost::enable_shared_from_this<TLS_As
       void handle_http_request(const HTTP_Parser::Request& request) override
          {
          std::ostringstream response;
-         if(request.verb() != "GET")
-            {
-            response << "HTTP/1.0 405 Method Not Allowed\r\n\r\nNo POST for you";
-            }
-         else
+         if(request.verb() == "GET")
             {
             if(request.location() == "/" || request.location() == "/status")
                {
@@ -294,9 +309,13 @@ class TLS_Asio_HTTP_Session final : public boost::enable_shared_from_this<TLS_As
                }
             else
                {
-               response << "HTTP/1.0 404 Not Found\r\n\r\nSorry, no";
+               response << "HTTP/1.0 404 Not Found\r\n\r\n";
                }
             }
+         else
+            {
+            response << "HTTP/1.0 405 Method Not Allowed\r\n\r\n";
+            }
 
          const std::string response_str = response.str();
          m_tls.send(response_str);
@@ -420,11 +439,13 @@ class TLS_Asio_HTTP_Server final
          boost::asio::io_service& io, unsigned short port,
          Botan::Credentials_Manager& creds,
          Botan::TLS::Policy& policy,
-         Botan::TLS::Session_Manager& session_mgr)
+         Botan::TLS::Session_Manager& session_mgr,
+         size_t max_clients)
          : m_acceptor(io, tcp::endpoint(tcp::v4(), port))
          , m_creds(creds)
          , m_policy(policy)
          , m_session_manager(session_mgr)
+         , m_status(max_clients)
          {
          session::pointer new_session = make_session();
 
@@ -455,13 +476,18 @@ class TLS_Asio_HTTP_Server final
             new_session->start();
             new_session = make_session();
 
-            m_acceptor.async_accept(
-               new_session->client_socket(),
-               boost::bind(
-                  &TLS_Asio_HTTP_Server::handle_accept,
-                  this,
-                  new_session,
-                  boost::asio::placeholders::error));
+            m_status.client_serviced();
+
+            if(m_status.should_exit() == false)
+               {
+               m_acceptor.async_accept(
+                  new_session->client_socket(),
+                  boost::bind(
+                     &TLS_Asio_HTTP_Server::handle_accept,
+                     this,
+                     new_session,
+                     boost::asio::placeholders::error));
+               }
             }
          }
 
@@ -470,6 +496,7 @@ class TLS_Asio_HTTP_Server final
       Botan::Credentials_Manager& m_creds;
       Botan::TLS::Policy& m_policy;
       Botan::TLS::Session_Manager& m_session_manager;
+      ServerStatus m_status;
    };
 
 }
@@ -478,7 +505,7 @@ class TLS_HTTP_Server final : public Command
    {
    public:
       TLS_HTTP_Server() : Command("tls_http_server server_cert server_key "
-                                  "--port=443 --policy= --threads=0 "
+                                  "--port=443 --policy= --threads=0 --max-clients=0 "
                                   "--session-db= --session-db-pass=") {}
 
       std::string group() const override
@@ -508,6 +535,7 @@ class TLS_HTTP_Server final : public Command
          const std::string server_key = get_arg("server_key");
 
          const size_t num_threads = thread_count();
+         const size_t max_clients = get_arg_sz("max-clients");
 
          Basic_Credentials_Manager creds(rng(), server_crt, server_key);
 
@@ -551,7 +579,7 @@ class TLS_HTTP_Server final : public Command
 
          boost::asio::io_service io;
 
-         TLS_Asio_HTTP_Server server(io, listen_port, creds, *policy, *session_mgr);
+         TLS_Asio_HTTP_Server server(io, listen_port, creds, *policy, *session_mgr, max_clients);
 
          std::vector<std::shared_ptr<std::thread>> threads;
 
diff --git a/src/cli/tls_proxy.cpp b/src/cli/tls_proxy.cpp
index 2770a7c48..596e4105c 100644
--- a/src/cli/tls_proxy.cpp
+++ b/src/cli/tls_proxy.cpp
@@ -14,6 +14,7 @@
 #include <string>
 #include <vector>
 #include <thread>
+#include <atomic>
 
 #define _GLIBCXX_HAVE_GTHR_DEFAULT
 #include <boost/asio.hpp>
@@ -45,17 +46,17 @@ namespace {
 
 using boost::asio::ip::tcp;
 
-inline void log_exception(const char* where, const std::exception& e)
+void log_exception(const char* where, const std::exception& e)
    {
    std::cout << where << ' ' << e.what() << std::endl;
    }
 
-inline void log_error(const char* where, const boost::system::error_code& error)
+void log_error(const char* where, const boost::system::error_code& error)
    {
    std::cout << where << ' ' << error.message() << std::endl;
    }
 
-inline void log_binary_message(const char* where, const uint8_t buf[], size_t buf_len)
+void log_binary_message(const char* where, const uint8_t buf[], size_t buf_len)
    {
    BOTAN_UNUSED(where, buf, buf_len);
    //std::cout << where << ' ' << Botan::hex_encode(buf, buf_len) << std::endl;
@@ -68,11 +69,32 @@ void log_text_message(const char* where,  const uint8_t buf[], size_t buf_len)
    //std::cout << where << ' ' << std::string(c, c + buf_len)  << std::endl;
    }
 
+class ServerStatus
+   {
+   public:
+      ServerStatus(size_t max_clients) : m_max_clients(max_clients), m_clients_serviced(0) {}
+
+      bool should_exit() const
+         {
+         if(m_max_clients == 0)
+            return false;
+
+         return clients_serviced() >= m_max_clients;
+         }
+
+      void client_serviced() { m_clients_serviced++; }
+
+      size_t clients_serviced() const { return m_clients_serviced.load(); }
+   private:
+      size_t m_max_clients;
+      std::atomic<size_t> m_clients_serviced;
+   };
+
 class tls_proxy_session final : public boost::enable_shared_from_this<tls_proxy_session>,
                                 public Botan::TLS::Callbacks
    {
    public:
-      enum { readbuf_size = 4 * 1024 };
+      enum { readbuf_size = 17 * 1024 };
 
       typedef boost::shared_ptr<tls_proxy_session> pointer;
 
@@ -106,12 +128,16 @@ class tls_proxy_session final : public boost::enable_shared_from_this<tls_proxy_
 
       void stop()
          {
-         /*
-         Don't need to talk to the server anymore
-         Client socket is closed during write callback
-         */
-         m_server_socket.close();
-         m_tls.close();
+         if(m_is_closed == false)
+            {
+            /*
+            Don't need to talk to the server anymore
+            Client socket is closed during write callback
+            */
+            m_server_socket.close();
+            m_tls.close();
+            m_is_closed = true;
+            }
          }
 
    private:
@@ -290,15 +316,8 @@ class tls_proxy_session final : public boost::enable_shared_from_this<tls_proxy_
 
       bool tls_session_established(const Botan::TLS::Session& session) override
          {
-         //std::cout << "Handshake from client complete" << std::endl;
-
          m_hostname = session.server_info().hostname();
 
-         if(m_hostname != "")
-            {
-            std::cout << "Client requested hostname '" << m_hostname << "'" << std::endl;
-            }
-
          auto onConnect = [this](boost::system::error_code ec, tcp::resolver::iterator /*endpoint*/)
             {
             if(ec)
@@ -320,10 +339,6 @@ class tls_proxy_session final : public boost::enable_shared_from_this<tls_proxy_
             m_tls.close();
             return;
             }
-         else
-            {
-            std::cout << "Alert " << alert.type_string() << std::endl;
-            }
          }
 
       boost::asio::io_service::strand m_strand;
@@ -344,6 +359,8 @@ class tls_proxy_session final : public boost::enable_shared_from_this<tls_proxy_
       std::vector<uint8_t> m_s2p;
       std::vector<uint8_t> m_p2s;
       std::vector<uint8_t> m_p2s_pending;
+
+      bool m_is_closed = false;
    };
 
 class tls_proxy_server final
@@ -356,12 +373,14 @@ class tls_proxy_server final
          tcp::resolver::iterator endpoints,
          Botan::Credentials_Manager& creds,
          Botan::TLS::Policy& policy,
-         Botan::TLS::Session_Manager& session_mgr)
+         Botan::TLS::Session_Manager& session_mgr,
+         size_t max_clients)
          : m_acceptor(io, tcp::endpoint(tcp::v4(), port))
          , m_server_endpoints(endpoints)
          , m_creds(creds)
          , m_policy(policy)
          , m_session_manager(session_mgr)
+         , m_status(max_clients)
          {
          session::pointer new_session = make_session();
 
@@ -393,13 +412,18 @@ class tls_proxy_server final
             new_session->start();
             new_session = make_session();
 
-            m_acceptor.async_accept(
-               new_session->client_socket(),
-               boost::bind(
-                  &tls_proxy_server::handle_accept,
-                  this,
-                  new_session,
-                  boost::asio::placeholders::error));
+            m_status.client_serviced();
+
+            if(m_status.should_exit() == false)
+               {
+               m_acceptor.async_accept(
+                  new_session->client_socket(),
+                  boost::bind(
+                     &tls_proxy_server::handle_accept,
+                     this,
+                     new_session,
+                     boost::asio::placeholders::error));
+               }
             }
          }
 
@@ -409,6 +433,7 @@ class tls_proxy_server final
       Botan::Credentials_Manager& m_creds;
       Botan::TLS::Policy& m_policy;
       Botan::TLS::Session_Manager& m_session_manager;
+      ServerStatus m_status;
    };
 
 }
@@ -417,7 +442,7 @@ class TLS_Proxy final : public Command
    {
    public:
       TLS_Proxy() : Command("tls_proxy listen_port target_host target_port server_cert server_key "
-                               "--threads=0 --session-db= --session-db-pass=") {}
+                               "--threads=0 --max-clients=0 --session-db= --session-db-pass=") {}
 
       std::string group() const override
          {
@@ -448,6 +473,7 @@ class TLS_Proxy final : public Command
          const std::string server_key = get_arg("server_key");
 
          const size_t num_threads = thread_count();
+         const size_t max_clients = get_arg_sz("max-clients");
 
          Basic_Credentials_Manager creds(rng(), server_crt, server_key);
 
@@ -474,7 +500,7 @@ class TLS_Proxy final : public Command
             session_mgr.reset(new Botan::TLS::Session_Manager_In_Memory(rng()));
             }
 
-         tls_proxy_server server(io, listen_port, server_endpoint_iterator, creds, policy, *session_mgr);
+         tls_proxy_server server(io, listen_port, server_endpoint_iterator, creds, policy, *session_mgr, max_clients);
 
          std::vector<std::shared_ptr<std::thread>> threads;
 
diff --git a/src/scripts/test_cli.py b/src/scripts/test_cli.py
index bb9bb8c0a..f7c298f39 100755
--- a/src/scripts/test_cli.py
+++ b/src/scripts/test_cli.py
@@ -1,5 +1,11 @@
 #!/usr/bin/python
 
+"""
+(C) 2018,2019 Jack Lloyd
+
+Botan is released under the Simplified BSD License (see license.txt)
+"""
+
 import subprocess
 import sys
 import os
@@ -646,7 +652,7 @@ def cli_tls_http_server_tests(tmp_dir):
 
     test_cli("sign_cert", "%s %s %s --output=%s" % (ca_cert, priv_key, crt_req, server_cert))
 
-    tls_server = subprocess.Popen([CLI_PATH, 'tls_http_server',
+    tls_server = subprocess.Popen([CLI_PATH, 'tls_http_server', '--max-clients=1',
                                    '--port=%d' % (server_port), server_cert, priv_key],
                                   stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
@@ -665,7 +671,9 @@ def cli_tls_http_server_tests(tmp_dir):
     if body.find('TLS negotiation with Botan 2.') < 0:
         logging.error('Unexpected response body')
 
-    tls_server.terminate()
+    rc = tls_server.wait()
+    if rc != 0:
+        logging.error("Unexpected return code from https_server %d", rc)
 
 def cli_tls_proxy_tests(tmp_dir):
     # pylint: disable=too-many-locals
@@ -714,7 +722,7 @@ def cli_tls_proxy_tests(tmp_dir):
     test_cli("sign_cert", "%s %s %s --output=%s" % (ca_cert, priv_key, crt_req, server_cert))
 
     tls_proxy = subprocess.Popen([CLI_PATH, 'tls_proxy', str(proxy_port), '127.0.0.1', str(server_port),
-                                  server_cert, priv_key, '--output=/tmp/proxy.err'],
+                                  server_cert, priv_key, '--output=/tmp/proxy.err', '--max-clients=1'],
                                  stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
     time.sleep(.5)
@@ -751,7 +759,9 @@ def cli_tls_proxy_tests(tmp_dir):
     if body != server_response:
         logging.error('Unexpected response from server %s' % (body))
 
-    tls_proxy.terminate()
+    rc = tls_proxy.wait()
+    if rc != 0:
+        logging.error('Unexpected return code %d', rc)
 
 def cli_trust_root_tests(tmp_dir):
     pem_file = os.path.join(tmp_dir, 'pems')