blob: b93edce37f0f1cfab3abac0f1036cdc74a5ceaa2 [file] [log] [blame]
/*
* Copyright (C) 2018 The Android Open Source Project
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.android.tradefed.util;
import com.android.tradefed.build.BuildRetrievalError;
import com.android.tradefed.build.IFileDownloader;
import com.android.tradefed.log.LogUtil.CLog;
import com.google.api.client.googleapis.json.GoogleJsonResponseException;
import com.google.api.services.storage.Storage;
import com.google.api.services.storage.model.Objects;
import com.google.api.services.storage.model.StorageObject;
import com.google.common.annotations.VisibleForTesting;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
/** File downloader to download file from google cloud storage (GCS). */
public class GCSFileDownloader extends GCSCommon implements IFileDownloader {
public static final String GCS_PREFIX = "gs://";
public static final String GCS_APPROX_PREFIX = "gs:/";
private static final Pattern GCS_PATH_PATTERN = Pattern.compile("gs://([^/]*)/(.*)");
private static final String PATH_SEP = "/";
private static final Collection<String> SCOPES =
Collections.singleton("https://www.googleapis.com/auth/devstorage.read_only");
private static final long LIST_BATCH_SIZE = 100;
public GCSFileDownloader(File jsonKeyFile) {
super(jsonKeyFile);
}
public GCSFileDownloader() {}
/**
* Download a file from a GCS bucket file.
*
* @param bucketName GCS bucket name
* @param filename the filename
* @return {@link InputStream} with the file content.
*/
public InputStream downloadFile(String bucketName, String filename) throws IOException {
InputStream remoteInput = null;
ByteArrayOutputStream tmpStream = null;
try {
remoteInput =
getStorage().objects().get(bucketName, filename).executeMediaAsInputStream();
// The input stream from api call can not be reset. Change it to ByteArrayInputStream.
tmpStream = new ByteArrayOutputStream();
StreamUtil.copyStreams(remoteInput, tmpStream);
return new ByteArrayInputStream(tmpStream.toByteArray());
} finally {
StreamUtil.close(remoteInput);
StreamUtil.close(tmpStream);
}
}
private Storage getStorage() throws IOException {
return getStorage(SCOPES);
}
@VisibleForTesting
StorageObject getRemoteFileMetaData(String bucketName, String remoteFilename)
throws IOException {
try {
return getStorage().objects().get(bucketName, remoteFilename).execute();
} catch (GoogleJsonResponseException e) {
if (e.getStatusCode() == 404) {
return null;
}
throw e;
}
}
/**
* Download file from GCS.
*
* <p>Right now only support GCS path.
*
* @param remoteFilePath gs://bucket/file/path format GCS path.
* @return local file
* @throws BuildRetrievalError
*/
@Override
public File downloadFile(String remoteFilePath) throws BuildRetrievalError {
File destFile = createTempFile(remoteFilePath, null);
try {
downloadFile(remoteFilePath, destFile);
return destFile;
} catch (BuildRetrievalError e) {
FileUtil.recursiveDelete(destFile);
throw e;
}
}
@Override
public void downloadFile(String remotePath, File destFile) throws BuildRetrievalError {
String[] pathParts = parseGcsPath(remotePath);
downloadFile(pathParts[0], pathParts[1], destFile);
}
private boolean isFileFresh(File localFile, StorageObject remoteFile) throws IOException {
if (localFile == null && remoteFile == null) {
return true;
}
if (localFile == null || remoteFile == null) {
return false;
}
if (!localFile.exists()) {
return false;
}
return remoteFile.getMd5Hash().equals(FileUtil.calculateBase64Md5(localFile));
}
@Override
public boolean isFresh(File localFile, String remotePath) throws BuildRetrievalError {
String[] pathParts = parseGcsPath(remotePath);
String bucketName = pathParts[0];
String remoteFilename = pathParts[1];
try {
StorageObject remoteFileMeta = getRemoteFileMetaData(bucketName, remoteFilename);
if (localFile == null || !localFile.exists()) {
if (!isRemoteFolder(bucketName, remoteFilename) && remoteFileMeta == null) {
// The local doesn't exist and the remote filename is not a folder or a file.
return true;
}
return false;
}
if (!localFile.isDirectory()) {
return isFileFresh(localFile, remoteFileMeta);
}
remoteFilename = sanitizeDirectoryName(remoteFilename);
return recursiveCheckFolderFreshness(bucketName, remoteFilename, localFile);
} catch (IOException e) {
throw new BuildRetrievalError(e.getMessage(), e);
}
}
/**
* Check if remote folder is the same as local folder, recursively. The remoteFolderName must
* end with "/".
*
* @param bucketName is the gcs bucket name.
* @param remoteFolderName is the relative path to the bucket.
* @param localFolder is the local folder
* @return true if local file is the same as remote file, otherwise false.
* @throws IOException
*/
private boolean recursiveCheckFolderFreshness(
String bucketName, String remoteFolderName, File localFolder) throws IOException {
Set<String> subFilenames = new HashSet<>(Arrays.asList(localFolder.list()));
List<String> subRemoteFolders = new ArrayList<>();
List<StorageObject> subRemoteFiles = new ArrayList<>();
listRemoteFilesUnderFolder(bucketName, remoteFolderName, subRemoteFiles, subRemoteFolders);
for (StorageObject subRemoteFile : subRemoteFiles) {
String subFilename = Paths.get(subRemoteFile.getName()).getFileName().toString();
if (!isFileFresh(new File(localFolder, subFilename), subRemoteFile)) {
return false;
}
subFilenames.remove(subFilename);
}
for (String subRemoteFolder : subRemoteFolders) {
String subFolderName = Paths.get(subRemoteFolder).getFileName().toString();
if (!recursiveCheckFolderFreshness(
bucketName, subRemoteFolder, new File(localFolder, subFolderName))) {
return false;
}
subFilenames.remove(subFolderName);
}
return subFilenames.isEmpty();
}
void listRemoteFilesUnderFolder(
String bucketName, String folder, List<StorageObject> subFiles, List<String> subFolders)
throws IOException {
String pageToken = null;
while (true) {
com.google.api.services.storage.Storage.Objects.List listOperation =
getStorage()
.objects()
.list(bucketName)
.setPrefix(folder)
.setDelimiter(PATH_SEP)
.setMaxResults(LIST_BATCH_SIZE);
if (pageToken != null) {
listOperation.setPageToken(pageToken);
}
Objects objects = listOperation.execute();
if (objects.getItems() != null && !objects.getItems().isEmpty()) {
subFiles.addAll(objects.getItems());
}
if (objects.getPrefixes() != null && !objects.getPrefixes().isEmpty()) {
subFolders.addAll(objects.getPrefixes());
}
pageToken = objects.getNextPageToken();
if (pageToken == null) {
return;
}
}
}
String[] parseGcsPath(String remotePath) throws BuildRetrievalError {
if (remotePath.startsWith(GCS_APPROX_PREFIX) && !remotePath.startsWith(GCS_PREFIX)) {
// File object remove double // so we have to rebuild it in some cases
remotePath = remotePath.replaceAll(GCS_APPROX_PREFIX, GCS_PREFIX);
}
Matcher m = GCS_PATH_PATTERN.matcher(remotePath);
if (!m.find()) {
throw new BuildRetrievalError(
String.format("Only GCS path is supported, %s is not supported", remotePath));
}
return new String[] {m.group(1), m.group(2)};
}
String sanitizeDirectoryName(String name) {
/** Folder name should end with "/" */
if (!name.endsWith(PATH_SEP)) {
name += PATH_SEP;
}
return name;
}
/** check given filename is a folder or not. */
@VisibleForTesting
boolean isRemoteFolder(String bucketName, String filename) throws IOException {
filename = sanitizeDirectoryName(filename);
Objects objects =
getStorage()
.objects()
.list(bucketName)
.setPrefix(filename)
.setDelimiter(PATH_SEP)
.setMaxResults(1l)
.execute();
if (objects.getItems() != null && !objects.getItems().isEmpty()) {
return true;
}
if (objects.getPrefixes() != null && !objects.getPrefixes().isEmpty()) {
return true;
}
return false;
}
@VisibleForTesting
void downloadFile(String bucketName, String remoteFilename, File localFile)
throws BuildRetrievalError {
try {
if (!isRemoteFolder(bucketName, remoteFilename)) {
fetchRemoteFile(bucketName, remoteFilename, localFile);
return;
}
remoteFilename = sanitizeDirectoryName(remoteFilename);
recursiveDownloadFolder(bucketName, remoteFilename, localFile);
} catch (IOException e) {
CLog.e("Failed to download gs://%s/%s, clean up.", bucketName, remoteFilename);
throw new BuildRetrievalError(e.getMessage(), e);
}
}
private void fetchRemoteFile(String bucketName, String remoteFilename, File localFile)
throws IOException {
getStorage()
.objects()
.get(bucketName, remoteFilename)
.executeMediaAndDownloadTo(new FileOutputStream(localFile));
}
/**
* Recursively download remote folder to local folder.
*
* @param bucketName the gcs bucket name
* @param remoteFolderName remote folder name, must end with "/"
* @param localFolder local folder
* @throws IOException
*/
private void recursiveDownloadFolder(
String bucketName, String remoteFolderName, File localFolder) throws IOException {
CLog.d("Downloading folder gs://%s/%s.", bucketName, remoteFolderName);
if (!localFolder.exists()) {
FileUtil.mkdirsRWX(localFolder);
}
Set<String> subFilenames = new HashSet<>(Arrays.asList(localFolder.list()));
List<String> subRemoteFolders = new ArrayList<>();
List<StorageObject> subRemoteFiles = new ArrayList<>();
listRemoteFilesUnderFolder(bucketName, remoteFolderName, subRemoteFiles, subRemoteFolders);
for (StorageObject subRemoteFile : subRemoteFiles) {
String subFilename = Paths.get(subRemoteFile.getName()).getFileName().toString();
fetchRemoteFile(
bucketName, subRemoteFile.getName(), new File(localFolder, subFilename));
subFilenames.remove(subFilename);
}
for (String subRemoteFolder : subRemoteFolders) {
String subFolderName = Paths.get(subRemoteFolder).getFileName().toString();
recursiveDownloadFolder(
bucketName, subRemoteFolder, new File(localFolder, subFolderName));
subFilenames.remove(subFolderName);
}
for (String subFilename : subFilenames) {
FileUtil.recursiveDelete(new File(localFolder, subFilename));
}
}
/**
* Creates a unique file on temporary disk to house downloaded file with given path.
*
* <p>Constructs the file name based on base file name from path
*
* @param remoteFilePath the remote path to construct the name from
*/
@VisibleForTesting
File createTempFile(String remoteFilePath, File rootDir) throws BuildRetrievalError {
try {
// create a unique file.
File tmpFile = FileUtil.createTempFileForRemote(remoteFilePath, rootDir);
// now delete it so name is available
tmpFile.delete();
return tmpFile;
} catch (IOException e) {
String msg = String.format("Failed to create tmp file for %s", remoteFilePath);
throw new BuildRetrievalError(msg, e);
}
}
}