Fix memory overrun when adding a single line and b64-encode inputs

MisterTea · May 17, 2018 · b8124da · b8124da
1 parent ec125e1
commit b8124da
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 5 deletions.
diff --git a/htm/HtmServer.cpp b/htm/HtmServer.cpp
@@ -48,8 +48,8 @@ void HtmServer::run() {
             RawSocketUtils::readAll(endpointFd, &uid[0], uid.length());
             length -= uid.length();
             LOG(ERROR) << "READING FROM " << uid << ":" << length;
-            string data(length, '\0');
-            RawSocketUtils::readAll(endpointFd, &data[0], length);
+            string data;
+            RawSocketUtils::readB64EncodedLength(endpointFd, &data, length);
             LOG(ERROR) << "READ FROM " << uid << ":" << data << " " << length;
             state.appendData(uid, data);
             break;

diff --git a/htm/TerminalHandler.cpp b/htm/TerminalHandler.cpp
@@ -113,7 +113,9 @@ string TerminalHandler::pollUserTerminal() {
           buffer.insert(buffer.end(), tokens.begin(), tokens.end());
         } else {
           buffer.back().append(tokens.front());
-          buffer.insert(buffer.end(), tokens.begin() + 1, tokens.end());
+          if (tokens.size() > 1) {
+            buffer.insert(buffer.end(), tokens.begin() + 1, tokens.end());
+          }
         }
         if (buffer.size() > MAX_BUFFER_LINES) {
           int amountToErase = buffer.size() - MAX_BUFFER_LINES;

diff --git a/terminal/RawSocketUtils.hpp b/terminal/RawSocketUtils.hpp
@@ -15,15 +15,27 @@ class RawSocketUtils {
   static inline void writeB64(int fd, const char* buf, size_t count) {
     int encodedLength = base64::Base64::EncodedLength(count);
     string s(encodedLength, '\0');
-    base64::Base64::Encode(buf, count, &s[0], s.length());
+    if (!base64::Base64::Encode(buf, count, &s[0], s.length())) {
+      throw runtime_error("b64 decode failed");
+    }
     writeAll(fd, &s[0], s.length());
   }
 
   static inline void readB64(int fd, char* buf, size_t count) {
     int encodedLength = base64::Base64::EncodedLength(count);
     string s(encodedLength, '\0');
     readAll(fd, &s[0], s.length());
-    base64::Base64::Decode((const char*)&s[0], s.length(), buf, count);
+    if(!base64::Base64::Decode((const char*)&s[0], s.length(), buf, count)) {
+      throw runtime_error("b64 decode failed");
+    }
+  }
+
+  static inline void readB64EncodedLength(int fd, string* out, size_t encodedLength) {
+    string s(encodedLength, '\0');
+    readAll(fd, &s[0], s.length());
+    if(!base64::Base64::Decode(s, out)) {
+      throw runtime_error("b64 decode failed");
+    }
   }
 
   static inline string readMessage(int fd) {