cloudreve/pkg/filesystem/archive.go

package filesystem

import (
	"archive/zip"
	"context"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"

	model "github.com/cloudreve/Cloudreve/v3/models"
	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/fsctx"
	"github.com/cloudreve/Cloudreve/v3/pkg/util"
	"github.com/gin-gonic/gin"
	"github.com/mholt/archiver/v4"
)

/* ===============
     压缩/解压缩
   ===============
*/

// Compress 创建给定目录和文件的压缩文件
func (fs *FileSystem) Compress(ctx context.Context, writer io.Writer, folderIDs, fileIDs []uint, isArchive bool) error {
	// 查找待压缩目录
	folders, err := model.GetFoldersByIDs(folderIDs, fs.User.ID)
	if err != nil && len(folderIDs) != 0 {
		return ErrDBListObjects
	}

	// 查找待压缩文件
	files, err := model.GetFilesByIDs(fileIDs, fs.User.ID)
	if err != nil && len(fileIDs) != 0 {
		return ErrDBListObjects
	}

	// 如果上下文限制了父目录，则进行检查
	if parent, ok := ctx.Value(fsctx.LimitParentCtx).(*model.Folder); ok {
		// 检查目录
		for _, folder := range folders {
			if *folder.ParentID != parent.ID {
				return ErrObjectNotExist
			}
		}

		// 检查文件
		for _, file := range files {
			if file.FolderID != parent.ID {
				return ErrObjectNotExist
			}
		}
	}

	// 尝试获取请求上下文，以便于后续检查用户取消任务
	reqContext := ctx
	ginCtx, ok := ctx.Value(fsctx.GinCtx).(*gin.Context)
	if ok {
		reqContext = ginCtx.Request.Context()
	}

	// 将顶级待处理对象的路径设为根路径
	for i := 0; i < len(folders); i++ {
		folders[i].Position = ""
	}
	for i := 0; i < len(files); i++ {
		files[i].Position = ""
	}

	// 创建压缩文件Writer
	zipWriter := zip.NewWriter(writer)
	defer zipWriter.Close()

	ctx = reqContext

	// 压缩各个目录及文件
	for i := 0; i < len(folders); i++ {
		select {
		case <-reqContext.Done():
			// 取消压缩请求
			return ErrClientCanceled
		default:
			fs.doCompress(reqContext, nil, &folders[i], zipWriter, isArchive)
		}

	}
	for i := 0; i < len(files); i++ {
		select {
		case <-reqContext.Done():
			// 取消压缩请求
			return ErrClientCanceled
		default:
			fs.doCompress(reqContext, &files[i], nil, zipWriter, isArchive)
		}
	}

	return nil
}

func (fs *FileSystem) doCompress(ctx context.Context, file *model.File, folder *model.Folder, zipWriter *zip.Writer, isArchive bool) {
	// 如果对象是文件
	if file != nil {
		// 切换上传策略
		fs.Policy = file.GetPolicy()
		err := fs.DispatchHandler()
		if err != nil {
			util.Log().Warning("无法压缩文件%s，%s", file.Name, err)
			return
		}

		// 获取文件内容
		fileToZip, err := fs.Handler.Get(
			context.WithValue(ctx, fsctx.FileModelCtx, *file),
			file.SourceName,
		)
		if err != nil {
			util.Log().Debug("Open%s，%s", file.Name, err)
			return
		}
		if closer, ok := fileToZip.(io.Closer); ok {
			defer closer.Close()
		}

		// 创建压缩文件头
		header := &zip.FileHeader{
			Name:               filepath.FromSlash(path.Join(file.Position, file.Name)),
			Modified:           file.UpdatedAt,
			UncompressedSize64: file.Size,
		}

		// 指定是压缩还是归档
		if isArchive {
			header.Method = zip.Store
		} else {
			header.Method = zip.Deflate
		}

		writer, err := zipWriter.CreateHeader(header)
		if err != nil {
			return
		}

		_, err = io.Copy(writer, fileToZip)
	} else if folder != nil {
		// 对象是目录
		// 获取子文件
		subFiles, err := folder.GetChildFiles()
		if err == nil && len(subFiles) > 0 {
			for i := 0; i < len(subFiles); i++ {
				fs.doCompress(ctx, &subFiles[i], nil, zipWriter, isArchive)
			}

		}
		// 获取子目录，继续递归遍历
		subFolders, err := folder.GetChildFolder()
		if err == nil && len(subFolders) > 0 {
			for i := 0; i < len(subFolders); i++ {
				fs.doCompress(ctx, nil, &subFolders[i], zipWriter, isArchive)
			}
		}
	}
}

