From 67b67f83826195c62d1e25ddbce0b3bc77257b03 Mon Sep 17 00:00:00 2001
From: James Roper <james@jazzy.id.au>
Date: Tue, 19 Dec 2023 06:02:51 +1100
Subject: [PATCH] Fix outbound message size checking (#10739)

This fixes two bugs in outbound message size checking:

* When thet checke failed, the thrown StatusRuntimeException with a status code
  of RESOURCE_EXHAUSTED was been caught and rewrapped in another
  StatusRuntimeException but this time with status code UNKNOWN.
* This applies the max outbound message size check to messages prior to, and
  after compression, since compression of a message that is smaller than the
  maximum send size can result in a larger message that exceeds the maximum
  send size.
---
 .../java/io/grpc/internal/MessageFramer.java    | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/core/src/main/java/io/grpc/internal/MessageFramer.java b/core/src/main/java/io/grpc/internal/MessageFramer.java
index 93d35250a0..5e75fa2e6f 100644
--- a/core/src/main/java/io/grpc/internal/MessageFramer.java
+++ b/core/src/main/java/io/grpc/internal/MessageFramer.java
@@ -27,6 +27,7 @@ import io.grpc.Compressor;
 import io.grpc.Drainable;
 import io.grpc.KnownLength;
 import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
@@ -147,6 +148,8 @@ public class MessageFramer implements Framer {
           .withDescription("Failed to frame message")
           .withCause(e)
           .asRuntimeException();
+    } catch (StatusRuntimeException e) {
+      throw e;
     } catch (RuntimeException e) {
       throw Status.INTERNAL
           .withDescription("Failed to frame message")
@@ -170,13 +173,6 @@ public class MessageFramer implements Framer {
     }
     BufferChainOutputStream bufferChain = new BufferChainOutputStream();
     int written = writeToOutputStream(message, bufferChain);
-    if (maxOutboundMessageSize >= 0 && written > maxOutboundMessageSize) {
-      throw Status.RESOURCE_EXHAUSTED
-          .withDescription(
-              String.format(
-                  Locale.US, "message too large %d > %d", written , maxOutboundMessageSize))
-          .asRuntimeException();
-    }
     writeBufferChain(bufferChain, false);
     return written;
   }
@@ -238,6 +234,13 @@ public class MessageFramer implements Framer {
    */
   private void writeBufferChain(BufferChainOutputStream bufferChain, boolean compressed) {
     int messageLength = bufferChain.readableBytes();
+    if (maxOutboundMessageSize >= 0 && messageLength > maxOutboundMessageSize) {
+      throw Status.RESOURCE_EXHAUSTED
+          .withDescription(
+              String.format(
+                  Locale.US, "message too large %d > %d", messageLength , maxOutboundMessageSize))
+          .asRuntimeException();
+    }
     headerScratch.clear();
     headerScratch.put(compressed ? COMPRESSED : UNCOMPRESSED).putInt(messageLength);
     WritableBuffer writeableHeader = bufferAllocator.allocate(HEADER_LENGTH);