core/src/main/java/com/cloud/storage/template/SimpleHttpMultiFileDownloader.java

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.cloud.storage.template;

import static com.cloud.utils.NumbersUtil.toHumanReadableSize;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.RandomAccessFile;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.cloudstack.managed.context.ManagedContextRunnable;
import org.apache.cloudstack.storage.command.DownloadCommand;
import org.apache.commons.httpclient.Header;
import org.apache.commons.httpclient.HttpClient;
import org.apache.commons.httpclient.HttpException;
import org.apache.commons.httpclient.HttpMethod;
import org.apache.commons.httpclient.HttpMethodRetryHandler;
import org.apache.commons.httpclient.HttpStatus;
import org.apache.commons.httpclient.MultiThreadedHttpConnectionManager;
import org.apache.commons.httpclient.NoHttpResponseException;
import org.apache.commons.httpclient.methods.GetMethod;
import org.apache.commons.httpclient.methods.HeadMethod;
import org.apache.commons.httpclient.params.HttpMethodParams;
import org.apache.commons.lang3.StringUtils;
import org.apache.log4j.Logger;

import com.cloud.storage.StorageLayer;

public class SimpleHttpMultiFileDownloader extends ManagedContextRunnable implements TemplateDownloader {
    public static final Logger s_logger = Logger.getLogger(SimpleHttpMultiFileDownloader.class.getName());
    private static final MultiThreadedHttpConnectionManager s_httpClientManager = new MultiThreadedHttpConnectionManager();

    private static final int CHUNK_SIZE = 1024 * 1024; //1M
    private String[] downloadUrls;
    private String currentToFile;
    public TemplateDownloader.Status currentStatus;
    public TemplateDownloader.Status status;
    private String errorString = null;
    private long totalRemoteSize = 0;
    private long currentRemoteSize = 0;
    public long downloadTime = 0;
    public long currentTotalBytes;
    public long totalBytes = 0;
    private final HttpClient client;
    private GetMethod request;
    private boolean resume = false;
    private DownloadCompleteCallback completionCallback;
    StorageLayer _storage;
    boolean inited = true;

    private String toDir;
    private long maxTemplateSizeInBytes;
    private DownloadCommand.ResourceType resourceType = DownloadCommand.ResourceType.TEMPLATE;
    private final HttpMethodRetryHandler retryHandler;

    private HashMap<String, String> urlFileMap;
    private boolean followRedirects = false;

    public SimpleHttpMultiFileDownloader(StorageLayer storageLayer, String[] downloadUrls, String toDir,
                                         DownloadCompleteCallback callback, long maxTemplateSizeInBytes,
                                         DownloadCommand.ResourceType resourceType) {
        _storage = storageLayer;
        this.downloadUrls = downloadUrls;
        this.toDir = toDir;
        this.resourceType = resourceType;
        this.maxTemplateSizeInBytes = maxTemplateSizeInBytes;
        completionCallback = callback;
        status = TemplateDownloader.Status.NOT_STARTED;
        currentStatus = TemplateDownloader.Status.NOT_STARTED;
        currentTotalBytes = 0;
        client = new HttpClient(s_httpClientManager);
        retryHandler = createRetryTwiceHandler();
        urlFileMap = new HashMap<>();
    }

    private GetMethod createRequest(String downloadUrl) {
        GetMethod request = new GetMethod(downloadUrl);
        request.getParams().setParameter(HttpMethodParams.RETRY_HANDLER, retryHandler);
        request.setFollowRedirects(followRedirects);
        return request;
    }

    private void checkTemporaryDestination(String toDir) {
        try {
            File f = File.createTempFile("dnld", "tmp_", new File(toDir));
            if (_storage != null) {
                _storage.setWorldReadableAndWriteable(f);
            }
            currentToFile = f.getAbsolutePath();
        } catch (IOException ex) {
            errorString = "Unable to start download -- check url? ";
            currentStatus = TemplateDownloader.Status.UNRECOVERABLE_ERROR;
            s_logger.warn("Exception in constructor -- " + ex.toString());
        }
    }

