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 
17 package com.android.tradefed.util;
18 
19 import com.android.tradefed.build.BuildRetrievalError;
20 import com.android.tradefed.build.IFileDownloader;
21 import com.android.tradefed.log.LogUtil.CLog;
22 import com.android.tradefed.result.error.InfraErrorIdentifier;
23 
24 import com.google.api.client.googleapis.json.GoogleJsonResponseException;
25 import com.google.api.services.storage.Storage;
26 import com.google.api.services.storage.model.Objects;
27 import com.google.api.services.storage.model.StorageObject;
28 import com.google.common.annotations.VisibleForTesting;
29 
30 import java.io.ByteArrayInputStream;
31 import java.io.ByteArrayOutputStream;
32 import java.io.File;
33 import java.io.FileOutputStream;
34 import java.io.IOException;
35 import java.io.InputStream;
36 import java.io.OutputStream;
37 import java.net.SocketException;
38 import java.nio.file.Paths;
39 import java.util.ArrayList;
40 import java.util.Arrays;
41 import java.util.Collection;
42 import java.util.Collections;
43 import java.util.HashSet;
44 import java.util.List;
45 import java.util.Set;
46 import java.util.regex.Matcher;
47 import java.util.regex.Pattern;
48 
49 /** File downloader to download file from google cloud storage (GCS). */
50 public class GCSFileDownloader extends GCSCommon implements IFileDownloader {
51     public static final String GCS_PREFIX = "gs://";
52     public static final String GCS_APPROX_PREFIX = "gs:/";
53 
54     private static final Pattern GCS_PATH_PATTERN = Pattern.compile("gs://([^/]*)/(.*)");
55     private static final String PATH_SEP = "/";
56     private static final Collection<String> SCOPES =
57             Collections.singleton("https://www.googleapis.com/auth/devstorage.read_only");
58     private static final long LIST_BATCH_SIZE = 100;
59 
GCSFileDownloader(File jsonKeyFile)60     public GCSFileDownloader(File jsonKeyFile) {
61         super(jsonKeyFile);
62     }
63 
GCSFileDownloader()64     public GCSFileDownloader() {}
65 
66     /**
67      * Download a file from a GCS bucket file.
68      *
69      * @param bucketName GCS bucket name
70      * @param filename the filename
71      * @return {@link InputStream} with the file content.
72      */
downloadFile(String bucketName, String filename)73     public InputStream downloadFile(String bucketName, String filename) throws IOException {
74         InputStream remoteInput = null;
75         ByteArrayOutputStream tmpStream = null;
76         try {
77             remoteInput =
78                     getStorage().objects().get(bucketName, filename).executeMediaAsInputStream();
79             // The input stream from api call can not be reset. Change it to ByteArrayInputStream.
80             tmpStream = new ByteArrayOutputStream();
81             StreamUtil.copyStreams(remoteInput, tmpStream);
82             return new ByteArrayInputStream(tmpStream.toByteArray());
83         } finally {
84             StreamUtil.close(remoteInput);
85             StreamUtil.close(tmpStream);
86         }
87     }
88 
getStorage()89     private Storage getStorage() throws IOException {
90         return getStorage(SCOPES);
91     }
92 
93     @VisibleForTesting
getRemoteFileMetaData(String bucketName, String remoteFilename)94     StorageObject getRemoteFileMetaData(String bucketName, String remoteFilename)
95             throws IOException {
96         try {
97             return getStorage().objects().get(bucketName, remoteFilename).execute();
98         } catch (GoogleJsonResponseException e) {
99             if (e.getStatusCode() == 404) {
100                 return null;
101             }
102             throw e;
103         }
104     }
105 
106     /**
107      * Download file from GCS.
108      *
109      * <p>Right now only support GCS path.
110      *
111      * @param remoteFilePath gs://bucket/file/path format GCS path.
112      * @return local file
113      * @throws BuildRetrievalError
114      */
115     @Override
downloadFile(String remoteFilePath)116     public File downloadFile(String remoteFilePath) throws BuildRetrievalError {
117         File destFile = createTempFile(remoteFilePath, null);
118         try {
119             downloadFile(remoteFilePath, destFile);
120             return destFile;
121         } catch (BuildRetrievalError e) {
122             FileUtil.recursiveDelete(destFile);
123             throw e;
124         }
125     }
126 
127     @Override
downloadFile(String remotePath, File destFile)128     public void downloadFile(String remotePath, File destFile) throws BuildRetrievalError {
129         String[] pathParts = parseGcsPath(remotePath);
130         downloadFile(pathParts[0], pathParts[1], destFile);
131     }
132 
isFileFresh(File localFile, StorageObject remoteFile)133     private boolean isFileFresh(File localFile, StorageObject remoteFile) throws IOException {
134         if (localFile == null && remoteFile == null) {
135             return true;
136         }
137         if (localFile == null || remoteFile == null) {
138             return false;
139         }
140         if (!localFile.exists()) {
141             return false;
142         }
143         return remoteFile.getMd5Hash().equals(FileUtil.calculateBase64Md5(localFile));
144     }
145 
146     @Override
isFresh(File localFile, String remotePath)147     public boolean isFresh(File localFile, String remotePath) throws BuildRetrievalError {
148         String[] pathParts = parseGcsPath(remotePath);
149         String bucketName = pathParts[0];
150         String remoteFilename = pathParts[1];
151         try {
152             StorageObject remoteFileMeta = getRemoteFileMetaData(bucketName, remoteFilename);
153             if (localFile == null || !localFile.exists()) {
154                 if (!isRemoteFolder(bucketName, remoteFilename) && remoteFileMeta == null) {
155                     // The local doesn't exist and the remote filename is not a folder or a file.
156                     return true;
157                 }
158                 return false;
159             }
160             if (!localFile.isDirectory()) {
161                 return isFileFresh(localFile, remoteFileMeta);
162             }
163             remoteFilename = sanitizeDirectoryName(remoteFilename);
164             return recursiveCheckFolderFreshness(bucketName, remoteFilename, localFile);
165         } catch (IOException e) {
166             throw new BuildRetrievalError(e.getMessage(), e);
167         }
168     }
169 
170     /**
171      * Check if remote folder is the same as local folder, recursively. The remoteFolderName must
172      * end with "/".
173      *
174      * @param bucketName is the gcs bucket name.
175      * @param remoteFolderName is the relative path to the bucket.
176      * @param localFolder is the local folder
177      * @return true if local file is the same as remote file, otherwise false.
178      * @throws IOException
179      */
recursiveCheckFolderFreshness( String bucketName, String remoteFolderName, File localFolder)180     private boolean recursiveCheckFolderFreshness(
181             String bucketName, String remoteFolderName, File localFolder) throws IOException {
182         Set<String> subFilenames = new HashSet<>(Arrays.asList(localFolder.list()));
183         List<String> subRemoteFolders = new ArrayList<>();
184         List<StorageObject> subRemoteFiles = new ArrayList<>();
185         listRemoteFilesUnderFolder(bucketName, remoteFolderName, subRemoteFiles, subRemoteFolders);
186         for (StorageObject subRemoteFile : subRemoteFiles) {
187             String subFilename = Paths.get(subRemoteFile.getName()).getFileName().toString();
188             if (!isFileFresh(new File(localFolder, subFilename), subRemoteFile)) {
189                 return false;
190             }
191             subFilenames.remove(subFilename);
192         }
193         for (String subRemoteFolder : subRemoteFolders) {
194             String subFolderName = Paths.get(subRemoteFolder).getFileName().toString();
195             File subFolder = new File(localFolder, subFolderName);
196             if (new File(localFolder, subFolderName).exists()
197                     && !new File(localFolder, subFolderName).isDirectory()) {
198                 CLog.w("%s exists as a non-directory.", subFolder);
199                 subFolder = new File(localFolder, subFolderName + "_folder");
200             }
201             if (!recursiveCheckFolderFreshness(bucketName, subRemoteFolder, subFolder)) {
202                 return false;
203             }
204             subFilenames.remove(subFolder.getName());
205         }
206         return subFilenames.isEmpty();
207     }
208 
listRemoteFilesUnderFolder( String bucketName, String folder, List<StorageObject> subFiles, List<String> subFolders)209     void listRemoteFilesUnderFolder(
210             String bucketName, String folder, List<StorageObject> subFiles, List<String> subFolders)
211             throws IOException {
212         String pageToken = null;
213         while (true) {
214             com.google.api.services.storage.Storage.Objects.List listOperation =
215                     getStorage()
216                             .objects()
217                             .list(bucketName)
218                             .setPrefix(folder)
219                             .setDelimiter(PATH_SEP)
220                             .setMaxResults(LIST_BATCH_SIZE);
221             if (pageToken != null) {
222                 listOperation.setPageToken(pageToken);
223             }
224             Objects objects = listOperation.execute();
225             if (objects.getItems() != null && !objects.getItems().isEmpty()) {
226                 subFiles.addAll(objects.getItems());
227             }
228             if (objects.getPrefixes() != null && !objects.getPrefixes().isEmpty()) {
229                 subFolders.addAll(objects.getPrefixes());
230             }
231             pageToken = objects.getNextPageToken();
232             if (pageToken == null) {
233                 return;
234             }
235         }
236     }
237 
parseGcsPath(String remotePath)238     String[] parseGcsPath(String remotePath) throws BuildRetrievalError {
239         if (remotePath.startsWith(GCS_APPROX_PREFIX) && !remotePath.startsWith(GCS_PREFIX)) {
240             // File object remove double // so we have to rebuild it in some cases
241             remotePath = remotePath.replaceAll(GCS_APPROX_PREFIX, GCS_PREFIX);
242         }
243         Matcher m = GCS_PATH_PATTERN.matcher(remotePath);
244         if (!m.find()) {
245             throw new BuildRetrievalError(
246                     String.format("Only GCS path is supported, %s is not supported", remotePath),
247                     InfraErrorIdentifier.ARTIFACT_UNSUPPORTED_PATH);
248         }
249         return new String[] {m.group(1), m.group(2)};
250     }
251 
sanitizeDirectoryName(String name)252     String sanitizeDirectoryName(String name) {
253         /** Folder name should end with "/" */
254         if (!name.endsWith(PATH_SEP)) {
255             name += PATH_SEP;
256         }
257         return name;
258     }
259 
260     /** check given filename is a folder or not. */
261     @VisibleForTesting
isRemoteFolder(String bucketName, String filename)262     boolean isRemoteFolder(String bucketName, String filename) throws IOException {
263         filename = sanitizeDirectoryName(filename);
264         Objects objects =
265                 getStorage()
266                         .objects()
267                         .list(bucketName)
268                         .setPrefix(filename)
269                         .setDelimiter(PATH_SEP)
270                         .setMaxResults(1l)
271                         .execute();
272         if (objects.getItems() != null && !objects.getItems().isEmpty()) {
273             return true;
274         }
275         if (objects.getPrefixes() != null && !objects.getPrefixes().isEmpty()) {
276             return true;
277         }
278         return false;
279     }
280 
281     @VisibleForTesting
downloadFile(String bucketName, String remoteFilename, File localFile)282     void downloadFile(String bucketName, String remoteFilename, File localFile)
283             throws BuildRetrievalError {
284         int i = 0;
285         try {
286             do {
287                 i++;
288                 try {
289                     if (!isRemoteFolder(bucketName, remoteFilename)) {
290                         fetchRemoteFile(bucketName, remoteFilename, localFile);
291                         return;
292                     }
293                     remoteFilename = sanitizeDirectoryName(remoteFilename);
294                     recursiveDownloadFolder(bucketName, remoteFilename, localFile);
295                     return;
296                 } catch (SocketException se) {
297                     // Allow one retry in case of flaky connection.
298                     if (i >= 2) {
299                         throw se;
300                     }
301                     CLog.e(
302                             "Error '%s' while downloading gs://%s/%s. retrying.",
303                             se.getMessage(), bucketName, remoteFilename);
304                 }
305             } while (true);
306         } catch (IOException e) {
307             String message =
308                     String.format(
309                             "Failed to download gs://%s/%s due to: %s",
310                             bucketName, remoteFilename, e.getMessage());
311             CLog.e(message);
312             throw new BuildRetrievalError(message, e, InfraErrorIdentifier.GCS_ERROR);
313         }
314     }
315 
fetchRemoteFile(String bucketName, String remoteFilename, File localFile)316     private void fetchRemoteFile(String bucketName, String remoteFilename, File localFile)
317             throws IOException {
318         try (OutputStream writeStream = new FileOutputStream(localFile)) {
319             getStorage()
320                     .objects()
321                     .get(bucketName, remoteFilename)
322                     .executeMediaAndDownloadTo(writeStream);
323         }
324     }
325 
326     /**
327      * Recursively download remote folder to local folder.
328      *
329      * @param bucketName the gcs bucket name
330      * @param remoteFolderName remote folder name, must end with "/"
331      * @param localFolder local folder
332      * @throws IOException
333      */
recursiveDownloadFolder( String bucketName, String remoteFolderName, File localFolder)334     private void recursiveDownloadFolder(
335             String bucketName, String remoteFolderName, File localFolder) throws IOException {
336         CLog.d("Downloading folder gs://%s/%s.", bucketName, remoteFolderName);
337         if (!localFolder.exists()) {
338             FileUtil.mkdirsRWX(localFolder);
339         }
340         if (!localFolder.isDirectory()) {
341             String error =
342                     String.format(
343                             "%s is not a folder. (gs://%s/%s)",
344                             localFolder, bucketName, remoteFolderName);
345             CLog.e(error);
346             throw new IOException(error);
347         }
348         Set<String> subFilenames = new HashSet<>(Arrays.asList(localFolder.list()));
349         List<String> subRemoteFolders = new ArrayList<>();
350         List<StorageObject> subRemoteFiles = new ArrayList<>();
351         listRemoteFilesUnderFolder(bucketName, remoteFolderName, subRemoteFiles, subRemoteFolders);
352         for (StorageObject subRemoteFile : subRemoteFiles) {
353             String subFilename = Paths.get(subRemoteFile.getName()).getFileName().toString();
354             fetchRemoteFile(
355                     bucketName, subRemoteFile.getName(), new File(localFolder, subFilename));
356             subFilenames.remove(subFilename);
357         }
358         for (String subRemoteFolder : subRemoteFolders) {
359             String subFolderName = Paths.get(subRemoteFolder).getFileName().toString();
360             File subFolder = new File(localFolder, subFolderName);
361             if (new File(localFolder, subFolderName).exists()
362                     && !new File(localFolder, subFolderName).isDirectory()) {
363                 CLog.w("%s exists as a non-directory.", subFolder);
364                 subFolder = new File(localFolder, subFolderName + "_folder");
365             }
366             recursiveDownloadFolder(bucketName, subRemoteFolder, subFolder);
367             subFilenames.remove(subFolder.getName());
368         }
369         for (String subFilename : subFilenames) {
370             FileUtil.recursiveDelete(new File(localFolder, subFilename));
371         }
372     }
373 
374     /**
375      * Creates a unique file on temporary disk to house downloaded file with given path.
376      *
377      * <p>Constructs the file name based on base file name from path
378      *
379      * @param remoteFilePath the remote path to construct the name from
380      */
381     @VisibleForTesting
createTempFile(String remoteFilePath, File rootDir)382     File createTempFile(String remoteFilePath, File rootDir) throws BuildRetrievalError {
383         try {
384             // create a unique file.
385             File tmpFile = FileUtil.createTempFileForRemote(remoteFilePath, rootDir);
386             // now delete it so name is available
387             tmpFile.delete();
388             return tmpFile;
389         } catch (IOException e) {
390             String msg = String.format("Failed to create tmp file for %s", remoteFilePath);
391             throw new BuildRetrievalError(msg, e);
392         }
393     }
394 }
395