ArchiveService.java

/*
 * This file is part of vospace-file-service
 * Copyright (C) 2021 Istituto Nazionale di Astrofisica
 * SPDX-License-Identifier: GPL-3.0-or-later
 */
package it.inaf.ia2.transfer.service;

import it.inaf.ia2.aa.ServletRapClient;
import it.inaf.ia2.aa.data.User;
import it.inaf.ia2.rap.client.call.TokenExchangeRequest;
import it.inaf.ia2.transfer.auth.TokenPrincipal;
import it.inaf.ia2.transfer.persistence.FileDAO;
import it.inaf.ia2.transfer.persistence.JobDAO;
import it.inaf.ia2.transfer.persistence.LocationDAO;
import it.inaf.ia2.transfer.persistence.model.FileInfo;
import it.inaf.oats.vospace.exception.InternalFaultException;
import it.inaf.oats.vospace.exception.PermissionDeniedException;
import it.inaf.oats.vospace.exception.QuotaExceededException;
import it.inaf.oats.vospace.parent.persistence.LinkedServiceDAO;
import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.UncheckedIOException;
import java.nio.file.Files;
import java.security.Principal;
import java.util.List;
import java.util.Map;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;
import javax.annotation.PostConstruct;
import javax.servlet.http.HttpServletRequest;
import net.ivoa.xml.uws.v1.ExecutionPhase;
import org.kamranzafar.jtar.TarEntry;
import org.kamranzafar.jtar.TarOutputStream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.stereotype.Service;
import org.springframework.util.FileSystemUtils;
import org.springframework.util.unit.DataSize;
import org.springframework.web.client.RestTemplate;

@Service
public class ArchiveService {

    private static final Logger LOG = LoggerFactory.getLogger(ArchiveService.class);

    @Autowired
    private FileDAO fileDAO;

    @Autowired
    private LocationDAO locationDAO;
    
    @Autowired
    private LinkedServiceDAO linkedServiceDAO;

    @Autowired
    private JobDAO jobDAO;

    @Autowired
    private AuthorizationService authorizationService;

    @Autowired
    private RestTemplate restTemplate;
    
    @Autowired
    private ServletRapClient rapClient;

    @Value("${upload_location_id}")
    private int uploadLocationId;

    // Directory containing temporary files generated by jobs.
    @Value("${generated.dir}")
    private String generatedDirString;
    private File generatedDir;

    // Maximum size of the working directory for each registered user
    @Value("${generated.dir.max-size}")
    private DataSize generatedDirMaxSize;

    @PostConstruct
    public void init() {
        this.generatedDir = new File(generatedDirString);
        if (!this.generatedDir.exists()) {
            if (!this.generatedDir.mkdirs()) {
                throw new IllegalStateException("Unable to create directory " + this.generatedDir.getAbsolutePath());
            }
        }
    }

    public <O extends OutputStream, E> void createArchive(ArchiveJob job, HttpServletRequest servletRequest) {

        jobDAO.updateJobPhase(ExecutionPhase.EXECUTING, job.getJobId());

        LOG.trace("Started archive job " + job.getJobId());

        try {
            // TODO: check total size limit
            File archiveFile = getArchiveFile(job);

            String commonParent = getCommonParent(job.getVosPaths());
            // support directory used to generate folder inside tar files (path is redefined each time by TarEntry class)
            File supportDir = Files.createTempDirectory("dir").toFile();

            // it will be initialized only when necessary
            Map<Integer, String> portalLocationUrls = null;

            try (ArchiveHandler<O, E> handler = getArchiveHandler(archiveFile, job.getType())) {

                for (FileInfo fileInfo : fileDAO.getArchiveFileInfos(job.getVosPaths())) {

                    String relPath = fileInfo.getVirtualPath().substring(commonParent.length());

                    if (fileInfo.isDirectory()) {
                        handler.putNextEntry(supportDir, relPath);
                        continue;
                    }

                    // I expect only external links
                    // local links have been resolved before calling this endpoint                    
                    if (fileInfo.isLink()) {
                        downloadExternalLinkIntoArchive(fileInfo, relPath, 
                                job.getPrincipal(), handler, servletRequest);
                        continue;
                    }

                    if (fileInfo.getLocationId() != null && "portal".equals(fileInfo.getLocationType())) {
                        // remote file
                        if (portalLocationUrls == null) {
                            portalLocationUrls = locationDAO.getPortalLocationUrls();
                        }
                        String url = portalLocationUrls.get(fileInfo.getLocationId());
                        downloadRemoteLocationFileIntoArchive(fileInfo, relPath, job.getPrincipal(), handler, url);
                    } else {
                        // local file or virtual directory
                        writeFileIntoArchive(fileInfo, relPath, job.getPrincipal(), handler);
                    }
                }
            } finally {
                FileSystemUtils.deleteRecursively(supportDir);
            }

        } catch (IOException ex) {
            throw new UncheckedIOException(ex);
        }
    }