    private HttpMethodRetryHandler createRetryTwiceHandler() {
        return new HttpMethodRetryHandler() {
            @Override
            public boolean retryMethod(final HttpMethod method, final IOException exception, int executionCount) {
                if (executionCount >= 2) {
                    // Do not retry if over max retry count
                    return false;
                }
                if (exception instanceof NoHttpResponseException) {
                    // Retry if the server dropped connection on us
                    return true;
                }
                if (!method.isRequestSent()) {
                    // Retry if the request has not been sent fully or
                    // if it's OK to retry methods that have been sent
                    return true;
                }
                // otherwise do not retry
                return false;
            }
        };
    }

    private void tryAndGetTotalRemoteSize() {
        for (String downloadUrl : downloadUrls) {
            if (StringUtils.isBlank(downloadUrl)) {
                continue;
            }
            HeadMethod headMethod = new HeadMethod(downloadUrl);
            try {
                if (client.executeMethod(headMethod) != HttpStatus.SC_OK) {
                    continue;
                }
                Header contentLengthHeader = headMethod.getResponseHeader("content-length");
                if (contentLengthHeader == null) {
                    continue;
                }
                totalRemoteSize += Long.parseLong(contentLengthHeader.getValue());
            } catch (IOException e) {
                s_logger.warn(String.format("Cannot reach URL: %s while trying to get remote sizes due to: %s", downloadUrl, e.getMessage()), e);
            } finally {
                headMethod.releaseConnection();
            }
        }
    }

    private long downloadFile(String downloadUrl) {
        s_logger.debug("Starting download for " + downloadUrl);
        currentTotalBytes = 0;
        currentRemoteSize = 0;
        File file = null;
        request = null;
        try {
            request = createRequest(downloadUrl);
            checkTemporaryDestination(toDir);
            urlFileMap.put(downloadUrl, currentToFile);
            file = new File(currentToFile);
            long localFileSize = checkLocalFileSizeForResume(resume, file);
            if (checkServerResponse(localFileSize, downloadUrl)) return 0;
            if (!tryAndGetRemoteSize()) return 0;
            if (!canHandleDownloadSize()) return 0;
            checkAndSetDownloadSize();
            try (InputStream in = request.getResponseBodyAsStream();
                 RandomAccessFile out = new RandomAccessFile(file, "rw");
            ) {
                out.seek(localFileSize);
                s_logger.info("Starting download from " + downloadUrl + " to " + currentToFile + " remoteSize=" + toHumanReadableSize(currentRemoteSize) + " , max size=" + toHumanReadableSize(maxTemplateSizeInBytes));
                if (copyBytes(file, in, out)) return 0;
                checkDownloadCompletion();
            }
            return currentTotalBytes;
        } catch (HttpException hte) {
            currentStatus = TemplateDownloader.Status.UNRECOVERABLE_ERROR;
            errorString = hte.getMessage();
        } catch (IOException ioe) {
            currentStatus = TemplateDownloader.Status.UNRECOVERABLE_ERROR; //probably a file write error?
            // Let's not overwrite the original error message.
            if (errorString == null) {
                errorString = ioe.getMessage();
            }
        } finally {
            if (currentStatus == Status.UNRECOVERABLE_ERROR && file != null && file.exists() && !file.isDirectory()) {
                file.delete();
            }
            if (request != null) {
                request.releaseConnection();
            }
        }
        return 0;
    }

    @Override
    public long download(boolean resume, DownloadCompleteCallback callback) {
        if (skipDownloadOnStatus()) return 0;
        if (resume) {
            s_logger.error("Resume not allowed for this downloader");
            status = Status.UNRECOVERABLE_ERROR;
            return 0;
        }
        s_logger.debug("Starting downloads");
        status = Status.IN_PROGRESS;
        Date start = new Date();
        tryAndGetTotalRemoteSize();
        for (String downloadUrl : downloadUrls) {
            if (StringUtils.isBlank(downloadUrl)) {
                continue;
            }
            long bytes = downloadFile(downloadUrl);
            if (currentStatus != Status.DOWNLOAD_FINISHED) {
                break;
            }
            totalBytes += bytes;
        }
        status = currentStatus;
        Date finish = new Date();
        downloadTime += finish.getTime() - start.getTime();
        if (callback != null) {
            callback.downloadComplete(status);
        }
        return 0;
    }

