1 /*
2  * Copyright (C) 2018 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package com.android.tradefed.util;
17 
18 import com.android.tradefed.auth.ICredentialFactory;
19 import com.android.tradefed.config.GlobalConfiguration;
20 import com.android.tradefed.host.HostOptions;
21 import com.android.tradefed.log.LogUtil.CLog;
22 
23 import com.google.api.client.auth.oauth2.Credential;
24 import com.google.api.client.http.HttpBackOffUnsuccessfulResponseHandler;
25 import com.google.api.client.http.HttpRequest;
26 import com.google.api.client.http.HttpRequestInitializer;
27 import com.google.api.client.http.HttpResponse;
28 import com.google.api.client.http.HttpUnsuccessfulResponseHandler;
29 import com.google.api.client.util.ExponentialBackOff;
30 import com.google.auth.Credentials;
31 import com.google.auth.oauth2.ComputeEngineCredentials;
32 import com.google.auth.oauth2.GoogleCredentials;
33 import com.google.common.annotations.VisibleForTesting;
34 
35 import java.io.File;
36 import java.io.FileInputStream;
37 import java.io.FileNotFoundException;
38 import java.io.IOException;
39 import java.security.GeneralSecurityException;
40 import java.util.ArrayList;
41 import java.util.Arrays;
42 import java.util.Collection;
43 import java.util.List;
44 
45 /** Utils for create Google API client. */
46 public class GoogleApiClientUtil {
47 
48     public static final String APP_NAME = "tradefed";
49     private static GoogleApiClientUtil sInstance = null;
50 
getInstance()51     private static GoogleApiClientUtil getInstance() {
52         if (sInstance == null) {
53             sInstance = new GoogleApiClientUtil();
54         }
55         return sInstance;
56     }
57 
58     /**
59      * Create credential from json key file.
60      *
61      * @param file is the p12 key file
62      * @param scopes is the API's scope.
63      * @return a {@link Credential}.
64      * @throws FileNotFoundException
65      * @throws IOException
66      * @throws GeneralSecurityException
67      */
createCredentialFromJsonKeyFile(File file, Collection<String> scopes)68     public static Credentials createCredentialFromJsonKeyFile(File file, Collection<String> scopes)
69             throws IOException, GeneralSecurityException {
70         return getInstance().doCreateCredentialFromJsonKeyFile(file, scopes);
71     }
72 
73     @VisibleForTesting
doCreateCredentialFromJsonKeyFile(File file, Collection<String> scopes)74     Credentials doCreateCredentialFromJsonKeyFile(File file, Collection<String> scopes)
75             throws IOException, GeneralSecurityException {
76         Credentials credentail =
77                 GoogleCredentials.fromStream(new FileInputStream(file)).createScoped(scopes);
78         return credentail;
79     }
80 
81     /**
82      * Try to create credential with different key files or from local host.
83      *
84      * <p>1. If primaryKeyFile is set, try to use it to create credential. 2. Try to get
85      * corresponding key files from {@link HostOptions}. 3. Try to use backup key files. 4. Use
86      * local default credential.
87      *
88      * @param scopes scopes for the credential.
89      * @param primaryKeyFile the primary json key file; it can be null.
90      * @param hostOptionKeyFileName {@link HostOptions}'service-account-json-key-file option's key;
91      *     it can be null.
92      * @param backupKeyFiles backup key files.
93      * @return a {@link Credential}
94      * @throws IOException
95      * @throws GeneralSecurityException
96      */
createCredential( Collection<String> scopes, File primaryKeyFile, String hostOptionKeyFileName, File... backupKeyFiles)97     public static Credentials createCredential(
98             Collection<String> scopes,
99             File primaryKeyFile,
100             String hostOptionKeyFileName,
101             File... backupKeyFiles)
102             throws IOException, GeneralSecurityException {
103         return getInstance()
104                 .doCreateCredential(scopes, primaryKeyFile, hostOptionKeyFileName, backupKeyFiles);
105     }
106 
107     /**
108      * Try to create credential with different key files or from local host.
109      *
110      * <p>1. Use {@link ICredentialFactory} if useCredentialFactory is true and a {@link
111      * ICredentialFactory} is configured. If primaryKeyFile is set, try to use it to create
112      * credential. 2. Try to get corresponding key files from {@link HostOptions}. 3. Try to use
113      * backup key files. 4. Use local default credential.
114      *
115      * @param scopes scopes for the credential.
116      * @param useCredentialFactory use credential factory if it's configured.
117      * @param primaryKeyFile the primary json key file; it can be null.
118      * @param hostOptionKeyFileName {@link HostOptions}'service-account-json-key-file option's key;
119      *     it can be null.
120      * @param backupKeyFiles backup key files.
121      * @return a {@link Credential}
122      * @throws IOException
123      * @throws GeneralSecurityException
124      */
createCredential( Collection<String> scopes, boolean useCredentialFactory, File primaryKeyFile, String hostOptionKeyFileName, File... backupKeyFiles)125     public static Credentials createCredential(
126             Collection<String> scopes,
127             boolean useCredentialFactory,
128             File primaryKeyFile,
129             String hostOptionKeyFileName,
130             File... backupKeyFiles)
131             throws IOException, GeneralSecurityException {
132         Credentials credential = null;
133         if (useCredentialFactory) {
134             credential = getInstance().doCreateCredentialFromCredentialFactory(scopes);
135             // TODO(b/186766552): Throw exception once all hosts configured CredentialFactory.
136             if (credential != null) {
137                 return credential;
138             }
139             CLog.i("No CredentialFactory configured, fallback to key files.");
140         }
141         return getInstance()
142                 .doCreateCredential(scopes, primaryKeyFile, hostOptionKeyFileName, backupKeyFiles);
143     }
144 
145     @VisibleForTesting
doCreateCredential( Collection<String> scopes, File primaryKeyFile, String hostOptionKeyFileName, File... backupKeyFiles)146     Credentials doCreateCredential(
147             Collection<String> scopes,
148             File primaryKeyFile,
149             String hostOptionKeyFileName,
150             File... backupKeyFiles)
151             throws IOException, GeneralSecurityException {
152 
153         List<File> keyFiles = new ArrayList<File>();
154         if (primaryKeyFile != null) {
155             keyFiles.add(primaryKeyFile);
156         }
157         File hostOptionKeyFile = null;
158         if (hostOptionKeyFileName != null) {
159             try {
160                 hostOptionKeyFile =
161                         GlobalConfiguration.getInstance()
162                                 .getHostOptions()
163                                 .getServiceAccountJsonKeyFiles()
164                                 .get(hostOptionKeyFileName);
165                 if (hostOptionKeyFile != null) {
166                     keyFiles.add(hostOptionKeyFile);
167                 }
168             } catch (IllegalStateException e) {
169                 CLog.d("Global configuration haven't been initialized.");
170             }
171         }
172         keyFiles.addAll(Arrays.asList(backupKeyFiles));
173         for (File keyFile : keyFiles) {
174             if (keyFile != null) {
175                 if (keyFile.exists() && keyFile.canRead()) {
176                     CLog.d("Using %s.", keyFile.getAbsolutePath());
177                     return doCreateCredentialFromJsonKeyFile(keyFile, scopes);
178                 } else {
179                     CLog.i("No access to %s.", keyFile.getAbsolutePath());
180                 }
181             }
182         }
183         return doCreateDefaultCredential(scopes);
184     }
185 
186     @VisibleForTesting
doCreateCredentialFromCredentialFactory(Collection<String> scopes)187     Credentials doCreateCredentialFromCredentialFactory(Collection<String> scopes)
188             throws IOException {
189         try {
190             if (GlobalConfiguration.getInstance().getCredentialFactory() != null) {
191                 return GlobalConfiguration.getInstance()
192                         .getCredentialFactory()
193                         .createCredential(scopes);
194             }
195             CLog.w("No CredentialFactory configured.");
196         } catch (IllegalStateException e) {
197             System.out.println(
198                     "GlobalConfiguration is not initialized yet,"
199                             + "can not get CredentialFactory.");
200         }
201         return null;
202     }
203 
204     @VisibleForTesting
doCreateDefaultCredential(Collection<String> scopes)205     Credentials doCreateDefaultCredential(Collection<String> scopes) throws IOException {
206         try {
207             CLog.d("Using local authentication.");
208             return ComputeEngineCredentials.getApplicationDefault().createScoped(scopes);
209         } catch (IOException e) {
210             CLog.e(
211                     "Try 'gcloud auth application-default login' to login for "
212                             + "personal account; Or 'export "
213                             + "GOOGLE_APPLICATION_CREDENTIALS=/path/to/key.json' "
214                             + "for service account.");
215             throw e;
216         }
217     }
218 
219     /**
220      * @param requestInitializer a {@link HttpRequestInitializer}, normally it's {@link Credential}.
221      * @param connectTimeout connect timeout in milliseconds.
222      * @param readTimeout read timeout in milliseconds.
223      * @return a {@link HttpRequestInitializer} with timeout.
224      */
setHttpTimeout( final HttpRequestInitializer requestInitializer, int connectTimeout, int readTimeout)225     public static HttpRequestInitializer setHttpTimeout(
226             final HttpRequestInitializer requestInitializer, int connectTimeout, int readTimeout) {
227         return new HttpRequestInitializer() {
228             @Override
229             public void initialize(HttpRequest request) throws IOException {
230                 requestInitializer.initialize(request);
231                 request.setConnectTimeout(connectTimeout);
232                 request.setReadTimeout(readTimeout);
233             }
234         };
235     }
236 
237     /**
238      * Setup a retry strategy for the provided HttpRequestInitializer. In case of server errors
239      * requests will be automatically retried with an exponential backoff.
240      *
241      * @param initializer - an initializer which will setup a retry strategy.
242      * @return an initializer that will retry failed requests automatically.
243      */
244     public static HttpRequestInitializer configureRetryStrategyAndTimeout(
245             HttpRequestInitializer initializer, int connectTimeout, int readTimeout) {
246         return new HttpRequestInitializer() {
247             @Override
248             public void initialize(HttpRequest request) throws IOException {
249                 initializer.initialize(request);
250                 request.setConnectTimeout(connectTimeout);
251                 request.setReadTimeout(readTimeout);
252                 request.setUnsuccessfulResponseHandler(new RetryResponseHandler());
253             }
254         };
255     }
256 
257     /**
258      * Setup a retry strategy for the provided HttpRequestInitializer. In case of server errors
259      * requests will be automatically retried with an exponential backoff.
260      *
261      * @param initializer - an initializer which will setup a retry strategy.
262      * @return an initializer that will retry failed requests automatically.
263      */
264     public static HttpRequestInitializer configureRetryStrategy(
265             HttpRequestInitializer initializer) {
266         return new HttpRequestInitializer() {
267             @Override
268             public void initialize(HttpRequest request) throws IOException {
269                 initializer.initialize(request);
270                 request.setUnsuccessfulResponseHandler(new RetryResponseHandler());
271             }
272         };
273     }
274 
275     private static class RetryResponseHandler implements HttpUnsuccessfulResponseHandler {
276         // Initial interval to wait before retrying if a request fails.
277         private static final int INITIAL_RETRY_INTERVAL = 1000;
278         private static final int MAX_RETRY_INTERVAL = 3 * 60000; // Set max interval to 3 minutes.
279 
280         private final HttpUnsuccessfulResponseHandler backOffHandler;
281 
282         public RetryResponseHandler() {
283             backOffHandler =
284                     new HttpBackOffUnsuccessfulResponseHandler(
285                             new ExponentialBackOff.Builder()
286                                     .setInitialIntervalMillis(INITIAL_RETRY_INTERVAL)
287                                     .setMaxIntervalMillis(MAX_RETRY_INTERVAL)
288                                     .build());
289         }
290 
291         @Override
292         public boolean handleResponse(
293                 HttpRequest request, HttpResponse response, boolean supportsRetry)
294                 throws IOException {
295             CLog.w(
296                     "Request to %s failed: %d %s",
297                     request.getUrl(), response.getStatusCode(), response.getStatusMessage());
298             if (response.getStatusCode() == 400) {
299                 return false;
300             }
301             return backOffHandler.handleResponse(request, response, supportsRetry);
302         }
303     }
304 }
305