From 22dfe61fe16a83ae09e70626e311e0271d31da19 Mon Sep 17 00:00:00 2001
From: Matej Ferencevic <matej.ferencevic@memgraph.io>
Date: Fri, 7 Apr 2017 09:52:19 +0200
Subject: [PATCH] Added socket write verification to bolt buffer.

Reviewers: buda

Reviewed By: buda

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D199
---
 src/communication/bolt/v1/bolt_exception.hpp  | 11 +++++++
 .../v1/encoder/chunked_encoder_buffer.hpp     |  4 ++-
 src/communication/bolt/v1/states/executor.hpp | 33 +++++++++++++++----
 3 files changed, 41 insertions(+), 7 deletions(-)
 create mode 100644 src/communication/bolt/v1/bolt_exception.hpp

diff --git a/src/communication/bolt/v1/bolt_exception.hpp b/src/communication/bolt/v1/bolt_exception.hpp
new file mode 100644
index 000000000..883c4986d
--- /dev/null
+++ b/src/communication/bolt/v1/bolt_exception.hpp
@@ -0,0 +1,11 @@
+#pragma once
+
+#include "utils/exceptions/basic_exception.hpp"
+
+namespace communication::bolt {
+
+class BoltException : public BasicException {
+ public:
+  using BasicException::BasicException;
+};
+}
diff --git a/src/communication/bolt/v1/encoder/chunked_encoder_buffer.hpp b/src/communication/bolt/v1/encoder/chunked_encoder_buffer.hpp
index a411da2af..c12554171 100644
--- a/src/communication/bolt/v1/encoder/chunked_encoder_buffer.hpp
+++ b/src/communication/bolt/v1/encoder/chunked_encoder_buffer.hpp
@@ -5,6 +5,7 @@
 #include <memory>
 #include <vector>
 
+#include "communication/bolt/v1/bolt_exception.hpp"
 #include "communication/bolt/v1/constants.hpp"
 #include "logging/loggable.hpp"
 #include "utils/bswap.hpp"
@@ -112,7 +113,8 @@ class ChunkedEncoderBuffer : public Loggable {
     if (size_ == 0) return;
 
     // Flush the whole buffer.
-    socket_.Write(buffer_.data(), size_);
+    bool written = socket_.Write(buffer_.data(), size_);
+    if (!written) throw BoltException("Socket write failed!");
     logger.trace("Flushed {} bytes.", size_);
 
     // Cleanup.
diff --git a/src/communication/bolt/v1/states/executor.hpp b/src/communication/bolt/v1/states/executor.hpp
index 204cf071f..29b917aa5 100644
--- a/src/communication/bolt/v1/states/executor.hpp
+++ b/src/communication/bolt/v1/states/executor.hpp
@@ -2,6 +2,7 @@
 
 #include <string>
 
+#include "communication/bolt/v1/bolt_exception.hpp"
 #include "communication/bolt/v1/messaging/codes.hpp"
 #include "communication/bolt/v1/state.hpp"
 #include "logging/default.hpp"
@@ -14,6 +15,17 @@ struct Query {
   std::string statement;
 };
 
+template <typename Session>
+void StateExecutorFailure(Session &session, Logger &logger,
+                          const std::map<std::string, TypedValue> &metadata) {
+  try {
+    session.encoder_.MessageFailure(metadata);
+  } catch (const BoltException &e) {
+    logger.debug("MessageFailure failed because: {}", e.what());
+    session.Close();
+  }
+}
+
 /**
  * TODO (mferencevic): finish & document
  */
@@ -41,7 +53,8 @@ State StateExecutorRun(Session &session) {
 
       if (!is_successfully_executed) {
         db_accessor->abort();
-        session.encoder_.MessageFailure(
+        StateExecutorFailure<Session>(
+            session, logger,
             {{"code", "Memgraph.QueryExecutionFail"},
              {"message",
               "Query execution has failed (probably there is no "
@@ -57,23 +70,31 @@ State StateExecutorRun(Session &session) {
       // !! QUERY ENGINE -> RUN METHOD -> EXCEPTION HANDLING !!
     } catch (const query::SyntaxException &e) {
       db_accessor->abort();
-      session.encoder_.MessageFailure(
+      StateExecutorFailure<Session>(
+          session, logger,
           {{"code", "Memgraph.SyntaxException"}, {"message", "Syntax error"}});
       return ERROR;
     } catch (const query::QueryEngineException &e) {
       db_accessor->abort();
-      session.encoder_.MessageFailure(
+      StateExecutorFailure<Session>(
+          session, logger,
           {{"code", "Memgraph.QueryEngineException"},
            {"message", "Query engine was unable to execute the query"}});
       return ERROR;
     } catch (const StacktraceException &e) {
       db_accessor->abort();
-      session.encoder_.MessageFailure({{"code", "Memgraph.StacktraceException"},
-                                       {"message", "Unknown exception"}});
+      StateExecutorFailure<Session>(session, logger,
+                                    {{"code", "Memgraph.StacktraceException"},
+                                     {"message", "Unknown exception"}});
       return ERROR;
+    } catch (const BoltException &e) {
+      db_accessor->abort();
+      logger.debug("Failed because: {}", e.what());
+      session.Close();
     } catch (std::exception &e) {
       db_accessor->abort();
-      session.encoder_.MessageFailure(
+      StateExecutorFailure<Session>(
+          session, logger,
           {{"code", "Memgraph.Exception"}, {"message", "Unknown exception"}});
       return ERROR;
     }