From 5856538706dc3abc44dd5ba1a8bb19a0dfb1ab25 Mon Sep 17 00:00:00 2001
From: lyuxuan <yuxuanli@google.com>
Date: Tue, 24 Oct 2017 10:20:54 -0700
Subject: [PATCH] cap max msg size to min(max_int, max_uint32) (#1598)

---
 rpc_util.go       |  2 ++
 server.go         | 14 ++++++++++++--
 service_config.go | 28 ++++++++++++++++++++++------
 3 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/rpc_util.go b/rpc_util.go
index e006b7e6..f5978c19 100644
--- a/rpc_util.go
+++ b/rpc_util.go
@@ -206,6 +206,7 @@ func FailFast(failFast bool) CallOption {
 }
 
 // MaxCallRecvMsgSize returns a CallOption which sets the maximum message size the client can receive.
+// Note that the maximum effective value is MaxUint32 due to protocol limitations.
 func MaxCallRecvMsgSize(s int) CallOption {
 	return beforeCall(func(o *callInfo) error {
 		o.maxReceiveMessageSize = &s
@@ -214,6 +215,7 @@ func MaxCallRecvMsgSize(s int) CallOption {
 }
 
 // MaxCallSendMsgSize returns a CallOption which sets the maximum message size the client can send.
+// Note that the maximum effective value is MaxUint32 due to protocol limitations.
 func MaxCallSendMsgSize(s int) CallOption {
 	return beforeCall(func(o *callInfo) error {
 		o.maxSendMessageSize = &s
diff --git a/server.go b/server.go
index def301a1..70311854 100644
--- a/server.go
+++ b/server.go
@@ -208,17 +208,27 @@ func MaxMsgSize(m int) ServerOption {
 
 // MaxRecvMsgSize returns a ServerOption to set the max message size in bytes the server can receive.
 // If this is not set, gRPC uses the default 4MB.
+// Note that the maximum effective value is MaxUint32 due to protocol limitations.
 func MaxRecvMsgSize(m int) ServerOption {
 	return func(o *options) {
-		o.maxReceiveMessageSize = m
+		if int64(m) > int64(math.MaxUint32) {
+			o.maxReceiveMessageSize = math.MaxUint32
+		} else {
+			o.maxReceiveMessageSize = m
+		}
 	}
 }
 
 // MaxSendMsgSize returns a ServerOption to set the max message size in bytes the server can send.
 // If this is not set, gRPC uses the default 4MB.
+// Note that the maximum effective value is MaxUint32 due to protocol limitations.
 func MaxSendMsgSize(m int) ServerOption {
 	return func(o *options) {
-		o.maxSendMessageSize = m
+		if int64(m) > int64(math.MaxUint32) {
+			o.maxSendMessageSize = math.MaxUint32
+		} else {
+			o.maxSendMessageSize = m
+		}
 	}
 }
 
diff --git a/service_config.go b/service_config.go
index 0631e7d2..719d725e 100644
--- a/service_config.go
+++ b/service_config.go
@@ -20,6 +20,7 @@ package grpc
 
 import (
 	"encoding/json"
+	"math"
 	"time"
 
 	"google.golang.org/grpc/grpclog"
@@ -148,24 +149,39 @@ func parseServiceConfig(js string) (ServiceConfig, error) {
 	return sc, nil
 }
 
-func min(a, b *int) *int {
-	if *a < *b {
+func min(a, b int) int {
+	if a < b {
 		return a
 	}
 	return b
 }
 
+const maxInt = int(^uint(0) >> 1)
+
 func getMaxSize(mcMax, doptMax *int, defaultVal int) *int {
+	res := getRawMaxSize(mcMax, doptMax, defaultVal)
+
+	// Cap the max size to maxInt of current machine due to slice length limit.
+	res = min(res, maxInt)
+	if int64(res) > int64(math.MaxUint32) {
+		// Only reach here on 64-bit machine, where we need to cap the max size
+		// to MaxUint32.
+		res = math.MaxUint32
+	}
+	return &res
+}
+
+func getRawMaxSize(mcMax, doptMax *int, defaultVal int) int {
 	if mcMax == nil && doptMax == nil {
-		return &defaultVal
+		return defaultVal
 	}
 	if mcMax != nil && doptMax != nil {
-		return min(mcMax, doptMax)
+		return min(*mcMax, *doptMax)
 	}
 	if mcMax != nil {
-		return mcMax
+		return *mcMax
 	}
-	return doptMax
+	return *doptMax
 }
 
 func newBool(b bool) *bool {