ceph/src/arrow/java/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/GeneratedBearerTokenAuthenticator.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.auth2;
  19
  20 import java.nio.ByteBuffer;
  21 import java.util.Base64;
  22 import java.util.UUID;
  23 import java.util.concurrent.TimeUnit;
  24
  25 import org.apache.arrow.flight.CallHeaders;
  26 import org.apache.arrow.flight.CallStatus;
  27 import org.apache.arrow.flight.grpc.MetadataAdapter;
  28
  29 import com.google.common.base.Strings;
  30 import com.google.common.cache.Cache;
  31 import com.google.common.cache.CacheBuilder;
  32
  33 import io.grpc.Metadata;
  34
  35 /**
  36  * Generates and caches bearer tokens from user credentials.
  37  */
  38 public class GeneratedBearerTokenAuthenticator extends BearerTokenAuthenticator {
  39   private final Cache<String, String> bearerToIdentityCache;
  40
  41   /**
  42    * Generate bearer tokens for the given basic call authenticator.
  43    * @param authenticator The authenticator to initial validate inputs with.
  44    */
  45   public GeneratedBearerTokenAuthenticator(CallHeaderAuthenticator authenticator) {
  46     this(authenticator, CacheBuilder.newBuilder().expireAfterAccess(2, TimeUnit.HOURS));
  47   }
  48
  49   /**
  50    * Generate bearer tokens for the given basic call authenticator.
  51    * @param authenticator The authenticator to initial validate inputs with.
  52    * @param timeoutMinutes The time before tokens expire after being accessed.
  53    */
  54   public GeneratedBearerTokenAuthenticator(CallHeaderAuthenticator authenticator, int timeoutMinutes) {
  55     this(authenticator, CacheBuilder.newBuilder().expireAfterAccess(timeoutMinutes, TimeUnit.MINUTES));
  56   }
  57
  58   /**
  59    * Generate bearer tokens for the given basic call authenticator.
  60    * @param authenticator The authenticator to initial validate inputs with.
  61    * @param cacheBuilder The configuration of the cache of bearer tokens.
  62    */
  63   public GeneratedBearerTokenAuthenticator(CallHeaderAuthenticator authenticator,
  64       CacheBuilder<Object, Object> cacheBuilder) {
  65     super(authenticator);
  66     bearerToIdentityCache = cacheBuilder.build();
  67   }
  68
  69   @Override
  70   protected AuthResult validateBearer(String bearerToken) {
  71     final String peerIdentity = bearerToIdentityCache.getIfPresent(bearerToken);
  72     if (peerIdentity == null) {
  73       throw CallStatus.UNAUTHENTICATED.toRuntimeException();
  74     }
  75
  76     return new AuthResult() {
  77       @Override
  78       public String getPeerIdentity() {
  79         return peerIdentity;
  80       }
  81
  82       @Override
  83       public void appendToOutgoingHeaders(CallHeaders outgoingHeaders) {
  84         if (null == AuthUtilities.getValueFromAuthHeader(outgoingHeaders, Auth2Constants.BEARER_PREFIX)) {
  85           outgoingHeaders.insert(Auth2Constants.AUTHORIZATION_HEADER, Auth2Constants.BEARER_PREFIX + bearerToken);
  86         }
  87       }
  88     };
  89   }
  90
  91   @Override
  92   protected AuthResult getAuthResultWithBearerToken(AuthResult authResult) {
  93     // We generate a dummy header and call appendToOutgoingHeaders with it.
  94     // We then inspect the dummy header and parse the bearer token if present in the header
  95     // and generate a new bearer token if a bearer token is not present in the header.
  96     final CallHeaders dummyHeaders = new MetadataAdapter(new Metadata());
  97     authResult.appendToOutgoingHeaders(dummyHeaders);
  98     String bearerToken =
  99             AuthUtilities.getValueFromAuthHeader(dummyHeaders, Auth2Constants.BEARER_PREFIX);
 100     final AuthResult authResultWithBearerToken;
 101     if (Strings.isNullOrEmpty(bearerToken)) {
 102       // Generate a new bearer token and return an AuthResult that can write it.
 103       final UUID uuid = UUID.randomUUID();
 104       final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[16]);
 105       byteBuffer.putLong(uuid.getMostSignificantBits());
 106       byteBuffer.putLong(uuid.getLeastSignificantBits());
 107       final String newToken = Base64.getEncoder().encodeToString(byteBuffer.array());
 108       bearerToken = newToken;
 109       authResultWithBearerToken = new AuthResult() {
 110         @Override
 111         public String getPeerIdentity() {
 112           return authResult.getPeerIdentity();
 113         }
 114
 115         @Override
 116         public void appendToOutgoingHeaders(CallHeaders outgoingHeaders) {
 117           authResult.appendToOutgoingHeaders(outgoingHeaders);
 118           outgoingHeaders.insert(Auth2Constants.AUTHORIZATION_HEADER, Auth2Constants.BEARER_PREFIX + newToken);
 119         }
 120       };
 121     } else {
 122       // Use the bearer token supplied by the original auth result.
 123       authResultWithBearerToken = authResult;
 124     }
 125     bearerToIdentityCache.put(bearerToken, authResult.getPeerIdentity());
 126     return authResultWithBearerToken;
 127   }
 128 }