ceph/src/arrow/java/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java

   1 /*
   2  * Licensed to the Apache Software Foundation (ASF) under one or more
   3  * contributor license agreements.  See the NOTICE file distributed with
   4  * this work for additional information regarding copyright ownership.
   5  * The ASF licenses this file to You under the Apache License, Version 2.0
   6  * (the "License"); you may not use this file except in compliance with
   7  * the License.  You may obtain a copy of the License at
   8  *
   9  *    http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 package org.apache.arrow.flight.auth;
  19
  20 import java.io.IOException;
  21 import java.nio.charset.StandardCharsets;
  22 import java.util.Arrays;
  23 import java.util.Optional;
  24
  25 import org.apache.arrow.flight.Criteria;
  26 import org.apache.arrow.flight.FlightClient;
  27 import org.apache.arrow.flight.FlightInfo;
  28 import org.apache.arrow.flight.FlightServer;
  29 import org.apache.arrow.flight.FlightStatusCode;
  30 import org.apache.arrow.flight.FlightStream;
  31 import org.apache.arrow.flight.FlightTestUtil;
  32 import org.apache.arrow.flight.NoOpFlightProducer;
  33 import org.apache.arrow.flight.Ticket;
  34 import org.apache.arrow.memory.BufferAllocator;
  35 import org.apache.arrow.memory.RootAllocator;
  36 import org.apache.arrow.util.AutoCloseables;
  37 import org.apache.arrow.vector.VectorSchemaRoot;
  38 import org.apache.arrow.vector.types.Types;
  39 import org.apache.arrow.vector.types.pojo.Field;
  40 import org.apache.arrow.vector.types.pojo.Schema;
  41 import org.junit.After;
  42 import org.junit.Assert;
  43 import org.junit.Before;
  44 import org.junit.Ignore;
  45 import org.junit.Test;
  46
  47 import com.google.common.collect.ImmutableList;
  48
  49 public class TestBasicAuth {
  50
  51   private static final String USERNAME = "flight";
  52   private static final String PASSWORD = "woohoo";
  53   private static final byte[] VALID_TOKEN = "my_token".getBytes(StandardCharsets.UTF_8);
  54
  55   private FlightClient client;
  56   private FlightServer server;
  57   private BufferAllocator allocator;
  58
  59   @Test
  60   public void validAuth() {
  61     client.authenticateBasic(USERNAME, PASSWORD);
  62     Assert.assertTrue(ImmutableList.copyOf(client.listFlights(Criteria.ALL)).size() == 0);
  63   }
  64
  65   // ARROW-7722: this test occasionally leaks memory
  66   @Ignore
  67   @Test
  68   public void asyncCall() throws Exception {
  69     client.authenticateBasic(USERNAME, PASSWORD);
  70     client.listFlights(Criteria.ALL);
  71     try (final FlightStream s = client.getStream(new Ticket(new byte[1]))) {
  72       while (s.next()) {
  73         Assert.assertEquals(4095, s.getRoot().getRowCount());
  74       }
  75     }
  76   }
  77
  78   @Test
  79   public void invalidAuth() {
  80     FlightTestUtil.assertCode(FlightStatusCode.UNAUTHENTICATED, () -> {
  81       client.authenticateBasic(USERNAME, "WRONG");
  82     });
  83
  84     FlightTestUtil.assertCode(FlightStatusCode.UNAUTHENTICATED, () -> {
  85       client.listFlights(Criteria.ALL).forEach(action -> Assert.fail());
  86     });
  87   }
  88
  89   @Test
  90   public void didntAuth() {
  91     FlightTestUtil.assertCode(FlightStatusCode.UNAUTHENTICATED, () -> {
  92       client.listFlights(Criteria.ALL).forEach(action -> Assert.fail());
  93     });
  94   }
  95
  96   @Before
  97   public void setup() throws IOException {
  98     allocator = new RootAllocator(Long.MAX_VALUE);
  99     final BasicServerAuthHandler.BasicAuthValidator validator = new BasicServerAuthHandler.BasicAuthValidator() {
 100
 101       @Override
 102       public Optional<String> isValid(byte[] token) {
 103         if (Arrays.equals(token, VALID_TOKEN)) {
 104           return Optional.of(USERNAME);
 105         }
 106         return Optional.empty();
 107       }
 108
 109       @Override
 110       public byte[] getToken(String username, String password) {
 111         if (USERNAME.equals(username) && PASSWORD.equals(password)) {
 112           return VALID_TOKEN;
 113         } else {
 114           throw new IllegalArgumentException("invalid credentials");
 115         }
 116       }
 117     };
 118
 119     server = FlightTestUtil.getStartedServer((location) -> FlightServer.builder(
 120         allocator,
 121         location,
 122         new NoOpFlightProducer() {
 123           @Override
 124           public void listFlights(CallContext context, Criteria criteria,
 125               StreamListener<FlightInfo> listener) {
 126             if (!context.peerIdentity().equals(USERNAME)) {
 127               listener.onError(new IllegalArgumentException("Invalid username"));
 128               return;
 129             }
 130             listener.onCompleted();
 131           }
 132
 133           @Override
 134           public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) {
 135             if (!context.peerIdentity().equals(USERNAME)) {
 136               listener.error(new IllegalArgumentException("Invalid username"));
 137               return;
 138             }
 139             final Schema pojoSchema = new Schema(ImmutableList.of(Field.nullable("a",
 140                 Types.MinorType.BIGINT.getType())));
 141             try (VectorSchemaRoot root = VectorSchemaRoot.create(pojoSchema, allocator)) {
 142               listener.start(root);
 143               root.allocateNew();
 144               root.setRowCount(4095);
 145               listener.putNext();
 146               listener.completed();
 147             }
 148           }
 149         }).authHandler(new BasicServerAuthHandler(validator)).build());
 150     client = FlightClient.builder(allocator, server.getLocation()).build();
 151   }
 152
 153   @After
 154   public void shutdown() throws Exception {
 155     AutoCloseables.close(client, server, allocator);
 156   }
 157
 158 }