feat(api-server): add getBodyRaw method to NanotronClientRequest for …

…handling request body
Alwatr · Sep 13, 2024 · 828416b · 828416b
1 parent e4eb63e
commit 828416b
Showing 1 changed file with 72 additions and 0 deletions.
diff --git a/packages/api-server/src/api-client-request.ts b/packages/api-server/src/api-client-request.ts
@@ -56,4 +56,76 @@ export class NanotronClientRequest {
     // Create server response.
     this.serverResponse = new NanotronServerResponse(this, nativeServerResponse);
   }
+
+  getBodyRaw(): Promise<Buffer> {
+    return new Promise((resolve, reject) => {
+      // method must be POST or PUT or PATCH
+      const method = this.url.method;
+      if (!(method === 'POST' || method === 'PUT' || method === 'PATCH')) {
+        this.serverResponse.headers.connection = 'close';
+        this.serverResponse.statusCode = HttpStatusCodes.Error_Client_405_Method_Not_Allowed;
+        return reject(new Error('body_not_allowed'));
+      }
+
+      const contentLength = +this.raw_.headers['content-length']!;
+
+      if (Number.isNaN(contentLength) || contentLength === 0) {
+        this.serverResponse.headers.connection = 'close';
+        this.serverResponse.statusCode = HttpStatusCodes.Error_Client_411_Length_Required;
+        return reject(new Error('body_length_required'));
+      }
+
+      if (contentLength > this.routeOption!.bodyLimit) {
+        // TODO: routeOption when null?!
+        this.serverResponse.headers.connection = 'close';
+        this.serverResponse.statusCode = HttpStatusCodes.Error_Client_413_Payload_Too_Large;
+        return reject(new Error('body_too_large'));
+      }
+
+      const dataChunks: Uint8Array[] = [];
+      let receivedLength = 0;
+
+      const onData = (chunk: Uint8Array) => {
+        receivedLength += chunk.length;
+        dataChunks.push(chunk);
+
+        if (receivedLength > this.routeOption!.bodyLimit) {
+          this.raw_.removeListener('data', onData);
+          this.raw_.removeListener('end', onEnd);
+          this.raw_.removeListener('error', onEnd);
+          dataChunks.length = 0; // free memory
+          this.serverResponse.headers.connection = 'close';
+          this.serverResponse.statusCode = HttpStatusCodes.Error_Client_413_Payload_Too_Large;
+          return reject(new Error('body_too_large'));
+        }
+      };
+
+      const onEnd = (err?: Error) => {
+        this.raw_.removeListener('data', onData);
+        this.raw_.removeListener('end', onEnd);
+        this.raw_.removeListener('error', onEnd);
+
+        if (err !== undefined) {
+          dataChunks.length = 0; // free memory
+          this.serverResponse.statusCode = HttpStatusCodes.Error_Client_400_Bad_Request;
+          return reject(err);
+        }
+
+        if (receivedLength !== contentLength) {
+          dataChunks.length = 0; // free memory
+          this.serverResponse.statusCode = HttpStatusCodes.Error_Client_413_Payload_Too_Large;
+          return reject(new Error('body_too_large'));
+        }
+
+        const body = Buffer.concat(dataChunks);
+
+        resolve(body);
+      };
+
+      this.raw_.on('data', onData);
+      this.raw_.on('end', onEnd);
+      this.raw_.on('error', onEnd);
+      this.raw_.resume();
+    });
+  }
 }