From d28535bb202c08bc36730276d7d5e605b96d3a08 Mon Sep 17 00:00:00 2001
From: Kun Zhang <zhangkun@google.com>
Date: Thu, 16 Jul 2015 18:17:07 -0700
Subject: [PATCH] Fail call if more than responses received while only one is
 expected

---
 stub/src/main/java/io/grpc/stub/ClientCalls.java | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/stub/src/main/java/io/grpc/stub/ClientCalls.java b/stub/src/main/java/io/grpc/stub/ClientCalls.java
index d9dae75502..6c957ba165 100644
--- a/stub/src/main/java/io/grpc/stub/ClientCalls.java
+++ b/stub/src/main/java/io/grpc/stub/ClientCalls.java
@@ -167,7 +167,9 @@ public class ClientCalls {
       ReqT param,
       ClientCall.Listener<RespT> responseListener) {
     call.start(responseListener, new Metadata.Headers());
-    call.request(1);
+    // Initially ask for two responses from flow-control so that if a misbehaving server sends more
+    // than one responses, we can catch it.
+    call.request(2);
     try {
       call.sendPayload(param);
       call.halfClose();
@@ -182,7 +184,9 @@ public class ClientCalls {
       boolean streamingResponse) {
     call.start(new StreamObserverToCallListenerAdapter<RespT>(
         call, responseObserver, streamingResponse), new Metadata.Headers());
-    call.request(1);
+    // Initially ask for two responses from flow-control so that if a misbehaving server sends more
+    // than one responses, we can catch it.
+    call.request(2);
     return new CallToStreamObserverAdapter<ReqT>(call);
   }
 
@@ -215,6 +219,7 @@ public class ClientCalls {
     private final ClientCall<?, RespT> call;
     private final StreamObserver<RespT> observer;
     private final boolean streamingResponse;
+    private boolean firstResponseReceived;
 
     public StreamObserverToCallListenerAdapter(
         ClientCall<?, RespT> call, StreamObserver<RespT> observer, boolean streamingResponse) {
@@ -229,6 +234,12 @@ public class ClientCalls {
 
     @Override
     public void onPayload(RespT payload) {
+      if (firstResponseReceived && !streamingResponse) {
+        throw Status.INTERNAL
+            .withDescription("More than one responses received for unary or client-streaming call")
+            .asRuntimeException();
+      }
+      firstResponseReceived = true;
       observer.onValue(payload);
 
       if (streamingResponse) {