From d5e6c4c11acdde2f60b897232175889d5aa68f8d Mon Sep 17 00:00:00 2001
From: Chris Marsh <chris@discordapp.com>
Date: Mon, 31 Jul 2017 14:42:36 -0700
Subject: [PATCH] track open/close state in connection, disconnect on read
 error, clarify error codes a little

---
 src/connection.h        |  1 +
 src/connection_unix.cpp | 28 +++++++++++++++++++++++++---
 src/connection_win.cpp  |  8 ++++++++
 src/rpc_connection.cpp  |  9 +++++++--
 src/rpc_connection.h    |  6 ++++++
 5 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/src/connection.h b/src/connection.h
index 6a548ce..a8f99b9 100644
--- a/src/connection.h
+++ b/src/connection.h
@@ -11,6 +11,7 @@ int GetProcessId();
 struct BaseConnection {
     static BaseConnection* Create();
     static void Destroy(BaseConnection*&);
+    bool isOpen{false};
     bool Open();
     bool Close();
     bool Write(const void* data, size_t length);
diff --git a/src/connection_unix.cpp b/src/connection_unix.cpp
index 0249ab1..8fa01a8 100644
--- a/src/connection_unix.cpp
+++ b/src/connection_unix.cpp
@@ -1,5 +1,6 @@
 #include "connection.h"
 
+#include <errno.h>
 #include <fcntl.h>
 #include <stdio.h>
 #include <string.h>
@@ -18,7 +19,12 @@ struct BaseConnectionUnix : public BaseConnection {
 };
 
 static BaseConnectionUnix Connection;
-sockaddr_un PipeAddr{};
+static sockaddr_un PipeAddr{};
+#ifdef MSG_NOSIGNAL
+static int MsgFlags = MSG_NOSIGNAL;
+#else
+static int MsgFlags = 0;
+#endif
 
 static const char* GetTempPath()
 {
@@ -52,11 +58,17 @@ bool BaseConnection::Open()
         return false;
     }
     fcntl(self->sock, F_SETFL, O_NONBLOCK);
+#ifdef SO_NOSIGPIPE
+    int optval = 1;
+    setsockopt(fd, SOL_SOCKET, SO_NOSIGPIPE, &optval, sizeof(optval));
+ #endif
+
     for (int pipeNum = 0; pipeNum < 10; ++pipeNum) {
         snprintf(
           PipeAddr.sun_path, sizeof(PipeAddr.sun_path), "%s/discord-ipc-%d", tempPath, pipeNum);
         int err = connect(self->sock, (const sockaddr*)&PipeAddr, sizeof(PipeAddr));
         if (err == 0) {
+            self->isOpen = true;
             return true;
         }
     }
@@ -72,6 +84,7 @@ bool BaseConnection::Close()
     }
     close(self->sock);
     self->sock = -1;
+    self->isOpen = false;
     return true;
 }
 
@@ -83,7 +96,10 @@ bool BaseConnection::Write(const void* data, size_t length)
         return false;
     }
 
-    ssize_t sentBytes = send(self->sock, data, length, 0);
+    ssize_t sentBytes = send(self->sock, data, length, MsgFlags);
+    if (sentBytes < 0) {
+        Close();
+    }
     return sentBytes == (ssize_t)length;
 }
 
@@ -95,6 +111,12 @@ bool BaseConnection::Read(void* data, size_t length)
         return false;
     }
 
-    int res = recv(self->sock, data, length, 0);
+    int res = recv(self->sock, data, length, MsgFlags);
+    if (res < 0) {
+        if (errno == EAGAIN) {
+            return false;
+        }
+        Close();
+    }
     return res == (int)length;
 }
diff --git a/src/connection_win.cpp b/src/connection_win.cpp
index a2d7361..9cc77e5 100644
--- a/src/connection_win.cpp
+++ b/src/connection_win.cpp
@@ -39,6 +39,7 @@ bool BaseConnection::Open()
         self->pipe = ::CreateFileW(
           pipeName, GENERIC_READ | GENERIC_WRITE, 0, nullptr, OPEN_EXISTING, 0, nullptr);
         if (self->pipe != INVALID_HANDLE_VALUE) {
+            self->isOpen = true;
             return true;
         }
 
@@ -64,6 +65,7 @@ bool BaseConnection::Close()
     auto self = reinterpret_cast<BaseConnectionWin*>(this);
     ::CloseHandle(self->pipe);
     self->pipe = INVALID_HANDLE_VALUE;
+    self->isOpen = false;
     return true;
 }
 
@@ -82,7 +84,13 @@ bool BaseConnection::Read(void* data, size_t length)
             if (::ReadFile(self->pipe, data, length, nullptr, nullptr) == TRUE) {
                 return true;
             }
+            else {
+                Close();
+            }
         }
     }
+    else {
+        Close();
+    }
     return false;
 }
diff --git a/src/rpc_connection.cpp b/src/rpc_connection.cpp
index 1ff2d3b..d408fdd 100644
--- a/src/rpc_connection.cpp
+++ b/src/rpc_connection.cpp
@@ -97,13 +97,18 @@ bool RpcConnection::Read(JsonDocument& message)
     for (;;) {
         bool didRead = connection->Read(&readFrame, sizeof(MessageFrameHeader));
         if (!didRead) {
+            if (!connection->isOpen) {
+                lastErrorCode = (int)ErrorCode::PipeClosed;
+                StringCopy(lastErrorMessage, "Pipe closed");
+                Close();
+            }
             return false;
         }
 
         if (readFrame.length > 0) {
             didRead = connection->Read(readFrame.message, readFrame.length);
             if (!didRead) {
-                lastErrorCode = -2;
+                lastErrorCode = (int)ErrorCode::ReadCorrupt;
                 StringCopy(lastErrorMessage, "Partial data in frame");
                 Close();
                 return false;
@@ -132,7 +137,7 @@ bool RpcConnection::Read(JsonDocument& message)
             break;
         default:
             // something bad happened
-            lastErrorCode = -1;
+            lastErrorCode = (int)ErrorCode::ReadCorrupt;
             StringCopy(lastErrorMessage, "Bad ipc frame");
             Close();
             return false;
diff --git a/src/rpc_connection.h b/src/rpc_connection.h
index cafe8f5..d3c30d6 100644
--- a/src/rpc_connection.h
+++ b/src/rpc_connection.h
@@ -8,6 +8,12 @@
 constexpr size_t MaxRpcFrameSize = 64 * 1024;
 
 struct RpcConnection {
+    enum class ErrorCode : int {
+        Success = 0,
+        PipeClosed = 1,
+        ReadCorrupt = 2,
+    };
+
     enum class Opcode : uint32_t {
         Handshake = 0,
         Frame = 1,