xds: Add ExtAuthzClientCall

xds/src/main/java/io/grpc/xds/internal/extauthz/AuthzCallbackObserver.java

@@ -0,0 +1,170 @@
+/*
+ * Copyright 2025 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.xds.internal.extauthz;
+
+import com.google.common.collect.ImmutableList;
+import io.envoyproxy.envoy.service.auth.v3.CheckResponse;
+import io.grpc.CallOptions;
+import io.grpc.Channel;
+import io.grpc.ClientCall;
+import io.grpc.Context;
+import io.grpc.Metadata;
+import io.grpc.MethodDescriptor;
+import io.grpc.Status;
+import io.grpc.internal.DelayedClientCall;
+import io.grpc.stub.StreamObserver;
+import io.grpc.xds.internal.grpcservice.HeaderValue;
+import io.grpc.xds.internal.headermutations.HeaderMutations;
+import io.grpc.xds.internal.headermutations.HeaderMutator;
+import io.grpc.xds.internal.headermutations.HeaderValueOption;
+import java.util.concurrent.Executor;
+
+/**
+ * A {@link StreamObserver} that handles the response from the external authorization service,
+ * transitioning the delayed gRPC call based on the authorization decision.
+ */
+final class AuthzCallbackObserver<ReqT, RespT> implements StreamObserver<CheckResponse> {
+
+  private static final Metadata.Key<String> X_ENVOY_AUTH_FAILURE_MODE_ALLOWED =
+      Metadata.Key.of("x-envoy-auth-failure-mode-allowed", Metadata.ASCII_STRING_MARSHALLER);
+
+  private final DelayedClientCall<ReqT, RespT> delayedCall;
+  private final Channel next;
+  private final MethodDescriptor<ReqT, RespT> method;
+  private final CallOptions callOptions;
+  private final Executor callExecutor;
+  private final CheckResponseHandler responseHandler;
+  // TODO(sauravzg): Make this static seeing that it holds no state.
+  private final HeaderMutator headerMutator = HeaderMutator.create();
+  private final ExtAuthzConfig config;
+  private final Context.CancellableContext authzContext;
+
+  AuthzCallbackObserver(
+      DelayedClientCall<ReqT, RespT> delayedCall,
+      Channel next,
+      MethodDescriptor<ReqT, RespT> method,
+      CallOptions callOptions,
+      Executor callExecutor,
+      CheckResponseHandler responseHandler,
+      ExtAuthzConfig config,
+      Context.CancellableContext authzContext) {
+    this.delayedCall = delayedCall;
+    this.next = next;
+    this.method = method;
+    this.callOptions = callOptions;
+    this.callExecutor = callExecutor;
+    this.responseHandler = responseHandler;
+    this.config = config;
+    this.authzContext = authzContext;
+  }
+
+  @Override
+  public void onNext(CheckResponse value) {
+    // Note: This implementation is currently exception-safe.
+    //
+    // TODO(sauravz): Revisit hardening if this invariant changes in the future.
+    // If an unhandled RuntimeException escapes onNext(), gRPC cancels the stream and
+    // invokes onError(). Under failure_mode_allow: true, this causes the call to fail
+    // open, which could inadvertently permit an unauthorized request.
+    AuthzResponse authzResponse = responseHandler.handleResponse(value);
+    if (authzResponse.decision() == AuthzResponse.Decision.ALLOW) {
+      ClientCall<ReqT, RespT> delegate = next.newCall(method, callOptions);
+      ClientCall<ReqT, RespT> finalCall;
+      if (isHeaderMutationsEmpty(authzResponse.requestHeaderMutations())
+          && isHeaderMutationsEmpty(authzResponse.responseHeaderMutations())) {
+        finalCall = delegate;
+      } else {
+        finalCall = new MutatingClientCall<>(
+            delegate,
+            authzResponse.requestHeaderMutations(),
+            authzResponse.responseHeaderMutations(),
+            headerMutator);
+      }
+      setCallAndDrain(finalCall);
+    } else {
+      // Defensive programming: status is always present on DENY branches since every
+      // deny() factory method requires a Status parameter. This orElseThrow satisfies
+      // static analysis and documents the invariant.
+      Status status = authzResponse.status().orElseThrow(
+          () -> new IllegalArgumentException("DENY response missing status"));
+      ClientCall<ReqT, RespT> failingCall;
+      if (isHeaderMutationsEmpty(authzResponse.responseHeaderMutations())) {
+        failingCall = new FailingClientCall<>(status);
+      } else {
+        failingCall = new FailingCallWithTrailerMutations<>(
+            status,
+            authzResponse.responseHeaderMutations(),
+            headerMutator);
+      }
+      setCallAndDrain(failingCall);
+    }
+  }
+
+  @Override
+  public void onError(Throwable t) {
+    if (config.failureModeAllow()) {
+      ClientCall<ReqT, RespT> delegate = next.newCall(method, callOptions);
+      ClientCall<ReqT, RespT> finalCall;
+      if (config.failureModeAllowHeaderAdd()) {
+        HeaderMutations requestMutations = HeaderMutations.create(
+            ImmutableList.of(HeaderValueOption.create(
+                HeaderValue.create(X_ENVOY_AUTH_FAILURE_MODE_ALLOWED.name(), "true"),
+                HeaderValueOption.HeaderAppendAction.APPEND_IF_EXISTS_OR_ADD)),
+            ImmutableList.of());
+        // TODO(sauravz): Consider using a static EMPTY_MUTATIONS constant
+        // to avoid repeated allocation of empty HeaderMutations.
+        HeaderMutations responseMutations =
+            HeaderMutations.create(ImmutableList.of(), ImmutableList.of());
+        finalCall = new MutatingClientCall<>(
+            delegate, requestMutations, responseMutations, headerMutator);
+      } else {
+        finalCall = delegate;
+      }
+      setCallAndDrain(finalCall);
+      authzContext.close();
+    } else {
+      Status statusToReturn = config.statusOnError().withCause(t);
+      setCallAndDrain(new FailingClientCall<>(statusToReturn));
+      authzContext.close();
+    }
+  }
+
+  @Override
+  public void onCompleted() {
+    // Defensive: if the authz server completed without sending a CheckResponse
+    // (buggy server), resolve the DelayedClientCall with a failure and close
+    // the authzContext. If onNext() already called setCall(), this is a no-op
+    // since DelayedClientCall.setCall() ignores subsequent calls.
+    Status status = config.statusOnError()
+        .withDescription("Authz server completed without sending CheckResponse");
+    setCallAndDrain(new FailingClientCall<>(status));
+    authzContext.close();
+  }
+
+  private void setCallAndDrain(ClientCall<ReqT, RespT> call) {
+    Runnable drain = delayedCall.setCall(call);
+    // drain is null if the call was already set (e.g. onCompleted() following onNext())
+    // or if the call was cancelled in-flight.
+    if (drain != null) {
+      callExecutor.execute(drain);
+    }
+  }
+
+  private static boolean isHeaderMutationsEmpty(HeaderMutations mutations) {
+    return mutations.headers().isEmpty() && mutations.headersToRemove().isEmpty();
+  }
+}