    private boolean copyBytes(File file, InputStream in, RandomAccessFile out) throws IOException {
        int bytes;
        byte[] block = new byte[CHUNK_SIZE];
        long offset = 0;
        boolean done = false;
        currentStatus = Status.IN_PROGRESS;
        while (!done && currentStatus != Status.ABORTED && offset <= currentRemoteSize) {
            if ((bytes = in.read(block, 0, CHUNK_SIZE)) > -1) {
                offset = writeBlock(bytes, out, block, offset);
            } else {
                done = true;
            }
        }
        out.getFD().sync();
        return false;
    }

    private long writeBlock(int bytes, RandomAccessFile out, byte[] block, long offset) throws IOException {
        out.write(block, 0, bytes);
        offset += bytes;
        out.seek(offset);
        currentTotalBytes += bytes;
        return offset;
    }

    private void checkDownloadCompletion() {
        String downloaded = "(incomplete download)";
        if (currentTotalBytes >= currentRemoteSize) {
            currentStatus = Status.DOWNLOAD_FINISHED;
            downloaded = "(download complete remote=" + toHumanReadableSize(currentRemoteSize) + " bytes)";
        }
        errorString = "Downloaded " + toHumanReadableSize(currentTotalBytes) + " bytes " + downloaded;
    }

    private boolean canHandleDownloadSize() {
        if (currentRemoteSize > maxTemplateSizeInBytes) {
            s_logger.info("Remote size is too large: " + toHumanReadableSize(currentRemoteSize) + " , max=" + toHumanReadableSize(maxTemplateSizeInBytes));
            currentStatus = Status.UNRECOVERABLE_ERROR;
            errorString = "Download file size is too large";
            return false;
        }
        return true;
    }

    private void checkAndSetDownloadSize() {
        if (currentRemoteSize == 0) {
            currentRemoteSize = maxTemplateSizeInBytes;
        }
        if (totalRemoteSize == 0) {
            totalRemoteSize = currentRemoteSize;
        }
    }

    private boolean tryAndGetRemoteSize() {
        Header contentLengthHeader = request.getResponseHeader("content-length");
        boolean chunked = false;
        long reportedRemoteSize = 0;
        if (contentLengthHeader == null) {
            Header chunkedHeader = request.getResponseHeader("Transfer-Encoding");
            if (chunkedHeader == null || !"chunked".equalsIgnoreCase(chunkedHeader.getValue())) {
                currentStatus = Status.UNRECOVERABLE_ERROR;
                errorString = " Failed to receive length of download ";
                return false;
            } else if ("chunked".equalsIgnoreCase(chunkedHeader.getValue())) {
                chunked = true;
            }
        } else {
            reportedRemoteSize = Long.parseLong(contentLengthHeader.getValue());
            if (reportedRemoteSize == 0) {
                currentStatus = Status.DOWNLOAD_FINISHED;
                String downloaded = "(download complete remote=" + currentRemoteSize + "bytes)";
                errorString = "Downloaded " + currentTotalBytes + " bytes " + downloaded;
                downloadTime = 0;
                return false;
            }
        }

        if (currentRemoteSize == 0) {
            currentRemoteSize = reportedRemoteSize;
        }
        return true;
    }

