From aa0b7cc9af4769a8cfd5350293f2bb7a355f51f3 Mon Sep 17 00:00:00 2001
From: zhangkun <zhangkun@google.com>
Date: Tue, 11 Nov 2014 18:03:27 -0800
Subject: [PATCH] Forbid calling next more than once in interceptors.

Will throw IllegalStateException.
-------------
Created by MOE: http://code.google.com/p/moe-java
MOE_MIGRATED_REVID=79717262
---
 .../com/google/net/stubby/ClientInterceptors.java |  8 ++++++--
 .../com/google/net/stubby/ServerInterceptors.java |  3 +++
 .../google/net/stubby/ClientInterceptorsTest.java | 14 ++++++++++++++
 .../google/net/stubby/ServerInterceptorsTest.java | 15 +++++++++++++++
 4 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/core/src/main/java/com/google/net/stubby/ClientInterceptors.java b/core/src/main/java/com/google/net/stubby/ClientInterceptors.java
index 1e40972de0..8f2f45e9d4 100644
--- a/core/src/main/java/com/google/net/stubby/ClientInterceptors.java
+++ b/core/src/main/java/com/google/net/stubby/ClientInterceptors.java
@@ -53,7 +53,7 @@ public class ClientInterceptors {
 
   private static class ProcessInterceptorChannel implements Channel {
     private final Channel channel;
-    private final Iterator<ClientInterceptor> interceptors;
+    private Iterator<ClientInterceptor> interceptors;
 
     private ProcessInterceptorChannel(Channel channel, Iterable<ClientInterceptor> interceptors) {
       this.channel = channel;
@@ -62,9 +62,13 @@ public class ClientInterceptors {
 
     @Override
     public <ReqT, RespT> Call<ReqT, RespT> newCall(MethodDescriptor<ReqT, RespT> method) {
-      if (interceptors.hasNext()) {
+      if (interceptors != null && interceptors.hasNext()) {
         return interceptors.next().interceptCall(method, this);
       } else {
+        Preconditions.checkState(interceptors != null,
+            "The channel has already been called. "
+            + "Some interceptor must have called on \"next\" twice.");
+        interceptors = null;
         return channel.newCall(method);
       }
     }
diff --git a/core/src/main/java/com/google/net/stubby/ServerInterceptors.java b/core/src/main/java/com/google/net/stubby/ServerInterceptors.java
index 1e4de74cdf..0b710c35f5 100644
--- a/core/src/main/java/com/google/net/stubby/ServerInterceptors.java
+++ b/core/src/main/java/com/google/net/stubby/ServerInterceptors.java
@@ -93,6 +93,9 @@ public class ServerInterceptors {
       if (interceptors != null && interceptors.hasNext()) {
         return interceptors.next().interceptCall(method, call, headers, this);
       } else {
+        Preconditions.checkState(interceptors != null,
+            "The call handler has already been called. "
+            + "Some interceptor must have called on \"next\" twice.");
         interceptors = null;
         return callHandler.startCall(method, call, headers);
       }
diff --git a/core/src/test/java/com/google/net/stubby/ClientInterceptorsTest.java b/core/src/test/java/com/google/net/stubby/ClientInterceptorsTest.java
index 635b31ec09..990f365fe0 100644
--- a/core/src/test/java/com/google/net/stubby/ClientInterceptorsTest.java
+++ b/core/src/test/java/com/google/net/stubby/ClientInterceptorsTest.java
@@ -81,6 +81,20 @@ public class ClientInterceptorsTest {
     verifyNoMoreInteractions(channel, interceptor);
   }
 
+  @Test(expected = IllegalStateException.class)
+  public void callNextTwice() {
+    ClientInterceptor interceptor = new ClientInterceptor() {
+      @Override
+      public <ReqT, RespT> Call<ReqT, RespT> interceptCall(MethodDescriptor<ReqT, RespT> method,
+          Channel next) {
+        next.newCall(method);
+        return next.newCall(method);
+      }
+    };
+    Channel intercepted = ClientInterceptors.intercept(channel, interceptor);
+    intercepted.newCall(method);
+  }
+
   @Test
   public void ordered() {
     final List<String> order = new ArrayList<String>();
diff --git a/core/src/test/java/com/google/net/stubby/ServerInterceptorsTest.java b/core/src/test/java/com/google/net/stubby/ServerInterceptorsTest.java
index 6efc5cb4c2..663206b5d8 100644
--- a/core/src/test/java/com/google/net/stubby/ServerInterceptorsTest.java
+++ b/core/src/test/java/com/google/net/stubby/ServerInterceptorsTest.java
@@ -116,6 +116,21 @@ public class ServerInterceptorsTest {
     verifyNoMoreInteractions(handler2);
   }
 
+  @Test(expected = IllegalStateException.class)
+  public void callNextTwice() {
+    ServerInterceptor interceptor = new ServerInterceptor() {
+      @Override
+      public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(String method,
+          ServerCall<RespT> call, Headers headers, ServerCallHandler<ReqT, RespT> next) {
+        next.startCall(method, call, headers);
+        return next.startCall(method, call, headers);
+      }
+    };
+    ServerServiceDefinition intercepted = ServerInterceptors.intercept(serviceDefinition,
+        interceptor);
+    getSoleMethod(intercepted).getServerCallHandler().startCall(methodName, call, headers);
+  }
+
   @Test
   public void ordered() {
     final List<String> order = new ArrayList<String>();