xds/src/main/java/io/grpc/xds/internal/extauthz/ExtAuthzClientCall.java

@@ -0,0 +1,123 @@
+/*
+ * Copyright 2025 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.xds.internal.extauthz;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.protobuf.util.Timestamps;
+import io.envoyproxy.envoy.service.auth.v3.AuthorizationGrpc;
+import io.envoyproxy.envoy.service.auth.v3.CheckRequest;
+import io.grpc.CallOptions;
+import io.grpc.Channel;
+import io.grpc.ClientCall;
+import io.grpc.Context;
+import io.grpc.Deadline;
+import io.grpc.ForwardingClientCall;
+import io.grpc.Metadata;
+import io.grpc.MethodDescriptor;
+import io.grpc.internal.DelayedClientCall;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ScheduledExecutorService;
+import javax.annotation.Nullable;
+
+/**
+ * A {@link ForwardingClientCall} that delegates to a lightweight {@link DelayedClientCall}
+ * to buffer RPC requests safely while waiting for the asynchronous authorization decision.
+ */
+public final class ExtAuthzClientCall<ReqT, RespT> extends ForwardingClientCall<ReqT, RespT> {
+
+  private final Channel next;
+  private final MethodDescriptor<ReqT, RespT> method;
+  private final CallOptions callOptions;
+  private final AuthorizationGrpc.AuthorizationStub authzStub;
+  private final CheckRequestBuilder checkRequestBuilder;
+  private final CheckResponseHandler responseHandler;
+  private final ExtAuthzConfig config;
+  private final Executor callExecutor;
+  private final Context.CancellableContext authzContext;
+
+  private final DelayedAuthzCall<ReqT, RespT> delegate;
+
+  public ExtAuthzClientCall(
+      Executor executor,
+      ScheduledExecutorService scheduler,
+      CallOptions callOptions,
+      Channel next,
+      MethodDescriptor<ReqT, RespT> method,
+      AuthorizationGrpc.AuthorizationStub authzStub,
+      CheckRequestBuilder checkRequestBuilder,
+      CheckResponseHandler responseHandler,
+      ExtAuthzConfig config) {
+    this.callOptions = callOptions;
+    this.next = next;
+    this.method = method;
+    this.authzStub = authzStub;
+    this.checkRequestBuilder = checkRequestBuilder;
+    this.responseHandler = responseHandler;
+    this.config = config;
+    this.callExecutor = executor;
+    this.authzContext = Context.current().withCancellation();
+    this.delegate = new DelayedAuthzCall<>(executor, scheduler, callOptions.getDeadline());
+  }
+
+  @Override
+  protected ClientCall<ReqT, RespT> delegate() {
+    return delegate;
+  }
+
+  @Override
+  public void start(Listener<RespT> responseListener, Metadata headers) {
+    // 1. Construct the check request FIRST before handoff (respects metadata thread safety)
+    CheckRequest request = checkRequestBuilder.buildRequest(method, headers,
+        Timestamps.fromMillis(System.currentTimeMillis()));
+
+    // 2. Start the delayed call (buffers headers and outbound payloads)
+    super.start(responseListener, headers);
+
+    // 3. Trigger the async authorization check under the cancellable context
+    authzContext.run(() -> {
+      authzStub.check(request, new AuthzCallbackObserver<>(
+          delegate, next, method, callOptions, callExecutor, responseHandler,
+          config, authzContext));
+    });
+  }
+
+  @Override
+  public void cancel(
+      @Nullable String message, @Nullable Throwable cause) {
+    super.cancel(message, cause);
+    authzContext.cancel(cause);
+  }
+
+  /**
+   * Returns the cancellable context used for the authorization check.
+   * Visible for testing.
+   */
+  @VisibleForTesting
+  Context.CancellableContext getAuthzContextForTest() {
+    return authzContext;
+  }
+
+  /**
+   * A lightweight package-private DelayedClientCall subclass to expose
+   * the protected constructor.
+   */
+  private static final class DelayedAuthzCall<ReqT, RespT> extends DelayedClientCall<ReqT, RespT> {
+    DelayedAuthzCall(Executor executor, ScheduledExecutorService scheduler, Deadline deadline) {
+      super("ExtAuthzClientCall", executor, scheduler, deadline);
+    }
+  }
+}

