grpc-js: Avoid buffering significantly more than max_receive_message_…

…size per received message (1.8.x)
grpc · Jun 7, 2024 · 8e62222 · 8e62222
1 parent 9d83947
commit 8e62222
Show file tree

Hide file tree

Showing 9 changed files with 186 additions and 148 deletions.
diff --git a/packages/grpc-js/src/compression-filter.ts b/packages/grpc-js/src/compression-filter.ts
@@ -21,7 +21,7 @@ import { WriteObject, WriteFlags } from './call-interface';
 import { Channel } from './channel';
 import { ChannelOptions } from './channel-options';
 import { CompressionAlgorithms } from './compression-algorithms';
-import { LogVerbosity } from './constants';
+import { DEFAULT_MAX_RECEIVE_MESSAGE_LENGTH, LogVerbosity, Status } from './constants';
 import { BaseFilter, Filter, FilterFactory } from './filter';
 import * as logging from './logging';
 import { Metadata, MetadataValue } from './metadata';
@@ -94,6 +94,10 @@ class IdentityHandler extends CompressionHandler {
 }
 
 class DeflateHandler extends CompressionHandler {
+  constructor(private maxRecvMessageLength: number) {
+    super();
+  }
+
   compressMessage(message: Buffer) {
     return new Promise<Buffer>((resolve, reject) => {
       zlib.deflate(message, (err, output) => {
@@ -108,18 +112,34 @@ class DeflateHandler extends CompressionHandler {
 
   decompressMessage(message: Buffer) {
     return new Promise<Buffer>((resolve, reject) => {
-      zlib.inflate(message, (err, output) => {
-        if (err) {
-          reject(err);
-        } else {
-          resolve(output);
+      let totalLength = 0;
+      const messageParts: Buffer[] = [];
+      const decompresser = zlib.createInflate();
+      decompresser.on('data', (chunk: Buffer) => {
+        messageParts.push(chunk);
+        totalLength += chunk.byteLength;
+        if (this.maxRecvMessageLength !== -1 && totalLength > this.maxRecvMessageLength) {
+          decompresser.destroy();
+          reject({
+            code: Status.RESOURCE_EXHAUSTED,
+            details: `Received message that decompresses to a size larger than ${this.maxRecvMessageLength}`
+          });
         }
       });
+      decompresser.on('end', () => {
+        resolve(Buffer.concat(messageParts));
+      });
+      decompresser.write(message);
+      decompresser.end();
     });
   }
 }
 
 class GzipHandler extends CompressionHandler {
+  constructor(private maxRecvMessageLength: number) {
+    super();
+  }
+
   compressMessage(message: Buffer) {
     return new Promise<Buffer>((resolve, reject) => {
       zlib.gzip(message, (err, output) => {
@@ -134,13 +154,25 @@ class GzipHandler extends CompressionHandler {
 
   decompressMessage(message: Buffer) {
     return new Promise<Buffer>((resolve, reject) => {
-      zlib.unzip(message, (err, output) => {
-        if (err) {
-          reject(err);
-        } else {
-          resolve(output);
+      let totalLength = 0;
+      const messageParts: Buffer[] = [];
+      const decompresser = zlib.createGunzip();
+      decompresser.on('data', (chunk: Buffer) => {
+        messageParts.push(chunk);
+        totalLength += chunk.byteLength;
+        if (this.maxRecvMessageLength !== -1 && totalLength > this.maxRecvMessageLength) {
+          decompresser.destroy();
+          reject({
+            code: Status.RESOURCE_EXHAUSTED,
+            details: `Received message that decompresses to a size larger than ${this.maxRecvMessageLength}`
+          });
         }
       });
+      decompresser.on('end', () => {
+        resolve(Buffer.concat(messageParts));
+      });
+      decompresser.write(message);
+      decompresser.end();
     });
   }
 }
@@ -165,14 +197,14 @@ class UnknownHandler extends CompressionHandler {
   }
 }
 
-function getCompressionHandler(compressionName: string): CompressionHandler {
+function getCompressionHandler(compressionName: string, maxReceiveMessageSize: number): CompressionHandler {
   switch (compressionName) {
     case 'identity':
       return new IdentityHandler();
     case 'deflate':
-      return new DeflateHandler();
+      return new DeflateHandler(maxReceiveMessageSize);
     case 'gzip':
-      return new GzipHandler();
+      return new GzipHandler(maxReceiveMessageSize);
     default:
       return new UnknownHandler(compressionName);
   }
@@ -182,11 +214,14 @@ export class CompressionFilter extends BaseFilter implements Filter {
   private sendCompression: CompressionHandler = new IdentityHandler();
   private receiveCompression: CompressionHandler = new IdentityHandler();
   private currentCompressionAlgorithm: CompressionAlgorithm = 'identity';
+  private maxReceiveMessageLength: number;
 
   constructor(channelOptions: ChannelOptions, private sharedFilterConfig: SharedCompressionFilterConfig) {
     super();
 
-    const compressionAlgorithmKey = channelOptions['grpc.default_compression_algorithm'];
+    const compressionAlgorithmKey =
+      channelOptions['grpc.default_compression_algorithm'];
+    this.maxReceiveMessageLength = channelOptions['grpc.max_receive_message_length'] ?? DEFAULT_MAX_RECEIVE_MESSAGE_LENGTH
     if (compressionAlgorithmKey !== undefined) {
       if (isCompressionAlgorithmKey(compressionAlgorithmKey)) {
         const clientSelectedEncoding = CompressionAlgorithms[compressionAlgorithmKey] as CompressionAlgorithm;
@@ -200,7 +235,10 @@ export class CompressionFilter extends BaseFilter implements Filter {
          */
         if (!serverSupportedEncodings || serverSupportedEncodings.includes(clientSelectedEncoding)) {
           this.currentCompressionAlgorithm = clientSelectedEncoding;
-          this.sendCompression = getCompressionHandler(this.currentCompressionAlgorithm);
+          this.sendCompression = getCompressionHandler(
+            this.currentCompressionAlgorithm,
+            -1
+          );
         }
       } else {
         logging.log(LogVerbosity.ERROR, `Invalid value provided for grpc.default_compression_algorithm option: ${compressionAlgorithmKey}`);
@@ -228,7 +266,7 @@ export class CompressionFilter extends BaseFilter implements Filter {
     if (receiveEncoding.length > 0) {
       const encoding: MetadataValue = receiveEncoding[0];
       if (typeof encoding === 'string') {
-        this.receiveCompression = getCompressionHandler(encoding);
+        this.receiveCompression = getCompressionHandler(encoding, this.maxReceiveMessageLength);
       }
     }
     metadata.remove('grpc-encoding');

diff --git a/packages/grpc-js/src/internal-channel.ts b/packages/grpc-js/src/internal-channel.ts
@@ -33,7 +33,6 @@ import {
 } from './resolver';
 import { trace } from './logging';
 import { SubchannelAddress } from './subchannel-address';
-import { MaxMessageSizeFilterFactory } from './max-message-size-filter';
 import { mapProxyName } from './http_proxy';
 import { GrpcUri, parseUri, uriToString } from './uri-parser';
 import { ServerSurfaceCall } from './server-call';
@@ -310,7 +309,6 @@ export class InternalChannel {
       }
     );
     this.filterStackFactory = new FilterStackFactory([
-      new MaxMessageSizeFilterFactory(this.options),
       new CompressionFilterFactory(this, this.options),
     ]);
     this.trace('Channel constructed with options ' + JSON.stringify(options, undefined, 2));

diff --git a/packages/grpc-js/src/max-message-size-filter.ts b/packages/grpc-js/src/max-message-size-filter.ts
diff --git a/packages/grpc-js/src/server-call.ts b/packages/grpc-js/src/server-call.ts
@@ -19,7 +19,6 @@ import { EventEmitter } from 'events';
 import * as http2 from 'http2';
 import { Duplex, Readable, Writable } from 'stream';
 import * as zlib from 'zlib';
-import { promisify } from 'util';
 
 import {
   Status,
@@ -38,8 +37,6 @@ import { Deadline } from './deadline';
 import { getErrorCode, getErrorMessage } from './error';
 
 const TRACER_NAME = 'server_call';
-const unzip = promisify(zlib.unzip);
-const inflate = promisify(zlib.inflate);
 
 function trace(text: string): void {
   logging.trace(LogVerbosity.DEBUG, TRACER_NAME, text);
@@ -478,19 +475,42 @@ export class Http2ServerCallStream<
   private getDecompressedMessage(
     message: Buffer,
     encoding: string
-  ): Buffer | Promise<Buffer> {
-    if (encoding === 'deflate') {
-      return inflate(message.subarray(5));
-    } else if (encoding === 'gzip') {
-      return unzip(message.subarray(5));
-    } else if (encoding === 'identity') {
-      return message.subarray(5);
+  ): Buffer | Promise<Buffer> {    const messageContents = message.subarray(5);
+    if (encoding === 'identity') {
+      return messageContents;
+    } else if (encoding === 'deflate' || encoding === 'gzip') {
+      let decompresser: zlib.Gunzip | zlib.Deflate;
+      if (encoding === 'deflate') {
+        decompresser = zlib.createInflate();
+      } else {
+        decompresser = zlib.createGunzip();
+      }
+      return new Promise((resolve, reject) => {
+        let totalLength = 0
+        const messageParts: Buffer[] = [];
+        decompresser.on('data', (chunk: Buffer) => {
+          messageParts.push(chunk);
+          totalLength += chunk.byteLength;
+          if (this.maxReceiveMessageSize !== -1 && totalLength > this.maxReceiveMessageSize) {
+            decompresser.destroy();
+            reject({
+              code: Status.RESOURCE_EXHAUSTED,
+              details: `Received message that decompresses to a size larger than ${this.maxReceiveMessageSize}`
+            });
+          }
+        });
+        decompresser.on('end', () => {
+          resolve(Buffer.concat(messageParts));
+        });
+        decompresser.write(messageContents);
+        decompresser.end();
+      });
+    } else {
+      return Promise.reject({
+        code: Status.UNIMPLEMENTED,
+        details: `Received message compressed with unsupported encoding "${encoding}"`,
+      });
     }
-
-    return Promise.reject({
-      code: Status.UNIMPLEMENTED,
-      details: `Received message compressed with unsupported encoding "${encoding}"`,
-    });
   }
 
   sendMetadata(customMetadata?: Metadata) {
@@ -807,7 +827,7 @@ export class Http2ServerCallStream<
       | ServerDuplexStream<RequestType, ResponseType>,
     encoding: string
   ) {
-    const decoder = new StreamDecoder();
+    const decoder = new StreamDecoder(this.maxReceiveMessageSize);
 
     let readsDone = false;
 
@@ -823,29 +843,34 @@ export class Http2ServerCallStream<
     };
 
     this.stream.on('data', async (data: Buffer) => {
-      const messages = decoder.write(data);
+      let messages: Buffer[];
+      try {
+        messages = decoder.write(data);
+      } catch (e) {
+        this.sendError({
+          code: Status.RESOURCE_EXHAUSTED,
+          details: (e as Error).message
+        });
+        return;
+      }
 
       pendingMessageProcessing = true;
       this.stream.pause();
       for (const message of messages) {
-        if (
-          this.maxReceiveMessageSize !== -1 &&
-          message.length > this.maxReceiveMessageSize
-        ) {
-          this.sendError({
-            code: Status.RESOURCE_EXHAUSTED,
-            details: `Received message larger than max (${message.length} vs. ${this.maxReceiveMessageSize})`,
-          });
-          return;
-        }
         this.emit('receiveMessage');
 
         const compressed = message.readUInt8(0) === 1;
         const compressedMessageEncoding = compressed ? encoding : 'identity';
-        const decompressedMessage = await this.getDecompressedMessage(
-          message,
-          compressedMessageEncoding
-        );
+        let decompressedMessage: Buffer;
+        try {
+          decompressedMessage = await this.getDecompressedMessage(
+            message,
+            compressedMessageEncoding
+          );
+        } catch (e) {
+          this.sendError(e as Partial<StatusObject>);
+          return;
+        }
 
         // Encountered an error with decompression; it'll already have been propogated back
         // Just return early