oss/src/main/java/cn/reghao/dfs/store/service/FileUploadService.java

package cn.reghao.dfs.store.service;

import cn.reghao.dfs.store.model.dto.*;
import cn.reghao.dfs.store.util.redis.ds.RedisSet;
import cn.reghao.jutil.jdk.security.DigestUtil;
import cn.reghao.dfs.store.config.PathUrl;
import cn.reghao.dfs.store.db.repository.FileRepository;
import cn.reghao.dfs.store.model.dto.*;
import cn.reghao.dfs.store.model.po.FileInfo;
import cn.reghao.dfs.store.util.media.ImageOps;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.multipart.MultipartFile;

import java.io.*;

/**
 * @author reghao
 * @date 2021-08-04 17:29:47
 */
@Slf4j
@Service
public class FileUploadService {
    // 10MiB
    private final static long PART_SIZE = 1024*1024*10;

    private final FileUrlService fileUrlService;
    private final FileStoreService fileStoreService;
    private final FileTypeService fileTypeService;
    private final FileRepository fileRepository;
    private final RedisSet redisSet;

    public FileUploadService(FileUrlService fileUrlService, FileStoreService fileStoreService,
                             FileTypeService fileTypeService, FileRepository fileRepository,
                             RedisSet redisSet) {
        this.fileUrlService = fileUrlService;
        this.fileStoreService = fileStoreService;
        this.fileTypeService = fileTypeService;
        this.fileRepository = fileRepository;
        this.redisSet = redisSet;
    }

    public synchronized UploadPrepareRet prepareUpload(UploadPrepare uploadPrepare) {
        // TODO 避免用户通过 sha256sum 查询文件是否存在, 这是一个不安全的行为, 用户可能会伪造 sha256sum 来获取不属于他权限的文件
        // TODO 解决办法是检查用户已上传文件的 sha256sum, 这样可确保他拥有 sha256sum 所属文件的权限, 但这会带来服务器带宽和磁盘 IO 的性能开销
        String sha256sum = uploadPrepare.getFileSha256sum();
        FileInfo fileInfo = fileRepository.getFileInfo(sha256sum);
        if (fileInfo == null) {
            String uploadId = fileRepository.createFileUploadId(uploadPrepare);
            return new UploadPrepareRet(uploadId, PART_SIZE, false);
        }

        String uploadId = fileRepository.getOrCreateUploadId(fileInfo.getFileId());
        // uploaded 为 true 则 exist 也应该为 true, 但这里统一将 exist 设置为 false, 表示客户端总是要传输文件流
        boolean uploaded = fileInfo.getUploaded();
        return new UploadPrepareRet(uploadId, PART_SIZE, false);
    }

    /**
     * 将文件存放到 FileStore，并在数据库中记录文件相关信息
     *
     * @param
     * @return
     * @date 2021-12-08 下午3:30
     */
    public void putFile(UploadFile uploadFile) throws Exception {
        String uploadId = uploadFile.getUploadId();
        FileInfo fileInfo = fileRepository.getFileInfoByUploadId(uploadId);
        if (!fileInfo.getUploaded()) {
            put(uploadFile.getFile(), fileInfo);
        }
    }

    @Transactional(rollbackFor = Exception.class)
    public synchronized String put(MultipartFile multipartFile, FileInfo fileInfo) throws Exception {
        byte[] bytes = multipartFile.getBytes();
        String sha256sum = DigestUtil.sha256sum(bytes);
        if (!sha256sum.equals(fileInfo.getSha256sum())) {
            throw new Exception("uploadId 和 sha256sum 不匹配!");
        }

        String contentType = multipartFile.getContentType();
        String fileId = fileInfo.getFileId();
        String suffix = fileInfo.getSuffix();
        PathUrl pathUrl;
        if (contentType != null && contentType.startsWith("image")) {
            ImageOps.Size size1 = ImageOps.info(new ByteArrayInputStream(bytes));
            pathUrl = fileUrlService.getImagePathAndUrl(sha256sum, fileId, suffix, size1.getWidth(), size1.getHeight());
        } else {
            pathUrl = fileUrlService.genPathAndUrl(sha256sum, fileId, suffix);
        }

        FileContentType fileType = fileTypeService.getFileType(contentType, pathUrl.getFilePath());
        fileStoreService.saveFile(pathUrl.getFilePath(), bytes);
        fileRepository.saveFile(fileId, fileType, pathUrl);
        return pathUrl.getUrl();
    }

    /**
     * 处理通过 HTTP 请求提交的分片文件
     *
     * @param
     * @return
     * @date 2021-12-08 下午3:31
     */
    @Transactional(rollbackFor = Exception.class)
    public synchronized UploadFilePartRet putFilePart(UploadFilePart uploadFilePart) throws Exception {
        int partIndex = uploadFilePart.getSplitIndex();
        int partNum = uploadFilePart.getSplitNum();
        MultipartFile multipartFile = uploadFilePart.getFile();
        String uploadId = uploadFilePart.getUploadId();
        FileInfo fileInfo = fileRepository.getFileInfoByFileId(uploadId);
        // 没有检查文件的 sha256sum, 存在安全隐患
        if (fileInfo.getUploaded()) {
            return new UploadFilePartRet(uploadId, true);
        }

        String fileId = fileInfo.getFileId();
        String sha256sum = fileInfo.getSha256sum();
        String suffix = fileInfo.getSuffix();
        PathUrl pathUrl = fileUrlService.genPathAndUrl(sha256sum, fileId, suffix);
        String key = RedisKey.filePartKey(fileId);
        if (!redisSet.sismember(key, partIndex)) {
            redisSet.sadd(key, partIndex);

            long size = fileInfo.getSize();
            File file = fileStoreService.createFile(pathUrl.getFilePath(), size);
            long pos = partIndex * PART_SIZE;
            fileStoreService.writeToFile(multipartFile.getInputStream(), file, pos);
        }

        long len = redisSet.scard(key);
        if (len != partNum) {
            return new UploadFilePartRet(fileId, false);
        } else {
            String path = pathUrl.getFilePath();
            FileInputStream fis = new FileInputStream(path);
            String sha256sumMerged = DigestUtil.sha256sum(fis);
            if (!sha256sum.equals(sha256sumMerged)) {
                throw new Exception("uploadId 和 sha256sum 不匹配!");
            }

            FileContentType fileType = fileTypeService.getFileType(fileInfo.getContentType(), pathUrl.getFilePath());
            fileRepository.saveFile(fileId, fileType, pathUrl);
            redisSet.spop(key, len);
            return new UploadFilePartRet(uploadId, true);
        }
    }
}