    private File getArchiveFile(ArchiveJob job) throws IOException {

        File parentDir = getArchiveParentDir(job.getPrincipal());

        if (!parentDir.exists()) {
            if (!parentDir.mkdirs()) {
                LOG.error("Unable to create directory " + parentDir.getAbsolutePath());
                throw new InternalFaultException("Unable to create temporary directory for job");
            }
        }

        checkQuotaLimit(parentDir);

        File archiveFile = parentDir.toPath().resolve(job.getJobId() + "." + job.getType().getExtension()).toFile();

        if (!archiveFile.createNewFile()) {
            LOG.error("Unable to create file " + archiveFile.getAbsolutePath());
            throw new InternalFaultException("Unable to create archive file");
        }

        return archiveFile;
    }

    /**
     * If used working space exceeds quota limit throws an
     * InsufficientStorageException.
     */
    private void checkQuotaLimit(File parentDir) throws IOException {
        long usedSpace = Files.walk(parentDir.toPath()).mapToLong(p -> p.toFile().length()).sum();

        if (usedSpace > generatedDirMaxSize.toBytes()) {
            throw new QuotaExceededException("Archive size limit exceeded.");
        }
    }

    public File getArchiveParentDir(Principal principal) {
        return generatedDir.toPath().resolve(principal.getName()).toFile();
    }

    private String getCommonParent(List<String> vosPaths) {
        String commonParent = null;
        for (String vosPath : vosPaths) {
            if (commonParent == null) {
                commonParent = vosPath;
            } else {
                StringBuilder newCommonParent = new StringBuilder();
                boolean same = true;
                int lastSlashPos = vosPath.lastIndexOf("/");
                for (int i = 0; same && i < Math.min(commonParent.length(), vosPath.length()) && i <= lastSlashPos; i++) {
                    if (commonParent.charAt(i) == vosPath.charAt(i)) {
                        newCommonParent.append(commonParent.charAt(i));
                    } else {
                        same = false;
                    }
                }
                commonParent = newCommonParent.toString();
            }
        }
        return commonParent;
    }

    private abstract class ArchiveHandler<O extends OutputStream, E> implements AutoCloseable {

        private final O os;
        private final File parentDir;

        ArchiveHandler(O os, File parentDir) {
            this.os = os;
            this.parentDir = parentDir;
        }

        public abstract E getEntry(File file, String path);

        protected abstract void putNextEntry(E entry) throws IOException;

        public final void putNextEntry(File file, String path) throws IOException {
            putNextEntry(getEntry(file, path));
            checkQuotaLimit(parentDir);
        }

        public final O getOutputStream() {
            return os;
        }

        @Override
        public void close() throws IOException {
            os.close();
        }
    }

    private class TarArchiveHandler extends ArchiveHandler<TarOutputStream, TarEntry> {

        TarArchiveHandler(File archiveFile) throws IOException {
            super(new TarOutputStream(new BufferedOutputStream(new FileOutputStream(archiveFile))), archiveFile.getParentFile());
        }

        @Override
        public TarEntry getEntry(File file, String path) {
            return new TarEntry(file, path);
        }

        @Override
        public void putNextEntry(TarEntry entry) throws IOException {
            getOutputStream().putNextEntry(entry);
        }
    }

    private class ZipArchiveHandler extends ArchiveHandler<ZipOutputStream, ZipEntry> {