// Decompress 解压缩给定压缩文件到dst目录
func (fs *FileSystem) Decompress(ctx context.Context, src, dst, encoding string) error {
	err := fs.ResetFileIfNotExist(ctx, src)
	if err != nil {
		return err
	}

	tempZipFilePath := ""
	defer func() {
		// 结束时删除临时压缩文件
		if tempZipFilePath != "" {
			if err := os.Remove(tempZipFilePath); err != nil {
				util.Log().Warning("无法删除临时压缩文件 %s , %s", tempZipFilePath, err)
			}
		}
	}()

	// 下载压缩文件到临时目录
	fileStream, err := fs.Handler.Get(ctx, fs.FileTarget[0].SourceName)
	if err != nil {
		return err
	}

	defer fileStream.Close()

	tempZipFilePath = filepath.Join(
		util.RelativePath(model.GetSettingByName("temp_path")),
		"decompress",
		fmt.Sprintf("archive_%d.zip", time.Now().UnixNano()),
	)

	zipFile, err := util.CreatNestedFile(tempZipFilePath)
	if err != nil {
		util.Log().Warning("无法创建临时压缩文件 %s , %s", tempZipFilePath, err)
		tempZipFilePath = ""
		return err
	}
	defer zipFile.Close()

	// 下载前先判断是否是可解压的格式
	format, readStream, err := archiver.Identify(fs.FileTarget[0].SourceName, fileStream)
	if err != nil {
		util.Log().Warning("无法识别文件格式 %s , %s", fs.FileTarget[0].SourceName, err)
		return err
	}

	extractor, ok := format.(archiver.Extractor)
	if !ok {
		return fmt.Errorf("file not an extractor %s", fs.FileTarget[0].SourceName)
	}

	// 只有zip格式可以多个文件同时上传
	var isZip bool
	switch extractor.(type) {
	case archiver.Zip:
		extractor = archiver.Zip{TextEncoding: encoding}
		isZip = true
	}

	// 除了zip必须下载到本地，其余的可以边下载边解压
	reader := readStream
	if isZip {
		_, err = io.Copy(zipFile, readStream)
		if err != nil {
			util.Log().Warning("无法写入临时压缩文件 %s , %s", tempZipFilePath, err)
			return err
		}

		fileStream.Close()

		// 设置文件偏移量
		zipFile.Seek(0, io.SeekStart)
		reader = zipFile
	}

	// 重设存储策略
	fs.Policy = &fs.User.Policy
	err = fs.DispatchHandler()
	if err != nil {
		return err
	}

	var wg sync.WaitGroup
	parallel := model.GetIntSetting("max_parallel_transfer", 4)
	worker := make(chan int, parallel)
	for i := 0; i < parallel; i++ {
		worker <- i
	}

	// 上传文件函数
	uploadFunc := func(fileStream io.ReadCloser, size int64, savePath, rawPath string) {
		defer func() {
			if isZip {
				worker <- 1
				wg.Done()
			}
			if err := recover(); err != nil {
				util.Log().Warning("上传压缩包内文件时出错")
				fmt.Println(err)
			}
		}()

		err := fs.UploadFromStream(ctx, &fsctx.FileStream{
			File:        fileStream,
			Size:        uint64(size),
			Name:        path.Base(savePath),
			VirtualPath: path.Dir(savePath),
		}, true)
		fileStream.Close()
		if err != nil {
			util.Log().Debug("无法上传压缩包内的文件%s , %s , 跳过", rawPath, err)
		}
	}

	// 解压缩文件，回调函数如果出错会停止解压的下一步进行，全部return nil
	err = extractor.Extract(ctx, reader, nil, func(ctx context.Context, f archiver.File) error {
		rawPath := util.FormSlash(f.NameInArchive)
		savePath := path.Join(dst, rawPath)
		// 路径是否合法
		if !strings.HasPrefix(savePath, util.FillSlash(path.Clean(dst))) {
			util.Log().Warning("%s: illegal file path", f.NameInArchive)
			return nil
		}

		// 如果是目录
		if f.FileInfo.IsDir() {
			fs.CreateDirectory(ctx, savePath)
			return nil
		}

		// 上传文件
		fileStream, err := f.Open()
		if err != nil {
			util.Log().Warning("无法打开压缩包内文件%s , %s , 跳过", rawPath, err)
			return nil
		}

		if !isZip {
			uploadFunc(fileStream, f.FileInfo.Size(), savePath, rawPath)
		} else {
			<-worker
			wg.Add(1)
			go uploadFunc(fileStream, f.FileInfo.Size(), savePath, rawPath)
		}
		return nil
	})
	wg.Wait()
	return err

}