From 0dff8a7d3f0da388a6327ba9c797691ae855312f Mon Sep 17 00:00:00 2001
From: Maxime Coste <frrrwww@gmail.com>
Date: Tue, 29 Nov 2016 19:12:56 +0000
Subject: [PATCH] Small remote code cleanup

---
 src/main.cc   | 13 ++++---------
 src/remote.cc | 26 +++++++++++++++++---------
 src/remote.hh |  7 -------
 3 files changed, 21 insertions(+), 25 deletions(-)

diff --git a/src/main.cc b/src/main.cc
index f522c088..5dc28d5c 100644
--- a/src/main.cc
+++ b/src/main.cc
@@ -465,14 +465,9 @@ int run_client(StringView session, StringView init_command, UIType ui_type)
         while (true)
             event_manager.handle_next_events(EventMode::Normal);
     }
-    catch (remote_error&)
+    catch (remote_error& e)
     {
-        write_stderr("disconnected from server\n");
-        return -1;
-    }
-    catch (connection_failed& e)
-    {
-        write_stderr(format("{}\n", e.what()));
+        write_stderr(format("{}\ndisconnecting\n", e.what()));
         return -1;
     }
     return 0;
@@ -713,9 +708,9 @@ int run_pipe(StringView session)
     {
         send_command(session, command);
     }
-    catch (connection_failed& e)
+    catch (remote_error& e)
     {
-        write_stderr(format("{}\n", e.what()));
+        write_stderr(format("{}\ndisconnecting\n", e.what()));
         return -1;
     }
     return 0;
diff --git a/src/remote.cc b/src/remote.cc
index 9de0a13a..762db191 100644
--- a/src/remote.cc
+++ b/src/remote.cc
@@ -49,10 +49,18 @@ public:
 
     ~MsgWriter() noexcept(false)
     {
-        *reinterpret_cast<uint32_t*>(m_stream.data()+1) = (uint32_t)m_stream.size();
-        int res = ::write(m_socket, m_stream.data(), m_stream.size());
-        if (res == 0)
-            throw remote_error{"peer disconnected"};
+        uint32_t count = (uint32_t)m_stream.size();
+        char* data = m_stream.data();
+        *reinterpret_cast<uint32_t*>(data+1) = count;
+        while (count > 0)
+        {
+            int res = ::write(m_socket, data, count);
+            if (res <= 0)
+                throw remote_error{res ? "peer disconnected"
+                                       : format("socket write failed: {}", strerror(errno))};
+            data += res;
+            count -= res;
+        }
     }
 
     void write(const char* val, size_t size)
@@ -226,10 +234,9 @@ private:
     void read_from_socket(int sock, size_t size)
     {
         int res = ::read(sock, m_stream.data() + m_write_pos, size);
-        if (res == 0)
-            throw remote_error{"peer disconnected"};
-        if (res < 0)
-            throw remote_error{format("socket error, read returned {}", res)};
+        if (res <= 0)
+            throw remote_error{res ? "peer disconnected"
+                                   : format("socket read failed: {}", strerror(errno))};
         m_write_pos += res;
     }
 
@@ -238,6 +245,7 @@ private:
     uint32_t m_write_pos = 0;
     uint32_t m_read_pos = header_size;
 };
+
 template<>
 String MsgReader::read<String>()
 {
@@ -493,7 +501,7 @@ static int connect_to(StringView session)
     fcntl(sock, F_SETFD, FD_CLOEXEC);
     sockaddr_un addr = session_addr(session);
     if (connect(sock, (sockaddr*)&addr, sizeof(addr.sun_path)) == -1)
-        throw connection_failed(addr.sun_path);
+        throw remote_error(format("connect to {} failed", addr.sun_path));
     return sock;
 }
 
diff --git a/src/remote.hh b/src/remote.hh
index 7e171b17..a79c38b4 100644
--- a/src/remote.hh
+++ b/src/remote.hh
@@ -19,13 +19,6 @@ struct remote_error : runtime_error
     {}
 };
 
-struct connection_failed : runtime_error
-{
-    connection_failed(StringView filename)
-        : runtime_error{format("connect to {} failed", filename)}
-    {}
-};
-
 class FDWatcher;
 
 // A remote client handle communication between a client running on the server