1 files changed, 91 insertions, 32 deletions

diff --git a/client-qt/udm-client-qt/client_session.cpp b/client-qt/udm-client-qt/client_session.cpp
index d01d0f5..b5c6181 100644
--- a/client-qt/udm-client-qt/client_session.cpp
+++ b/client-qt/udm-client-qt/client_session.cpp
@@ -49,17 +49,93 @@ std::shared_ptr<char*> pack_msg(client_msg *msg, int *size_)
 }
 
 
+class socket_wraper
+{
+public:
+    socket_wraper(boost::asio::ssl::stream<boost::asio::ip::tcp::socket> *s) : is_ssl(true)
+    {
+        socket_ssl_ = s;
+    }
+
+    socket_wraper(tcp::socket *s) : is_ssl(false)
+    {
+        socket_ = s;
+    }
+
+    void operator=(boost::asio::ssl::stream<boost::asio::ip::tcp::socket> *s)
+    {
+        socket_ssl_ = s;
+        is_ssl = true;
+    }
+    void operator=(tcp::socket *s)
+    {
+        socket_ = s;
+        is_ssl = false;
+    }
+
+    boost::asio::ip::tcp::socket& get_socket()
+    {
+        return *socket_;
+    }
+
+    boost::asio::ssl::stream<boost::asio::ip::tcp::socket>& get_ssl_socket()
+    {
+        return *socket_ssl_;
+    }
+    template<typename MutableBufferSequence> std::size_t read_some(const MutableBufferSequence & buffers)
+    {
+        if(is_ssl)
+            return socket_ssl_->read_some(buffers);
+        else
+            return socket_->read_some(buffers);
+    }
+    template<typename MutableBufferSequence> std::size_t read_some(const MutableBufferSequence & buffers, boost::system::error_code & ec)
+    {
+        if(is_ssl)
+            return socket_ssl_->read_some(buffers, ec);
+        else
+            return socket_->read_some(buffers, ec);
+    }
+
+    template<typename MutableBufferSequence, typename ReadHandler>  void async_read_some(const MutableBufferSequence & buffers, ReadHandler handler)
+    {
+        if(is_ssl)
+            socket_ssl_->async_read_some(buffers, handler);
+        else
+            socket_->async_read_some(buffers, handler);
+    }
+
+    template<typename ConstBufferSequence, typename WriteHandler> void async_write_some(const ConstBufferSequence & buffers, WriteHandler handler)
+    {
+        if(is_ssl)
+            socket_ssl_->async_write_some(buffers, handler);
+        else
+            socket_->async_write_some(buffers, handler);
+    }
+    ~socket_wraper()
+    {
+        if(is_ssl)
+            delete socket_ssl_;
+        else
+            delete socket_;
+    }
+private:
+    bool is_ssl;
+    tcp::socket *socket_;
+    boost::asio::ssl::stream<boost::asio::ip::tcp::socket> *socket_ssl_;
+
+};
+
+
 client_session::client_session(QObject *parent, boost::asio::io_service *io_service) : QObject(parent), io_service_(*io_service)
 {
-    socket_ssl_ = nullptr;
-    ssl_enabled = false;
 }
 
 
-tcp::socket& client_session::socket()
+/*tcp::socket& client_session::socket()
 {
     return *socket_;
-}
+} */
 
 void client_session::client_connect(QString host, QString password, int port)
 {
@@ -79,12 +155,9 @@ void client_session::client_connect(QString host, QString password, int port)
     boost::asio::ip::tcp::resolver::query query(host.toUtf8().data(), port_buf);
     boost::asio::ip::tcp::resolver::iterator it = r.resolve(query);
 
-    socket_ = new tcp::socket(io_service_);
+    socket_ = new socket_wraper(new tcp::socket(io_service_));
 
-    if(!ssl_enabled)
-        boost::asio::async_connect(*socket_, it, boost::bind(&client_session::handle_connect, this, boost::asio::placeholders::error));
-    else
-        boost::asio::async_connect(socket_ssl_->lowest_layer(), it, boost::bind(&client_session::handle_connect, this, boost::asio::placeholders::error));
+    boost::asio::async_connect(socket_->get_socket(), it, boost::bind(&client_session::handle_connect, this, boost::asio::placeholders::error));
 //    socket_.async_connect(ep, boost::bind(&client_session::handle_connect, this, boost::asio::placeholders::error));
     io_service_.run();
 }
