android-10.0.0_r47/s

/*
 * 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);
        }
    }
}