server: add posibility to validate request and return custom error response (#105)

Signed-off-by: Marcin Falkowski <[email protected]>

Author: MarcinFalkowski

server/src/main/java/io/envoyproxy/controlplane/server/DiscoveryServer.java

@@ -8,6 +8,7 @@
 import io.envoyproxy.controlplane.cache.Resources;
 import io.envoyproxy.controlplane.cache.Response;
 import io.envoyproxy.controlplane.cache.Watch;
+import io.envoyproxy.controlplane.server.exception.RequestException;
 import io.envoyproxy.controlplane.server.serializer.DefaultProtoResourcesSerializer;
 import io.envoyproxy.controlplane.server.serializer.ProtoResourcesSerializer;
 import io.envoyproxy.envoy.api.v2.ClusterDiscoveryServiceGrpc.ClusterDiscoveryServiceImplBase;
@@ -29,6 +30,7 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 import org.slf4j.Logger;
@@ -199,6 +201,7 @@ private class DiscoveryRequestStreamObserver implements StreamObserver<Discovery
     private final long streamId;
     private final boolean ads;
     private final Executor executor;
+    private final AtomicBoolean isClosing = new AtomicBoolean();
 
     private AtomicLong streamNonce;
 
@@ -243,7 +246,12 @@ public void onNext(DiscoveryRequest request) {
           nonce,
           request.getVersionInfo());
 
-      callbacks.forEach(cb -> cb.onStreamRequest(streamId, request));
+      try {
+        callbacks.forEach(cb -> cb.onStreamRequest(streamId, request));
+      } catch (RequestException e) {
+        closeWithError(e);
+        return;
+      }
 
       for (String typeUrl : Resources.TYPE_URLS) {
         DiscoveryResponse response = latestResponse.get(typeUrl);
@@ -307,6 +315,13 @@ void onCancelled() {
       cancel();
     }
 
+    private void closeWithError(Throwable exception) {
+      if (isClosing.compareAndSet(false, true)) {
+        responseObserver.onError(exception);
+      }
+      cancel();
+    }
+
     private void cancel() {
       watches.values().forEach(Watch::cancel);
     }

server/src/main/java/io/envoyproxy/controlplane/server/DiscoveryServerCallbacks.java

@@ -1,5 +1,6 @@
 package io.envoyproxy.controlplane.server;
 
+import io.envoyproxy.controlplane.server.exception.RequestException;
 import io.envoyproxy.envoy.api.v2.DiscoveryRequest;
 import io.envoyproxy.envoy.api.v2.DiscoveryResponse;
 
@@ -48,6 +49,9 @@ default void onStreamOpen(long streamId, String typeUrl) {
    *
    * @param streamId an ID for this stream that is only unique to this discovery server instance
    * @param request the discovery request sent by the envoy instance
+   *
+   * @throws RequestException optionally can throw {@link RequestException} with custom status. That status
+   *     will be returned to the client and the stream will be closed with error.
    */
   default void onStreamRequest(long streamId, DiscoveryRequest request) {
 

server/src/main/java/io/envoyproxy/controlplane/server/exception/RequestException.java

@@ -0,0 +1,17 @@
+package io.envoyproxy.controlplane.server.exception;
+
+import io.grpc.Metadata;
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+
+import javax.annotation.Nullable;
+
+public class RequestException extends StatusRuntimeException {
+  public RequestException(Status status) {
+    this(status, null);
+  }
+
+  public RequestException(Status status, @Nullable Metadata trailers) {
+    super(status, trailers);
+  }
+}

server/src/test/java/io/envoyproxy/controlplane/server/DiscoveryServerTest.java

@@ -16,6 +16,7 @@
 import io.envoyproxy.controlplane.cache.TestResources;
 import io.envoyproxy.controlplane.cache.Watch;
 import io.envoyproxy.controlplane.cache.WatchCancelledException;
+import io.envoyproxy.controlplane.server.exception.RequestException;
 import io.envoyproxy.envoy.api.v2.Cluster;
 import io.envoyproxy.envoy.api.v2.ClusterDiscoveryServiceGrpc;
 import io.envoyproxy.envoy.api.v2.ClusterDiscoveryServiceGrpc.ClusterDiscoveryServiceStub;
@@ -56,6 +57,7 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
+
 import org.assertj.core.api.Condition;
 import org.junit.Rule;
 import org.junit.Test;
@@ -864,6 +866,90 @@ public Watch createWatch(boolean ads, DiscoveryRequest request, Set<String> know
     assertThat(callbacks.streamResponseCount).hasValue(1);
   }
 
+  @Test
+  public void testCallbacksRequestException() throws InterruptedException {
+    MockDiscoveryServerCallbacks callbacks = new MockDiscoveryServerCallbacks() {
+      @Override
+      public void onStreamRequest(long streamId, DiscoveryRequest request) {
+        super.onStreamRequest(streamId, request);
+        throw new RequestException(Status.INVALID_ARGUMENT.withDescription("request not valid"));
+      }
+    };
+
+    MockConfigWatcher configWatcher = new MockConfigWatcher(false, createResponses());
+    DiscoveryServer server = new DiscoveryServer(callbacks, configWatcher);
+
+    grpcServer.getServiceRegistry().addService(server.getAggregatedDiscoveryServiceImpl());
+    AggregatedDiscoveryServiceStub stub = AggregatedDiscoveryServiceGrpc.newStub(grpcServer.getChannel());
+
+    MockDiscoveryResponseObserver responseObserver = new MockDiscoveryResponseObserver();
+    StreamObserver<DiscoveryRequest> requestObserver = stub.streamAggregatedResources(responseObserver);
+
+    requestObserver.onNext(DiscoveryRequest.newBuilder()
+        .setNode(NODE)
+        .setTypeUrl(Resources.LISTENER_TYPE_URL)
+        .build());
+
+    if (!responseObserver.errorLatch.await(1, TimeUnit.SECONDS) || responseObserver.completed.get()) {
+      fail(format("failed to error before timeout, completed = %b", responseObserver.completed.get()));
+    }
+
+    callbacks.assertThatNoErrors();
+
+    assertThat(responseObserver.errorException).isInstanceOfSatisfying(StatusRuntimeException.class, ex -> {
+      assertThat(ex.getStatus().getCode()).isEqualTo(Status.Code.INVALID_ARGUMENT);
+      assertThat(ex.getStatus().getDescription()).isEqualTo("request not valid");
+    });
+
+    assertThat(callbacks.streamCloseCount).hasValue(0);
+    assertThat(callbacks.streamCloseWithErrorCount).hasValue(0);
+    assertThat(callbacks.streamOpenCount).hasValue(1);
+    assertThat(callbacks.streamRequestCount).hasValue(1);
+    assertThat(callbacks.streamResponseCount).hasValue(0);
+  }
+
+  @Test
+  public void testCallbacksOtherStatusException() throws InterruptedException {
+    MockDiscoveryServerCallbacks callbacks = new MockDiscoveryServerCallbacks() {
+      @Override
+      public void onStreamRequest(long streamId, DiscoveryRequest request) {
+        super.onStreamRequest(streamId, request);
+        throw new StatusRuntimeException(Status.INVALID_ARGUMENT.withDescription("request not valid"));
+      }
+    };
+
+    MockConfigWatcher configWatcher = new MockConfigWatcher(false, createResponses());
+    DiscoveryServer server = new DiscoveryServer(callbacks, configWatcher);
+
+    grpcServer.getServiceRegistry().addService(server.getAggregatedDiscoveryServiceImpl());
+    AggregatedDiscoveryServiceStub stub = AggregatedDiscoveryServiceGrpc.newStub(grpcServer.getChannel());
+
+    MockDiscoveryResponseObserver responseObserver = new MockDiscoveryResponseObserver();
+    StreamObserver<DiscoveryRequest> requestObserver = stub.streamAggregatedResources(responseObserver);
+
+    requestObserver.onNext(DiscoveryRequest.newBuilder()
+        .setNode(NODE)
+        .setTypeUrl(Resources.LISTENER_TYPE_URL)
+        .build());
+
+    if (!responseObserver.errorLatch.await(1, TimeUnit.SECONDS) || responseObserver.completed.get()) {
+      fail(format("failed to error before timeout, completed = %b", responseObserver.completed.get()));
+    }
+
+    callbacks.assertThatNoErrors();
+
+    assertThat(responseObserver.errorException).isInstanceOfSatisfying(StatusRuntimeException.class, ex -> {
+      assertThat(ex.getStatus().getCode()).isEqualTo(Status.Code.UNKNOWN);
+      assertThat(ex.getStatus().getDescription()).isNull();
+    });
+
+    assertThat(callbacks.streamCloseCount).hasValue(0);
+    assertThat(callbacks.streamCloseWithErrorCount).hasValue(0);
+    assertThat(callbacks.streamOpenCount).hasValue(1);
+    assertThat(callbacks.streamRequestCount).hasValue(1);
+    assertThat(callbacks.streamResponseCount).hasValue(0);
+  }
+
   private static Table<String, String, Collection<? extends Message>> createResponses() {
     return ImmutableTable.<String, String, Collection<? extends Message>>builder()
         .put(Resources.CLUSTER_TYPE_URL, VERSION, ImmutableList.of(CLUSTER))
@@ -1005,6 +1091,7 @@ private static class MockDiscoveryResponseObserver implements StreamObserver<Dis
     private final AtomicInteger nonce = new AtomicInteger();
     private final Collection<DiscoveryResponse> responses = new LinkedList<>();
 
+    private Throwable errorException;
     private boolean sendError = false;
 
     void assertThatNoErrors() {
@@ -1054,6 +1141,7 @@ public void onNext(DiscoveryResponse value) {
     @Override
     public void onError(Throwable t) {
       error.set(true);
+      errorException = t;
       errorLatch.countDown();
     }