diff --git a/src/sock/socket_session.cpp b/src/sock/socket_session.cpp
index d8cc1ed8..d073e78c 100644
--- a/src/sock/socket_session.cpp
+++ b/src/sock/socket_session.cpp
@@ -16,7 +16,8 @@ socket_session::socket_session(websocket::stream<beast::tcp_stream> &websocket,
 {
 }
 
-void socket_session::server_run(const unsigned short &port, const std::string &address)
+//port and address will be used to identify from which client the message recieved in the handler
+void socket_session::server_run(const std::uint16_t port, const std::string &address)
 {
     port_ = port;
     address_ = address;
@@ -29,15 +30,17 @@ void socket_session::server_run(const unsigned short &port, const std::string &a
         });
 }
 
-void socket_session::client_run(const unsigned short &port, const std::string &address, error ec)
+//port and address will be used to identify from which server the message recieved in the handler
+void socket_session::client_run(const std::uint16_t port, const std::string &address, error ec)
 {
     port_ = port;
     address_ = address;
 
-    sess_handler_.on_connect(this, ec);
     if (ec)
         return fail(ec, "handshake");
 
+    sess_handler_.on_connect(this);
+
     ws_.async_read(
         buffer_,
         [sp = shared_from_this()](
@@ -53,17 +56,17 @@ void socket_session::fail(error_code ec, char const *what)
         ec == websocket::error::closed)
         return;
 
-    std::cerr << what << ": " << ec.message() << "\n";
+    // std::cerr << what << ": " << ec.message() << "\n";
 }
 
 void socket_session::on_accept(error_code ec)
 {
-    sess_handler_.on_connect(this, ec);
-
     // Handle the error, if any
     if (ec)
         return fail(ec, "accept");
 
+    sess_handler_.on_connect(this);
+
     // Read a message
     ws_.async_read(
         buffer_,
@@ -75,12 +78,16 @@ void socket_session::on_accept(error_code ec)
 
 void socket_session::on_read(error_code ec, std::size_t)
 {
-    auto const string_message = std::make_shared<std::string const>(std::move(beast::buffers_to_string(buffer_.data())));
-    sess_handler_.on_message(this, string_message, ec);
-
     // Handle the error, if any
     if (ec)
+    {
+        //if something goes wrong when trying to read, socket connection will be closed and calling this to inform it to the handler
+        on_close(ec, 1);
         return fail(ec, "read");
+    }
+
+    std::string message = beast::buffers_to_string(buffer_.data());
+    sess_handler_.on_message(this, message);
 
     // Clear the buffer
     buffer_.consume(buffer_.size());
@@ -130,4 +137,26 @@ void socket_session::on_write(error_code ec, std::size_t)
                 sp->on_write(ec, bytes);
             });
 }
+
+void socket_session::close()
+{
+    // Close the WebSocket connection
+    ws_.async_close(websocket::close_code::normal,
+                    [sp = shared_from_this()](
+                        error_code ec) {
+                        sp->on_close(ec, 0);
+                    });
+}
+
+//type will be used identify whether the error is due to failure in closing the web socket or transfer of another exception to this method
+void socket_session::on_close(error_code ec, std::int8_t type)
+{
+    sess_handler_.on_close(this);
+
+    if (type == 1)
+        return;
+
+    if (ec)
+        return fail(ec, "close");
+}
 } // namespace sock
\ No newline at end of file
diff --git a/src/sock/socket_session.hpp b/src/sock/socket_session.hpp
index 4f1f7963..807c9a7c 100644
--- a/src/sock/socket_session.hpp
+++ b/src/sock/socket_session.hpp
@@ -39,17 +39,21 @@ class socket_session : public std::enable_shared_from_this<socket_session>
 
     void on_write(error ec, std::size_t bytes_transferred);
 
+    void on_close(error ec, std::int8_t type);
+
 public:
     socket_session(websocket::stream<beast::tcp_stream> &websocket, socket_session_handler &sess_handler);
 
-    unsigned short port_;
+    std::uint16_t port_;
     std::string address_;
 
-    void server_run(const unsigned short &port, const std::string &address);
-    void client_run(const unsigned short &port, const std::string &address, error ec);
+    void server_run(const std::uint16_t port, const std::string &address);
+    void client_run(const std::uint16_t port, const std::string &address, error ec);
 
     //Used to send message through an active websocket connection
     void send(std::shared_ptr<std::string const> const &ss);
+
+    void close();
 };
 } // namespace sock
 #endif
diff --git a/src/sock/socket_session_handler.hpp b/src/sock/socket_session_handler.hpp
index 3f07c7ad..6f455ecc 100644
--- a/src/sock/socket_session_handler.hpp
+++ b/src/sock/socket_session_handler.hpp
@@ -2,9 +2,6 @@
 #define _SOCK_SESSION_HANDLER_H_
 
 #include "socket_session.hpp"
-#include <boost/beast/core.hpp>
-
-using error = boost::system::error_code;
 
 namespace sock
 {
@@ -18,8 +15,8 @@ class socket_session;
 class socket_session_handler
 {
 public:
-    virtual void on_connect(socket_session *session, error ec) = 0;
-    virtual void on_message(socket_session *session, std::shared_ptr<std::string const> const &message, error ec) = 0;
+    virtual void on_connect(socket_session *session) = 0;
+    virtual void on_message(socket_session *session, const std::string &message) = 0;
     virtual void on_close(socket_session *session) = 0;
 };
 } // namespace sock