@@ -93,13 +166,14 @@ void client_session::client_connect_ssl(QString host, QString password, int port
 {
     ssl_enabled = true;
     this->pasword = password;
-    boost::asio::ssl::context ctx(boost::asio::ssl::context::tlsv12);
+    boost::asio::ssl::context ctx(boost::asio::ssl::context::tlsv12_client);
     ctx.load_verify_file(ssl_ca.toUtf8().data());
     ctx.use_certificate_file(ssl_crt.toUtf8().data(), boost::asio::ssl::context::pem);
     ctx.use_private_key_file(ssl_key.toUtf8().data(), boost::asio::ssl::context::pem);
-    socket_ssl_ = new boost::asio::ssl::stream<boost::asio::ip::tcp::socket>(io_service_, ctx);
+    auto socket_ssl_ = new boost::asio::ssl::stream<boost::asio::ip::tcp::socket>(io_service_, ctx);
     socket_ssl_->set_verify_mode(boost::asio::ssl::verify_peer);
     socket_ssl_->set_verify_callback(boost::bind(&client_session::verify_certificate, this, _1, _2 ));
+    socket_ = new socket_wraper(socket_ssl_);
 
     char port_buf[6];
     snprintf(port_buf, 5, "%d", port);
@@ -107,7 +181,7 @@ void client_session::client_connect_ssl(QString host, QString password, int port
     boost::asio::ip::tcp::resolver::query query(host.toUtf8().data(), port_buf);
     boost::asio::ip::tcp::resolver::iterator it = r.resolve(query);
 
-    boost::asio::async_connect(socket_ssl_->lowest_layer(), it, boost::bind(&client_session::handle_connect, this, boost::asio::placeholders::error));
+    boost::asio::async_connect(socket_->get_ssl_socket().lowest_layer(), it, boost::bind(&client_session::handle_connect, this, boost::asio::placeholders::error));
 }
 
 bool client_session::verify_certificate(bool preverified, boost::asio::ssl::verify_context& /*&ctx*/)
@@ -144,10 +218,7 @@ void client_session::handle_connect(const boost::system::error_code &e)
         emit client_connected(true, QString::fromStdString(e.message()));
         //listen for messages from server
         recv_data_ = new char[4];
-        if(!ssl_enabled)
-            boost::asio::async_read(*socket_, boost::asio::buffer(recv_data_, 4), boost::bind(&client_session::handle_read, this, boost::asio::placeholders::error, boost::asio::placeholders::bytes_transferred));
-        else
-            boost::asio::async_read(*socket_ssl_, boost::asio::buffer(recv_data_, 4), boost::bind(&client_session::handle_read, this, boost::asio::placeholders::error, boost::asio::placeholders::bytes_transferred));
+        boost::asio::async_read(*socket_, boost::asio::buffer(recv_data_, 4), boost::bind(&client_session::handle_read, this, boost::asio::placeholders::error, boost::asio::placeholders::bytes_transferred));
     }
 }
 
@@ -161,10 +232,7 @@ void client_session::handle_read(const boost::system::error_code& error, size_t
         delete [] recv_data_;
         char *buf = new char[size];
         boost::system::error_code ec;
-        if(!ssl_enabled)
-        socket_->read_some(boost::asio::buffer(buf, size), ec);
-        else
-            socket_ssl_->read_some(boost::asio::buffer(buf, size), ec);
+        boost::asio::read(*socket_, boost::asio::buffer(buf, size), ec);
         if(ec)
         {
 
@@ -189,10 +257,7 @@ void client_session::handle_read(const boost::system::error_code& error, size_t
             //check if someone connected to signal and delete message if noone
             //this will be limited to only one client connection because client will need to delete message manually
             recv_data_ = new char[4];
-            if(!ssl_enabled)
-                boost::asio::async_read(*socket_, boost::asio::buffer(recv_data_, 4), boost::bind(&client_session::handle_read, this, boost::asio::placeholders::error, boost::asio::placeholders::bytes_transferred));
-            else
-                boost::asio::async_read(*socket_ssl_, boost::asio::buffer(recv_data_, 4), boost::bind(&client_session::handle_read, this, boost::asio::placeholders::error, boost::asio::placeholders::bytes_transferred));
+            boost::asio::async_read(*socket_, boost::asio::buffer(recv_data_, 4), boost::bind(&client_session::handle_read, this, boost::asio::placeholders::error, boost::asio::placeholders::bytes_transferred));
         }
 //        if(!parsed)
         else
@@ -212,10 +277,7 @@ void client_session::send_message(client_msg &msg)
 {
     int size = 0;
     std::shared_ptr<char*> ptr = pack_msg(&msg, &size);
-    if(!ssl_enabled)
-        boost::asio::async_write(*socket_, boost::asio::buffer(*ptr, size), boost::bind(&client_session::handle_write, this, boost::asio::placeholders::error));
-    else
-        boost::asio::async_write(*socket_ssl_, boost::asio::buffer(*ptr, size), boost::bind(&client_session::handle_write, this, boost::asio::placeholders::error));
+    boost::asio::async_write(*socket_, boost::asio::buffer(*ptr, size), boost::bind(&client_session::handle_write, this, boost::asio::placeholders::error));
 }
 
 void client_session::handle_write(const boost::system::error_code& error)
@@ -231,10 +293,7 @@ client_session::~client_session()
 {
     //TODO: correct thread termination, reimplement
     io_service_.stop();
-    if(ssl_enabled)
-        delete socket_ssl_;
-    else
-        delete socket_;
+    delete socket_;
     boost::asio::io_service *s = &io_service_;
     delete s;
     //emit terminate_thread();