[Java] Implement auth API for Flight

java/flight/pom.xml

@@ -54,6 +54,11 @@
       <artifactId>grpc-core</artifactId>
       <version>${dep.grpc.version}</version>
     </dependency>
+    <dependency>
+      <groupId>io.grpc</groupId>
+      <artifactId>grpc-context</artifactId>
+      <version>${dep.grpc.version}</version>
+    </dependency>
     <dependency>
       <groupId>io.grpc</groupId>
       <artifactId>grpc-protobuf</artifactId>

java/flight/src/main/java/org/apache/arrow/flight/FlightClient.java

@@ -113,8 +113,9 @@ public class FlightClient implements AutoCloseable {
   }
 
   public void authenticate(ClientAuthHandler handler) {
-    Preconditions.checkArgument(!authInterceptor.hasToken(), "Auth already completed.");
-    authInterceptor.setToken(ClientAuthWrapper.doClientAuth(handler, asyncStub));
+    Preconditions.checkArgument(!authInterceptor.hasAuthHandler(), "Auth already completed.");
+    ClientAuthWrapper.doClientAuth(handler, asyncStub);
+    authInterceptor.setAuthHandler(handler);
   }
 
   /**

java/flight/src/main/java/org/apache/arrow/flight/FlightProducer.java

@@ -27,17 +27,22 @@ import org.apache.arrow.vector.VectorSchemaRoot;
  */
 public interface FlightProducer {
 
-  public void getStream(Ticket ticket, ServerStreamListener listener);
+  public void getStream(CallContext context, Ticket ticket,
+      ServerStreamListener listener);
 
-  public void listFlights(Criteria criteria, StreamListener<FlightInfo> listener);
+  public void listFlights(CallContext context, Criteria criteria,
+      StreamListener<FlightInfo> listener);
 
-  public FlightInfo getFlightInfo(FlightDescriptor descriptor);
+  public FlightInfo getFlightInfo(CallContext context,
+      FlightDescriptor descriptor);
 
-  public Callable<PutResult> acceptPut(FlightStream flightStream);
+  public Callable<PutResult> acceptPut(CallContext context,
+      FlightStream flightStream);
 
-  public Result doAction(Action action);
+  public Result doAction(CallContext context, Action action);
 
-  public void listActions(StreamListener<ActionType> listener);
+  public void listActions(CallContext context,
+      StreamListener<ActionType> listener);
 
   public interface ServerStreamListener {
 
@@ -65,5 +70,14 @@ public interface FlightProducer {
 
   }
 
+  /**
+   * Call-specific context.
+   */
+  interface CallContext {
+    /** The identity of the authenticated peer. May be the empty string if unknown. */
+    String peerIdentity();
 
+    /** Whether the call has been cancelled by the client. */
+    boolean isCancelled();
+  }
 }

java/flight/src/main/java/org/apache/arrow/flight/FlightService.java

@@ -19,8 +19,10 @@ package org.apache.arrow.flight;
 
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.function.BooleanSupplier;
 
 import org.apache.arrow.flight.FlightProducer.ServerStreamListener;
+import org.apache.arrow.flight.auth.AuthConstants;
 import org.apache.arrow.flight.auth.ServerAuthHandler;
 import org.apache.arrow.flight.auth.ServerAuthWrapper;
 import org.apache.arrow.flight.impl.Flight;
@@ -67,15 +69,22 @@ class FlightService extends FlightServiceImplBase {
   @Override
   public void listFlights(Flight.Criteria criteria, StreamObserver<FlightGetInfo> responseObserver) {
     try {
-      producer.listFlights(new Criteria(criteria), StreamPipe.wrap(responseObserver, t -> t.toProtocol()));
+      producer.listFlights(makeContext((ServerCallStreamObserver<?>) responseObserver), new Criteria(criteria),
+          StreamPipe.wrap(responseObserver, FlightInfo::toProtocol));
     } catch (Exception ex) {
       responseObserver.onError(ex);
     }
   }
 
+  private CallContext makeContext(ServerCallStreamObserver<?> responseObserver) {
+    return new CallContext(AuthConstants.PEER_IDENTITY_KEY.get(),
+        responseObserver::isCancelled);
+  }
+
   public void doGetCustom(Flight.Ticket ticket, StreamObserver<ArrowMessage> responseObserver) {
     try {
-      producer.getStream(new Ticket(ticket), new GetListener(responseObserver));
+      producer.getStream(makeContext((ServerCallStreamObserver<?>) responseObserver), new Ticket(ticket),
+          new GetListener(responseObserver));
     } catch (Exception ex) {
       responseObserver.onError(ex);
     }
@@ -84,7 +93,9 @@ class FlightService extends FlightServiceImplBase {
   @Override
   public void doAction(Flight.Action request, StreamObserver<Result> responseObserver) {
     try {
-      responseObserver.onNext(producer.doAction(new Action(request)).toProtocol());
+      responseObserver.onNext(
+          producer.doAction(makeContext((ServerCallStreamObserver<?>) responseObserver), new Action(request))
+              .toProtocol());
       responseObserver.onCompleted();
     } catch (Exception ex) {
       responseObserver.onError(ex);
@@ -94,7 +105,8 @@ class FlightService extends FlightServiceImplBase {
   @Override
   public void listActions(Empty request, StreamObserver<ActionType> responseObserver) {
     try {
-      producer.listActions(StreamPipe.wrap(responseObserver, t -> t.toProtocol()));
+      producer.listActions(makeContext((ServerCallStreamObserver<?>) responseObserver),
+          StreamPipe.wrap(responseObserver, t -> t.toProtocol()));
     } catch (Exception ex) {
       responseObserver.onError(ex);
     }
@@ -158,7 +170,7 @@ class FlightService extends FlightServiceImplBase {
     FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, null, (count) -> responseObserver.request(count));
     executors.submit(() -> {
       try {
-        responseObserver.onNext(producer.acceptPut(fs).call());
+        responseObserver.onNext(producer.acceptPut(makeContext(responseObserver), fs).call());
         responseObserver.onCompleted();
       } catch (Exception ex) {
         responseObserver.onError(ex);
@@ -171,7 +183,8 @@ class FlightService extends FlightServiceImplBase {
   @Override
   public void getFlightInfo(Flight.FlightDescriptor request, StreamObserver<FlightGetInfo> responseObserver) {
     try {
-      FlightInfo info = producer.getFlightInfo(new FlightDescriptor(request));
+      FlightInfo info = producer
+          .getFlightInfo(makeContext((ServerCallStreamObserver<?>) responseObserver), new FlightDescriptor(request));
       responseObserver.onNext(info.toProtocol());
       responseObserver.onCompleted();
     } catch (Exception ex) {
@@ -179,4 +192,24 @@ class FlightService extends FlightServiceImplBase {
     }
   }
 
+  static class CallContext implements FlightProducer.CallContext {
+
+    private final String peerIdentity;
+    private final BooleanSupplier isCancelled;
+
+    CallContext(final String peerIdentity, BooleanSupplier isCancelled) {
+      this.peerIdentity = peerIdentity;
+      this.isCancelled = isCancelled;
+    }
+
+    @Override
+    public String peerIdentity() {
+      return peerIdentity;
+    }
+
+    @Override
+    public boolean isCancelled() {
+      return this.isCancelled.getAsBoolean();
+    }
+  }
 }

java/flight/src/main/java/org/apache/arrow/flight/NoOpFlightProducer.java

@@ -24,32 +24,37 @@ import org.apache.arrow.flight.impl.Flight.PutResult;
 public class NoOpFlightProducer implements FlightProducer {
 
   @Override
-  public void getStream(Ticket ticket, ServerStreamListener listener) {
+  public void getStream(CallContext context, Ticket ticket,
+      ServerStreamListener listener) {
     listener.error(new UnsupportedOperationException("NYI"));
   }
 
   @Override
-  public void listFlights(Criteria criteria, StreamListener<FlightInfo> listener) {
+  public void listFlights(CallContext context, Criteria criteria,
+      StreamListener<FlightInfo> listener) {
     listener.onError(new UnsupportedOperationException("NYI"));
   }
 
   @Override
-  public FlightInfo getFlightInfo(FlightDescriptor descriptor) {
+  public FlightInfo getFlightInfo(CallContext context,
+      FlightDescriptor descriptor) {
     throw new UnsupportedOperationException("NYI");
   }
 
   @Override
-  public Callable<PutResult> acceptPut(FlightStream flightStream) {
+  public Callable<PutResult> acceptPut(CallContext context,
+      FlightStream flightStream) {
     throw new UnsupportedOperationException("NYI");
   }
 
   @Override
-  public Result doAction(Action action) {
+  public Result doAction(CallContext context, Action action) {
     throw new UnsupportedOperationException("NYI");
   }
 
   @Override
-  public void listActions(StreamListener<ActionType> listener) {
+  public void listActions(CallContext context,
+      StreamListener<ActionType> listener) {
     listener.onError(new UnsupportedOperationException("NYI"));
   }
 

java/flight/src/main/java/org/apache/arrow/flight/auth/AuthConstants.java

@@ -19,6 +19,7 @@ package org.apache.arrow.flight.auth;
 
 import org.apache.arrow.flight.FlightConstants;
 
+import io.grpc.Context;
 import io.grpc.Metadata.BinaryMarshaller;
 import io.grpc.Metadata.Key;
 import io.grpc.MethodDescriptor;
@@ -41,5 +42,7 @@ public final class AuthConstants {
     }
   });
 
+  public static final Context.Key<String> PEER_IDENTITY_KEY = Context.keyWithDefault("arrow-flight-peer-identity", "");
+
   private AuthConstants() {}
 }

java/flight/src/main/java/org/apache/arrow/flight/auth/BasicClientAuthHandler.java

@@ -28,6 +28,7 @@ public class BasicClientAuthHandler implements ClientAuthHandler {
 
   private final String name;
   private final String password;
+  private byte[] token = null;
 
   public BasicClientAuthHandler(String name, String password) {
     this.name = name;
@@ -35,7 +36,7 @@ public class BasicClientAuthHandler implements ClientAuthHandler {
   }
 
   @Override
-  public byte[] authenticate(ClientAuthSender outgoing, Iterator<byte[]> incoming) {
+  public void authenticate(ClientAuthSender outgoing, Iterator<byte[]> incoming) {
     BasicAuth.Builder builder = BasicAuth.newBuilder();
     if (name != null) {
       builder.setUsername(name);
@@ -46,7 +47,12 @@ public class BasicClientAuthHandler implements ClientAuthHandler {
     }
 
     outgoing.send(builder.build().toByteArray());
-    return incoming.next();
+    this.token = incoming.next();
+  }
+
+  @Override
+  public byte[] getCallToken() {
+    return token;
   }
 
 }

java/flight/src/main/java/org/apache/arrow/flight/auth/BasicServerAuthHandler.java

@@ -18,8 +18,10 @@
 package org.apache.arrow.flight.auth;
 
 import java.util.Iterator;
+import java.util.Optional;
 
 import org.apache.arrow.flight.impl.Flight.BasicAuth;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -42,7 +44,7 @@ public class BasicServerAuthHandler implements ServerAuthHandler {
 
     public byte[] getToken(String username, String password) throws Exception;
 
-    public boolean isValid(byte[] token);
+    public Optional<String> isValid(byte[] token);
 
   }
 
@@ -64,7 +66,7 @@ public class BasicServerAuthHandler implements ServerAuthHandler {
   }
 
   @Override
-  public boolean isValid(byte[] token) {
+  public Optional<String> isValid(byte[] token) {
     return authValidator.isValid(token);
   }
 

java/flight/src/main/java/org/apache/arrow/flight/auth/ClientAuthHandler.java

@@ -19,15 +19,30 @@ package org.apache.arrow.flight.auth;
 
 import java.util.Iterator;
 
+/**
+ * Implement authentication for Flight on the client side.
+ */
 public interface ClientAuthHandler {
+  /**
+   * Handle the initial handshake with the server.
+   * @param outgoing A channel to send data to the server.
+   * @param incoming An iterator of incoming data from the server.
+   */
+  void authenticate(ClientAuthSender outgoing, Iterator<byte[]> incoming);
 
-  public byte[] authenticate(ClientAuthSender outgoing, Iterator<byte[]> incoming);
+  /**
+   * Get the per-call authentication token.
+   */
+  byte[] getCallToken();
 
-  public interface ClientAuthSender {
+  /**
+   * A communication channel to the server during initial connection.
+   */
+  interface ClientAuthSender {
 
-    public void send(byte[] payload);
+    void send(byte[] payload);
 
-    public void onError(String message, Throwable cause);
+    void onError(String message, Throwable cause);
 
   }
 

java/flight/src/main/java/org/apache/arrow/flight/auth/ClientAuthInterceptor.java

@@ -26,20 +26,17 @@ import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
 
 public class ClientAuthInterceptor implements ClientInterceptor {
+  private volatile ClientAuthHandler authHandler = null;
 
-  private volatile Metadata headerWithAuth;
-
-  public void setToken(byte[] token) {
-    final Metadata md = new Metadata();
-    md.put(AuthConstants.TOKEN_KEY, token);
-    this.headerWithAuth = md;
+  public void setAuthHandler(ClientAuthHandler authHandler) {
+    this.authHandler = authHandler;
   }
 
   public ClientAuthInterceptor() {
   }
 
-  public boolean hasToken() {
-    return headerWithAuth != null;
+  public boolean hasAuthHandler() {
+    return authHandler != null;
   }
 
   @Override
@@ -48,7 +45,7 @@ public class ClientAuthInterceptor implements ClientInterceptor {
     ClientCall<ReqT, RespT> call = next.newCall(methodDescriptor, callOptions);
 
     // once we have an auth header, add that to the calls.
-    if (headerWithAuth != null) {
+    if (authHandler != null) {
       call = new HeaderAttachingClientCall<>(call);
     }
 
@@ -63,7 +60,9 @@ public class ClientAuthInterceptor implements ClientInterceptor {
 
     @Override
     public void start(Listener<RespT> responseListener, Metadata headers) {
-      headers.merge(headerWithAuth);
+      final Metadata authHeaders = new Metadata();
+      authHeaders.put(AuthConstants.TOKEN_KEY, authHandler.getCallToken());
+      headers.merge(authHeaders);
       super.start(responseListener, headers);
     }
   }

java/flight/src/main/java/org/apache/arrow/flight/auth/ClientAuthWrapper.java

@@ -38,13 +38,11 @@ public class ClientAuthWrapper {
    * @param stub The service stub.
    * @return The token if auth was successful.
    */
-  public static byte[] doClientAuth(ClientAuthHandler authHandler, FlightServiceStub stub) {
-
+  public static void doClientAuth(ClientAuthHandler authHandler, FlightServiceStub stub) {
     AuthObserver observer = new AuthObserver();
     observer.responseObserver = stub.handshake(observer);
-    byte[] bytes = authHandler.authenticate(observer.sender, observer.iter);
+    authHandler.authenticate(observer.sender, observer.iter);
     observer.responseObserver.onCompleted();
-    return bytes;
   }
 
   private static class AuthObserver implements StreamObserver<HandshakeResponse> {

java/flight/src/main/java/org/apache/arrow/flight/auth/ServerAuthHandler.java

@@ -18,12 +18,24 @@
 package org.apache.arrow.flight.auth;
 
 import java.util.Iterator;
+import java.util.Optional;
 
 public interface ServerAuthHandler {
 
-  public boolean isValid(byte[] token);
+  /**
+   * Validate the client token provided on each call.
+   *
+   * @return An empty optional if the client is not authenticated; the peer identity otherwise (may be the empty
+   *     string).
+   */
+  Optional<String> isValid(byte[] token);
 
-  public boolean authenticate(ServerAuthSender outgoing, Iterator<byte[]> incoming);
+  /**
+   * Handle the initial handshake with the client.
+   *
+   * @return true if client is authenticated, false otherwise.
+   */
+  boolean authenticate(ServerAuthSender outgoing, Iterator<byte[]> incoming);
 
   public interface ServerAuthSender {
 
@@ -33,11 +45,14 @@ public interface ServerAuthHandler {
 
   }
 
+  /**
+   * An auth handler that does nothing.
+   */
   ServerAuthHandler NO_OP = new ServerAuthHandler() {
 
     @Override
-    public boolean isValid(byte[] token) {
-      return true;
+    public Optional<String> isValid(byte[] token) {
+      return Optional.of("");
     }
 
     @Override

java/flight/src/main/java/org/apache/arrow/flight/auth/ServerAuthInterceptor.java

@@ -17,6 +17,10 @@
 
 package org.apache.arrow.flight.auth;
 
+import java.util.Optional;
+
+import io.grpc.Context;
+import io.grpc.Contexts;
 import io.grpc.Metadata;
 import io.grpc.ServerCall;
 import io.grpc.ServerCall.Listener;
@@ -35,18 +39,21 @@ public class ServerAuthInterceptor implements ServerInterceptor {
   @Override
   public <ReqT, RespT> Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call, Metadata headers,
       ServerCallHandler<ReqT, RespT> next) {
-    if (
-        !call.getMethodDescriptor().getFullMethodName().equals(AuthConstants.HANDSHAKE_DESCRIPTOR_NAME) &&
-        !isValid(headers)) {
-      call.close(Status.PERMISSION_DENIED, new Metadata());
-      // TODO: we should actually terminate here instead of causing an exception below.
-      return new NoopServerCallListener<>();
+    if (!call.getMethodDescriptor().getFullMethodName().equals(AuthConstants.HANDSHAKE_DESCRIPTOR_NAME)) {
+      final Optional<String> peerIdentity = isValid(headers);
+      if (!peerIdentity.isPresent()) {
+        call.close(Status.PERMISSION_DENIED, new Metadata());
+        // TODO: we should actually terminate here instead of causing an exception below.
+        return new NoopServerCallListener<>();
+      }
+      return Contexts.interceptCall(Context.current().withValue(AuthConstants.PEER_IDENTITY_KEY, peerIdentity.get()),
+          call, headers, next);
     }
 
     return next.startCall(call, headers);
   }
 
-  private final boolean isValid(Metadata headers) {
+  private final Optional<String> isValid(Metadata headers) {
     byte[] token = headers.get(AuthConstants.TOKEN_KEY);
     return authHandler.isValid(token);
   }

java/flight/src/main/java/org/apache/arrow/flight/example/InMemoryStore.java

@@ -55,7 +55,8 @@ public class InMemoryStore implements FlightProducer, AutoCloseable {
   }
 
   @Override
-  public void getStream(Ticket ticket, ServerStreamListener listener) {
+  public void getStream(CallContext context, Ticket ticket,
+      ServerStreamListener listener) {
     getStream(ticket).sendTo(allocator, listener);
   }
 
@@ -79,7 +80,8 @@ public class InMemoryStore implements FlightProducer, AutoCloseable {
   }
 
   @Override
-  public void listFlights(Criteria criteria, StreamListener<FlightInfo> listener) {
+  public void listFlights(CallContext context, Criteria criteria,
+      StreamListener<FlightInfo> listener) {
     try {
       for (FlightHolder h : holders.values()) {
         listener.onNext(h.getFlightInfo(location));
@@ -91,7 +93,8 @@ public class InMemoryStore implements FlightProducer, AutoCloseable {
   }
 
   @Override
-  public FlightInfo getFlightInfo(FlightDescriptor descriptor) {
+  public FlightInfo getFlightInfo(CallContext context,
+      FlightDescriptor descriptor) {
     FlightHolder h = holders.get(descriptor);
     if (h == null) {
       throw new IllegalStateException("Unknown descriptor.");
@@ -101,7 +104,8 @@ public class InMemoryStore implements FlightProducer, AutoCloseable {
   }
 
   @Override
-  public Callable<PutResult> acceptPut(final FlightStream flightStream) {
+  public Callable<PutResult> acceptPut(CallContext context,
+      final FlightStream flightStream) {
     return () -> {
       StreamCreator creator = null;
       boolean success = false;
@@ -130,7 +134,7 @@ public class InMemoryStore implements FlightProducer, AutoCloseable {
   }
 
   @Override
-  public Result doAction(Action action) {
+  public Result doAction(CallContext context, Action action) {
     switch (action.getType()) {
       case "drop":
         return new Result(new byte[0]);
@@ -141,7 +145,8 @@ public class InMemoryStore implements FlightProducer, AutoCloseable {
   }
 
   @Override
-  public void listActions(StreamListener<ActionType> listener) {
+  public void listActions(CallContext context,
+      StreamListener<ActionType> listener) {
     listener.onNext(new ActionType("get", "pull a stream. Action must be done via standard get mechanism"));
     listener.onNext(new ActionType("put", "push a stream. Action must be done via standard put mechanism"));
     listener.onNext(new ActionType("drop", "delete a flight. Action body is a JSON encoded path."));

java/flight/src/test/java/org/apache/arrow/flight/TestBackPressure.java

@@ -89,7 +89,8 @@ public class TestBackPressure {
       final FlightProducer producer = new NoOpFlightProducer() {
 
         @Override
-        public void getStream(Ticket ticket, ServerStreamListener listener) {
+        public void getStream(CallContext context, Ticket ticket,
+            ServerStreamListener listener) {
           int batches = 0;
           final Schema pojoSchema = new Schema(ImmutableList.of(Field.nullable("a", MinorType.BIGINT.getType())));
           VectorSchemaRoot root = VectorSchemaRoot.create(pojoSchema, allocator);

java/flight/src/test/java/org/apache/arrow/flight/TestBasicOperation.java

@@ -164,7 +164,8 @@ public class TestBasicOperation {
     }
 
     @Override
-    public void listFlights(Criteria criteria, StreamListener<FlightInfo> listener) {
+    public void listFlights(CallContext context, Criteria criteria,
+        StreamListener<FlightInfo> listener) {
       FlightGetInfo getInfo = FlightGetInfo.newBuilder()
           .setFlightDescriptor(Flight.FlightDescriptor.newBuilder()
               .setType(DescriptorType.CMD)
@@ -175,7 +176,8 @@ public class TestBasicOperation {
     }
 
     @Override
-    public Callable<PutResult> acceptPut(FlightStream flightStream) {
+    public Callable<PutResult> acceptPut(CallContext context,
+        FlightStream flightStream) {
       return () -> {
         try (VectorSchemaRoot root = flightStream.getRoot()) {
           while (flightStream.next()) {
@@ -187,7 +189,8 @@ public class TestBasicOperation {
     }
 
     @Override
-    public void getStream(Ticket ticket, ServerStreamListener listener) {
+    public void getStream(CallContext context, Ticket ticket,
+        ServerStreamListener listener) {
       final int size = 10;
 
       IntVector iv = new IntVector("c1", allocator);
@@ -222,7 +225,8 @@ public class TestBasicOperation {
     }
 
     @Override
-    public FlightInfo getFlightInfo(FlightDescriptor descriptor) {
+    public FlightInfo getFlightInfo(CallContext context,
+        FlightDescriptor descriptor) {
       FlightGetInfo getInfo = FlightGetInfo.newBuilder()
           .setFlightDescriptor(Flight.FlightDescriptor.newBuilder()
               .setType(DescriptorType.CMD)
@@ -232,7 +236,7 @@ public class TestBasicOperation {
     }
 
     @Override
-    public Result doAction(Action action) {
+    public Result doAction(CallContext context, Action action) {
       switch (action.getType()) {
         case "hello":
           return new Result("world".getBytes(Charsets.UTF_8));
@@ -242,7 +246,8 @@ public class TestBasicOperation {
     }
 
     @Override
-    public void listActions(StreamListener<ActionType> listener) {
+    public void listActions(CallContext context,
+        StreamListener<ActionType> listener) {
       listener.onNext(new ActionType("get", ""));
       listener.onNext(new ActionType("put", ""));
       listener.onNext(new ActionType("hello", ""));

java/flight/src/test/java/org/apache/arrow/flight/TestLargeMessage.java

@@ -17,7 +17,6 @@
 
 package org.apache.arrow.flight;
 
-import java.sql.Types;
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.Callable;
@@ -116,7 +115,8 @@ public class TestLargeMessage {
     }
 
     @Override
-    public void getStream(Ticket ticket, ServerStreamListener listener) {
+    public void getStream(CallContext context, Ticket ticket,
+        ServerStreamListener listener) {
       try (VectorSchemaRoot root = generateData(allocator)) {
         listener.start(root);
         listener.putNext();
@@ -125,17 +125,20 @@ public class TestLargeMessage {
     }
 
     @Override
-    public void listFlights(Criteria criteria, StreamListener<FlightInfo> listener) {
+    public void listFlights(CallContext context, Criteria criteria,
+        StreamListener<FlightInfo> listener) {
 
     }
 
     @Override
-    public FlightInfo getFlightInfo(FlightDescriptor descriptor) {
+    public FlightInfo getFlightInfo(CallContext context,
+        FlightDescriptor descriptor) {
       return null;
     }
 
     @Override
-    public Callable<Flight.PutResult> acceptPut(FlightStream flightStream) {
+    public Callable<Flight.PutResult> acceptPut(CallContext context,
+        FlightStream flightStream) {
       return () -> {
         try (VectorSchemaRoot root = flightStream.getRoot()) {
           while (flightStream.next()) {
@@ -147,12 +150,13 @@ public class TestLargeMessage {
     }
 
     @Override
-    public Result doAction(Action action) {
+    public Result doAction(CallContext context, Action action) {
       return null;
     }
 
     @Override
-    public void listActions(StreamListener<ActionType> listener) {
+    public void listActions(CallContext context,
+        StreamListener<ActionType> listener) {
 
     }
 

java/flight/src/test/java/org/apache/arrow/flight/auth/TestAuth.java

@@ -21,6 +21,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.Optional;
 
 import org.apache.arrow.flight.Criteria;
 import org.apache.arrow.flight.FlightClient;
@@ -32,6 +33,7 @@ import org.apache.arrow.flight.NoOpFlightProducer;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.util.AutoCloseables;
+
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
@@ -56,12 +58,10 @@ public class TestAuth {
   public void validAuth() {
     client.authenticateBasic(USERNAME, PASSWORD);
     Assert.assertTrue(ImmutableList.copyOf(client.listFlights(Criteria.ALL)).size() >= 0);
-
   }
 
   @Test
   public void invalidAuth() {
-
     assertThrows(StatusRuntimeException.class, () -> {
       client.authenticateBasic(USERNAME, "WRONG");
     }, PERMISSION_DENIED);
@@ -84,8 +84,11 @@ public class TestAuth {
     final BasicServerAuthHandler.BasicAuthValidator validator = new BasicServerAuthHandler.BasicAuthValidator() {
 
       @Override
-      public boolean isValid(byte[] token) {
-        return Arrays.equals(token, VALID_TOKEN);
+      public Optional<String> isValid(byte[] token) {
+        if (Arrays.equals(token, VALID_TOKEN)) {
+          return Optional.of(USERNAME);
+        }
+        return Optional.empty();
       }
 
       @Override
@@ -103,7 +106,12 @@ public class TestAuth {
         port,
         new NoOpFlightProducer() {
           @Override
-          public void listFlights(Criteria criteria, StreamListener<FlightInfo> listener) {
+          public void listFlights(CallContext context, Criteria criteria,
+              StreamListener<FlightInfo> listener) {
+            if (!context.peerIdentity().equals(USERNAME)) {
+              listener.onError(new IllegalArgumentException("Invalid username"));
+              return;
+            }
             listener.onCompleted();
           }
         },

java/flight/src/test/java/org/apache/arrow/flight/perf/PerformanceTestServer.java

@@ -83,7 +83,8 @@ public class PerformanceTestServer implements AutoCloseable {
   private final class PerfProducer implements FlightProducer {
 
     @Override
-    public void getStream(Ticket ticket, ServerStreamListener listener) {
+    public void getStream(CallContext context, Ticket ticket,
+        ServerStreamListener listener) {
       VectorSchemaRoot root = null;
       try {
         Token token = Token.parseFrom(ticket.getBytes());
@@ -146,11 +147,13 @@ public class PerformanceTestServer implements AutoCloseable {
     }
 
     @Override
-    public void listFlights(Criteria criteria, StreamListener<FlightInfo> listener) {
+    public void listFlights(CallContext context, Criteria criteria,
+        StreamListener<FlightInfo> listener) {
     }
 
     @Override
-    public FlightInfo getFlightInfo(FlightDescriptor descriptor) {
+    public FlightInfo getFlightInfo(CallContext context,
+        FlightDescriptor descriptor) {
       try {
         Preconditions.checkArgument(descriptor.isCommand());
         Perf exec = Perf.parseFrom(descriptor.getCommand());
@@ -181,17 +184,19 @@ public class PerformanceTestServer implements AutoCloseable {
     }
 
     @Override
-    public Callable<PutResult> acceptPut(FlightStream flightStream) {
+    public Callable<PutResult> acceptPut(CallContext context,
+        FlightStream flightStream) {
       return null;
     }
 
     @Override
-    public Result doAction(Action action) {
+    public Result doAction(CallContext context, Action action) {
       return null;
     }
 
     @Override
-    public void listActions(StreamListener<ActionType> listener) {
+    public void listActions(CallContext context,
+        StreamListener<ActionType> listener) {
     }
 
   }