grpc-js: Avoid buffering significantly more than max_receive_message_size per received message (1.10.x)

Author: murgatroid99

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';
@@ -98,6 +98,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) => {
@@ -112,18 +116,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) => {
@@ -138,13 +158,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();
     });
   }
 }
@@ -169,14 +201,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);
   }
@@ -186,6 +218,7 @@ 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,
@@ -195,6 +228,7 @@ export class CompressionFilter extends BaseFilter implements Filter {
 
     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[
@@ -215,7 +249,8 @@ export class CompressionFilter extends BaseFilter implements Filter {
         ) {
           this.currentCompressionAlgorithm = clientSelectedEncoding;
           this.sendCompression = getCompressionHandler(
-            this.currentCompressionAlgorithm
+            this.currentCompressionAlgorithm,
+            -1
           );
         }
       } else {
@@ -247,7 +282,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');

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';
@@ -402,7 +401,6 @@ export class InternalChannel {
       }
     );
     this.filterStackFactory = new FilterStackFactory([
-      new MaxMessageSizeFilterFactory(this.options),
       new CompressionFilterFactory(this, this.options),
     ]);
     this.trace(

packages/grpc-js/src/max-message-size-filter.ts

@@ -30,14 +30,10 @@ import {
 import * as http2 from 'http2';
 import { getErrorMessage } from './error';
 import * as zlib from 'zlib';
-import { promisify } from 'util';
 import { StreamDecoder } from './stream-decoder';
 import { CallEventTracker } from './transport';
 import * as logging from './logging';
 
-const unzip = promisify(zlib.unzip);
-const inflate = promisify(zlib.inflate);
-
 const TRACER_NAME = 'server_call';
 
 function trace(text: string) {
@@ -496,7 +492,7 @@ export class BaseServerInterceptingCall
   private wantTrailers = false;
   private cancelNotified = false;
   private incomingEncoding = 'identity';
-  private decoder = new StreamDecoder();
+  private decoder: StreamDecoder;
   private readQueue: ReadQueueEntry[] = [];
   private isReadPending = false;
   private receivedHalfClose = false;
@@ -554,6 +550,8 @@ export class BaseServerInterceptingCall
       this.maxReceiveMessageSize = options['grpc.max_receive_message_length']!;
     }
 
+    this.decoder = new StreamDecoder(this.maxReceiveMessageSize);
+
     const metadata = Metadata.fromHttp2Headers(headers);
 
     if (logging.isTracerEnabled(TRACER_NAME)) {
@@ -674,18 +672,41 @@ export class BaseServerInterceptingCall
     message: Buffer,
     encoding: string
   ): Buffer | Promise<Buffer> {
-    switch (encoding) {
-      case 'deflate':
-        return inflate(message.subarray(5));
-      case 'gzip':
-        return unzip(message.subarray(5));
-      case 'identity':
-        return message.subarray(5);
-      default:
-        return Promise.reject({
-          code: Status.UNIMPLEMENTED,
-          details: `Received message compressed with unsupported encoding "${encoding}"`,
+    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}"`,
+      });
     }
   }
 
@@ -698,10 +719,16 @@ export class BaseServerInterceptingCall
     const compressedMessageEncoding = compressed
       ? this.incomingEncoding
       : 'identity';
-    const decompressedMessage = await this.decompressMessage(
-      queueEntry.compressedMessage!,
-      compressedMessageEncoding
-    );
+    let decompressedMessage: Buffer;
+    try {
+      decompressedMessage = await this.decompressMessage(
+        queueEntry.compressedMessage!,
+        compressedMessageEncoding
+      );
+    } catch (err) {
+      this.sendStatus(err as PartialStatusObject);
+      return;
+    }
     try {
       queueEntry.parsedMessage = this.handler.deserialize(decompressedMessage);
     } catch (err) {
@@ -743,23 +770,16 @@ export class BaseServerInterceptingCall
         ' received data frame of size ' +
         data.length
     );
-    const rawMessages = this.decoder.write(data);
+    let rawMessages: Buffer[];
+    try {
+      rawMessages = this.decoder.write(data);
+    } catch (e) {
+      this.sendStatus({ code: Status.RESOURCE_EXHAUSTED, details: (e as Error).message });
+      return;
+    }
 
     for (const messageBytes of rawMessages) {
       this.stream.pause();
-      if (
-        this.maxReceiveMessageSize !== -1 &&
-        messageBytes.length - 5 > this.maxReceiveMessageSize
-      ) {
-        this.sendStatus({
-          code: Status.RESOURCE_EXHAUSTED,
-          details: `Received message larger than max (${
-            messageBytes.length - 5
-          } vs. ${this.maxReceiveMessageSize})`,
-          metadata: null,
-        });
-        return;
-      }
       const queueEntry: ReadQueueEntry = {
         type: 'COMPRESSED',
         compressedMessage: messageBytes,