fix uws polling

socketio · Feb 18, 2022 · 7d033f1 · 7d033f1
1 parent a463d26
commit 7d033f1
Showing 1 changed file with 42 additions and 12 deletions.
diff --git a/lib/transports-uws/polling.ts b/lib/transports-uws/polling.ts
@@ -122,6 +122,18 @@ export class Polling extends Transport {
       return;
     }
 
+    const contentLengthHeader = Number(req.headers["content-length"]);
+    if (!contentLengthHeader) {
+      this.onError("invalid content-length header");
+      res.writeStatus("400 Invalid content-length header").end();
+      return;
+    }
+    if (contentLengthHeader > this.maxHttpBufferSize) {
+      this.onError("payload too large");
+      res.writeStatus("413 Payload Too Large").end();
+      return;
+    }
+
     const isBinary = "application/octet-stream" === req.headers["content-type"];
 
     if (isBinary && this.protocol === 4) {
@@ -131,11 +143,11 @@ export class Polling extends Transport {
     this.dataReq = req;
     this.dataRes = res;
 
-    let chunks = [];
+    let buffer;
     let contentLength = 0;
 
     const cleanup = () => {
-      this.dataReq = this.dataRes = chunks = null;
+      this.dataReq = this.dataRes = null;
     };
 
     const onClose = () => {
@@ -154,8 +166,8 @@ export class Polling extends Transport {
       res.writeHeader(key, String(headers[key]));
     });
 
-    const onEnd = () => {
-      this.onData(Buffer.concat(chunks).toString());
+    const onEnd = (buffer) => {
+      this.onData(buffer.toString());
 
       if (this.readyState !== "closing") {
         res.end("ok");
@@ -165,18 +177,36 @@ export class Polling extends Transport {
 
     res.onAborted(onClose);
 
-    res.onData((chunk, isLast) => {
-      chunks.push(Buffer.from(chunk));
-      contentLength += Buffer.byteLength(chunk);
-      if (contentLength > this.maxHttpBufferSize) {
-        this.onError("payload too large");
-        res.writeStatus("413 Payload Too Large");
-        res.end();
+    res.onData((arrayBuffer, isLast) => {
+      const totalLength = contentLength + arrayBuffer.byteLength;
+      if (totalLength > contentLengthHeader) {
+        this.onError("content-length mismatch");
+        res.close(); // calls onAborted
         return;
       }
+
+      if (!buffer) {
+        if (isLast) {
+          onEnd(Buffer.from(arrayBuffer));
+          return;
+        }
+        buffer = Buffer.allocUnsafe(contentLengthHeader);
+      }
+
+      Buffer.from(arrayBuffer).copy(buffer, contentLength);
+
       if (isLast) {
-        onEnd();
+        if (totalLength != contentLengthHeader) {
+          this.onError("content-length mismatch");
+          res.writeStatus("400 content-length mismatch").end();
+          cleanup();
+          return;
+        }
+        onEnd(buffer);
+        return;
       }
+
+      contentLength = totalLength;
     });
   }