        ZipArchiveHandler(File archiveFile) throws IOException {
            super(new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(archiveFile))), archiveFile.getParentFile());
        }

        @Override
        public ZipEntry getEntry(File file, String path) {
            if (file.isDirectory()) {
                // ZipEntry assumes that paths ending with / are folders
                path += "/";
            }
            return new ZipEntry(path);
        }

        @Override
        public void putNextEntry(ZipEntry entry) throws IOException {
            getOutputStream().putNextEntry(entry);
        }
    }

    private ArchiveHandler getArchiveHandler(File archiveFile, ArchiveJob.Type type) throws IOException {
        switch (type) {
            case TAR:
                return new TarArchiveHandler(archiveFile);
            case ZIP:
                return new ZipArchiveHandler(archiveFile);
            default:
                throw new UnsupportedOperationException("Type " + type + " not supported yet");
        }
    }

    private <O extends OutputStream, E> void downloadFromUrlIntoArchive(String url, String relPath, TokenPrincipal tokenPrincipal, ArchiveHandler<O, E> handler) {
        LOG.trace("Downloading file from " + url);

        restTemplate.execute(url, HttpMethod.GET, req -> {
            HttpHeaders headers = req.getHeaders();
            if (tokenPrincipal.getToken() != null) {
                headers.setBearerAuth(tokenPrincipal.getToken());
            }
        }, res -> {
            File tmpFile = Files.createTempFile("download", null).toFile();
            try (FileOutputStream os = new FileOutputStream(tmpFile)) {
                res.getBody().transferTo(os);
                handler.putNextEntry(tmpFile, relPath);
                try (FileInputStream is = new FileInputStream(tmpFile)) {
                    is.transferTo(handler.getOutputStream());
                }
            } finally {
                tmpFile.delete();
            }
            return null;
        }, new Object[]{});
    }

    private <O extends OutputStream, E> void downloadRemoteLocationFileIntoArchive(
            FileInfo fileInfo, String relPath, TokenPrincipal tokenPrincipal,
            ArchiveHandler<O, E> handler, String baseUrl) {

        if (baseUrl == null) {
            LOG.error("Location URL not found for location " + fileInfo.getLocationId());
            throw new InternalFaultException("Unable to retrieve location of file "
                    + fileInfo.getVirtualPath());
        }

        String url = baseUrl + "/" + fileInfo.getVirtualName();

        downloadFromUrlIntoArchive(url, relPath, tokenPrincipal, handler);
    }

    private <O extends OutputStream, E> void downloadExternalLinkIntoArchive(
            FileInfo fileInfo, String relPath, TokenPrincipal tokenPrincipal,
            ArchiveHandler<O, E> handler, HttpServletRequest servletRequest) {

        String url = fileInfo.getTarget();

        if (url == null || url.isBlank()) {
            LOG.error("Target URL of link at path: {} is null or blank", fileInfo.getVirtualPath());
            throw new InternalFaultException("Target URL of link at path: "
                    + fileInfo.getVirtualPath() + " is null or blank");
        }
        
        // Append token if url is recognized
        
        if (linkedServiceDAO.isLinkedServiceUrl(url)) {
                url += "?token=" + getEndpointToken(tokenPrincipal, url, servletRequest);
            }
        
        downloadFromUrlIntoArchive(url, relPath, tokenPrincipal, handler);
        
    }

    private <O extends OutputStream, E> void writeFileIntoArchive(FileInfo fileInfo, String relPath, TokenPrincipal tokenPrincipal, ArchiveHandler<O, E> handler) throws IOException {
        if (!authorizationService.isDownloadable(fileInfo, tokenPrincipal)) {
            throw PermissionDeniedException.forPath(fileInfo.getVirtualPath());
        }

        File file = new File(fileInfo.getFilePath());
        LOG.trace("Adding file " + file.getAbsolutePath() + " to tar archive");

        try (InputStream is = new FileInputStream(file)) {
            handler.putNextEntry(file, relPath);
            is.transferTo(handler.getOutputStream());
        }
    }
    
    private String getEndpointToken(TokenPrincipal tokenPrincipal, 
            String endpoint, HttpServletRequest servletRequest) {

        String token = tokenPrincipal.getToken();

        if (token == null) {
            throw new PermissionDeniedException("Token is null");
        }

        TokenExchangeRequest exchangeRequest = new TokenExchangeRequest()
                .setSubjectToken(token)
                .setResource(endpoint);

        // TODO: add audience and scope
        return rapClient.exchangeToken(exchangeRequest, servletRequest);
    }
    
}