xds/src/main/java/io/grpc/xds/internal/extauthz/FailingCallWithTrailerMutations.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2025 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.xds.internal.extauthz;
+
+import io.grpc.ClientCall;
+import io.grpc.Metadata;
+import io.grpc.Status;
+import io.grpc.xds.internal.headermutations.HeaderMutations;
+import io.grpc.xds.internal.headermutations.HeaderMutator;
+
+/**
+ * A simple failing client call that lazily applies response mutations to trailers during start.
+ */
+final class FailingCallWithTrailerMutations<ReqT, RespT> extends ClientCall<ReqT, RespT> {
+  private final Status status;
+  private final HeaderMutations responseHeaderMutations;
+  private final HeaderMutator headerMutator;
+
+  FailingCallWithTrailerMutations(
+      Status status,
+      HeaderMutations responseHeaderMutations,
+      HeaderMutator headerMutator) {
+    this.status = status;
+    this.responseHeaderMutations = responseHeaderMutations;
+    this.headerMutator = headerMutator;
+  }
+
+  @Override
+  public void start(Listener<RespT> responseListener, Metadata headers) {
+    // Lazily allocate and apply response mutations to trailers copy on start!
+    Metadata trailers = new Metadata();
+    headerMutator.applyMutations(responseHeaderMutations, trailers);
+    responseListener.onClose(status, trailers);
+  }
+
+  @Override
+  public void request(int numMessages) {}
+
+  @Override
+  public void cancel(String message, Throwable cause) {}
+
+  @Override
+  public void halfClose() {}
+
+  @Override
+  public void sendMessage(ReqT message) {}
+}