Quit the proxy when reading zero or less from stdin

fishman · Aug 20, 2018 · 1d80bdd · 1d80bdd
1 parent 8db604e
commit 1d80bdd
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 12 deletions.
diff --git a/src/browser/NativeMessagingBase.h b/src/browser/NativeMessagingBase.h
@@ -52,7 +52,7 @@ protected slots:
 
 protected:
     virtual void    readLength() = 0;
-    virtual void    readStdIn(const quint32 length) = 0;
+    virtual bool    readStdIn(const quint32 length) = 0;
     void            readNativeMessages();
     QString         jsonToString(const QJsonObject& json) const;
     void            sendReply(const QJsonObject& json);

diff --git a/src/browser/NativeMessagingHost.cpp b/src/browser/NativeMessagingHost.cpp
@@ -114,10 +114,10 @@ void NativeMessagingHost::readLength()
     }
 }
 
-void NativeMessagingHost::readStdIn(const quint32 length)
+bool NativeMessagingHost::readStdIn(const quint32 length)
 {
     if (length <= 0) {
-        return;
+        return false;
     }
 
     QByteArray arr;
@@ -129,14 +129,15 @@ void NativeMessagingHost::readStdIn(const quint32 length)
         int c = std::getchar();
         if (c == EOF) {
             // message ended prematurely, ignore it and return
-            return;
+            return false;
         }
         arr.append(static_cast<char>(c));
     }
 
     if (arr.length() > 0) {
         sendReply(m_browserClients.readResponse(arr));
     }
+    return true;
 }
 
 void NativeMessagingHost::newLocalConnection()

diff --git a/src/browser/NativeMessagingHost.h b/src/browser/NativeMessagingHost.h
@@ -46,7 +46,7 @@ public slots:
 
 private:
     void        readLength();
-    void        readStdIn(const quint32 length);
+    bool        readStdIn(const quint32 length);
     void        sendReplyToAllClients(const QJsonObject& json);
 
 private slots:

diff --git a/src/proxy/NativeMessagingHost.cpp b/src/proxy/NativeMessagingHost.cpp
@@ -49,21 +49,36 @@ NativeMessagingHost::~NativeMessagingHost()
 #endif
 }
 
+void NativeMessagingHost::readNativeMessages()
+{
+#ifdef Q_OS_WIN
+    quint32 length = 0;
+    while (m_running.load() && !std::cin.eof()) {
+        length = 0;
+        std::cin.read(reinterpret_cast<char*>(&length), 4);
+        if (!readStdIn(length)) {
+            QCoreApplication::quit();
+        }
+        QThread::msleep(1);
+    }
+#endif
+}
+
 void NativeMessagingHost::readLength()
 {
     quint32 length = 0;
     std::cin.read(reinterpret_cast<char*>(&length), 4);
     if (!std::cin.eof() && length > 0) {
         readStdIn(length);
     } else {
-    	QCoreApplication::quit();
+        QCoreApplication::quit();
     }
 }
 
-void NativeMessagingHost::readStdIn(const quint32 length)
+bool NativeMessagingHost::readStdIn(const quint32 length)
 {
     if (length <= 0) {
-        return;
+        return false;
     }
 
     QByteArray arr;
@@ -73,7 +88,7 @@ void NativeMessagingHost::readStdIn(const quint32 length)
         int c = std::getchar();
         if (c == EOF) {
             // message ended prematurely, ignore it and return
-            return;
+            return false;
         }
         arr.append(static_cast<char>(c));
     }
@@ -82,6 +97,8 @@ void NativeMessagingHost::readStdIn(const quint32 length)
         m_localSocket->write(arr.constData(), arr.length());
         m_localSocket->flush();
     }
+
+    return true;
 }
 
 void NativeMessagingHost::newLocalMessage()
@@ -92,7 +109,7 @@ void NativeMessagingHost::newLocalMessage()
 
     QByteArray arr = m_localSocket->readAll();
     if (!arr.isEmpty()) {
-       sendReply(arr);
+        sendReply(arr);
     }
 }
 

diff --git a/src/proxy/NativeMessagingHost.h b/src/proxy/NativeMessagingHost.h
@@ -33,11 +33,12 @@ public slots:
     void socketStateChanged(QLocalSocket::LocalSocketState socketState);
 
 private:
+    void readNativeMessages();
     void readLength();
-    void readStdIn(const quint32 length);
+    bool readStdIn(const quint32 length);
 
 private:
-    QLocalSocket*                           m_localSocket;
+    QLocalSocket*   m_localSocket;
 };
 
 #endif // NATIVEMESSAGINGHOST_H