    private boolean checkServerResponse(long localFileSize, String downloadUrl) throws IOException {
        int responseCode = 0;

        if (localFileSize > 0) {
            // require partial content support for resume
            request.addRequestHeader("Range", "bytes=" + localFileSize + "-");
            if (client.executeMethod(request) != HttpStatus.SC_PARTIAL_CONTENT) {
                errorString = "HTTP Server does not support partial get";
                currentStatus = Status.UNRECOVERABLE_ERROR;
                return true;
            }
        } else if ((responseCode = client.executeMethod(request)) != HttpStatus.SC_OK) {
            currentStatus = Status.UNRECOVERABLE_ERROR;
            errorString = " HTTP Server returned " + responseCode + " (expected 200 OK) ";
            if (List.of(HttpStatus.SC_MOVED_PERMANENTLY, HttpStatus.SC_MOVED_TEMPORARILY).contains(responseCode)
                    && !followRedirects) {
                errorString = String.format("Failed to download %s due to redirection, response code: %d",
                        downloadUrl, responseCode);
                s_logger.error(errorString);
            }
            return true; //FIXME: retry?
        }
        return false;
    }

    private long checkLocalFileSizeForResume(boolean resume, File file) {
        // TODO check the status of this downloader as well?
        long localFileSize = 0;
        if (file.exists() && resume) {
            localFileSize = file.length();
            s_logger.info("Resuming download to file (current size)=" + toHumanReadableSize(localFileSize));
        }
        return localFileSize;
    }

    private boolean skipDownloadOnStatus() {
        switch (currentStatus) {
            case ABORTED:
            case UNRECOVERABLE_ERROR:
            case DOWNLOAD_FINISHED:
                return true;
            default:

        }
        return false;
    }

    public String[] getDownloadUrls() {
        return downloadUrls;
    }

    public String getCurrentToFile() {
        File file = new File(currentToFile);

        return file.getAbsolutePath();
    }

    @Override
    public TemplateDownloader.Status getStatus() {
        return currentStatus;
    }

    @Override
    public long getDownloadTime() {
        return downloadTime;
    }

    @Override
    public long getDownloadedBytes() {
        return totalBytes;
    }

    @Override
    @SuppressWarnings("fallthrough")
    public boolean stopDownload() {
        switch (getStatus()) {
            case IN_PROGRESS:
                if (request != null) {
                    request.abort();
                }
                currentStatus = TemplateDownloader.Status.ABORTED;
                return true;
            case UNKNOWN:
            case NOT_STARTED:
            case RECOVERABLE_ERROR:
            case UNRECOVERABLE_ERROR:
            case ABORTED:
                currentStatus = TemplateDownloader.Status.ABORTED;
            case DOWNLOAD_FINISHED:
                File f = new File(currentToFile);
                if (f.exists()) {
                    f.delete();
                }
                return true;

            default:
                return true;
        }
    }

    @Override
    public int getDownloadPercent() {
        if (totalRemoteSize == 0) {
            return 0;
        }

        return (int)(100.0 * totalBytes / totalRemoteSize);
    }

    @Override
    protected void runInContext() {
        try {
            download(resume, completionCallback);
        } catch (Throwable t) {
            s_logger.warn("Caught exception during download " + t.getMessage(), t);
            errorString = "Failed to install: " + t.getMessage();
            currentStatus = TemplateDownloader.Status.UNRECOVERABLE_ERROR;
        }

    }

    @Override
    public void setStatus(TemplateDownloader.Status status) {
        this.currentStatus = status;
    }

    public boolean isResume() {
        return resume;
    }

    @Override
    public String getDownloadError() {
        return errorString == null ? " " : errorString;
    }

    @Override
    public String getDownloadLocalPath() {
        return toDir;
    }

    @Override
    public void setResume(boolean resume) {
        this.resume = resume;
    }

    @Override
    public long getMaxTemplateSizeInBytes() {
        return maxTemplateSizeInBytes;
    }

    @Override
    public void setDownloadError(String error) {
        errorString = error;
    }

    @Override
    public boolean isInited() {
        return inited;
    }

    public DownloadCommand.ResourceType getResourceType() {
        return resourceType;
    }

    public Map<String, String> getDownloadedFilesMap() {
        return urlFileMap;
    }

    @Override
    public void setFollowRedirects(boolean followRedirects) {
        this.followRedirects = followRedirects;
        if (this.request != null) {
            this.request.setFollowRedirects(followRedirects);
        }
    }
}