]>
Commit | Line | Data |
---|---|---|
1d09f67e TL |
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 | } |