From 192b6983823892a2d934bddc61aa473519f204b1 Mon Sep 17 00:00:00 2001
From: Weidi Deng <weidi_deng@icloud.com>
Date: Thu, 3 Aug 2023 16:26:01 +0800
Subject: [PATCH] calculate file checksum when updating file using wopi or
 text/code editor or webdav upload

---
 pkg/filesystem/fsctx/stream.go | 24 ++++++++++++++++++++++++
 pkg/filesystem/hooks.go        | 31 ++++++++++++++++++++-----------
 pkg/webdav/webdav.go           |  7 ++++++-
 service/explorer/file.go       |  6 +++++-
 4 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/pkg/filesystem/fsctx/stream.go b/pkg/filesystem/fsctx/stream.go
index 512270b..1a20490 100644
--- a/pkg/filesystem/fsctx/stream.go
+++ b/pkg/filesystem/fsctx/stream.go
@@ -1,9 +1,12 @@
 package fsctx
 
 import (
+	"encoding/hex"
 	"errors"
 	"github.com/HFO4/aliyun-oss-go-sdk/oss"
+	"hash"
 	"io"
+	"strings"
 	"time"
 )
 
@@ -121,3 +124,24 @@ func (file *FileStream) SetSize(size uint64) {
 func (file *FileStream) SetModel(fileModel interface{}) {
 	file.Model = fileModel
 }
+
+type ChecksumFileStream struct {
+	Md5  hash.Hash
+	Sha1 hash.Hash
+	io.Reader
+	io.Closer
+}
+
+func (file *ChecksumFileStream) Hash() string {
+	var sb strings.Builder
+
+	sb.WriteString("MD5:")
+	sb.WriteString(hex.EncodeToString(file.Md5.Sum(nil)))
+
+	sb.WriteByte(' ')
+
+	sb.WriteString("SHA1:")
+	sb.WriteString(hex.EncodeToString(file.Sha1.Sum(nil)))
+
+	return sb.String()
+}
diff --git a/pkg/filesystem/hooks.go b/pkg/filesystem/hooks.go
index a2f9ed5..8ec6fd1 100644
--- a/pkg/filesystem/hooks.go
+++ b/pkg/filesystem/hooks.go
@@ -2,6 +2,8 @@ package filesystem
 
 import (
 	"context"
+	"crypto/md5"
+	"crypto/sha1"
 	model "github.com/cloudreve/Cloudreve/v3/models"
 	"github.com/cloudreve/Cloudreve/v3/pkg/cache"
 	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
@@ -9,6 +11,7 @@ import (
 	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/fsctx"
 	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
 	"github.com/cloudreve/Cloudreve/v3/pkg/util"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"strconv"
@@ -282,23 +285,29 @@ func NewWebdavAfterUploadHook(request *http.Request) func(ctx context.Context, f
 			modtime = time.Unix(timeUnix, 0)
 		}
 	}
-	checksum := request.Header.Get("OC-Checksum")
 
 	return func(ctx context.Context, fs *FileSystem, newFile fsctx.FileHeader) error {
 		file := newFile.Info().Model.(*model.File)
 		if !modtime.IsZero() {
-			err := model.DB.Model(file).UpdateColumn("updated_at", modtime).Error
-			if err != nil {
-				return err
-			}
-		}
-
-		if checksum != "" {
-			return file.UpdateMetadata(map[string]string{
-				model.ChecksumMetadataKey: checksum,
-			})
+			return model.DB.Model(file).UpdateColumn("updated_at", modtime).Error
 		}
 
 		return nil
 	}
 }
+
+// NewChecksumFileStreamAndAfterUploadHook 创建一个计算hash的数据流和相应的钩子函数
+func NewChecksumFileStreamAndAfterUploadHook(rc io.ReadCloser) (io.ReadCloser, Hook) {
+	cfs := &fsctx.ChecksumFileStream{
+		Md5:    md5.New(),
+		Sha1:   sha1.New(),
+		Closer: rc,
+	}
+	cfs.Reader = io.TeeReader(rc, io.MultiWriter(cfs.Md5, cfs.Sha1))
+	return cfs, func(ctx context.Context, fs *FileSystem, newFile fsctx.FileHeader) error {
+		file := newFile.Info().Model.(*model.File)
+		return file.UpdateMetadata(map[string]string{
+			model.ChecksumMetadataKey: cfs.Hash(),
+		})
+	}
+}
diff --git a/pkg/webdav/webdav.go b/pkg/webdav/webdav.go
index 9b2ff1b..e82a575 100644
--- a/pkg/webdav/webdav.go
+++ b/pkg/webdav/webdav.go
@@ -344,9 +344,13 @@ func (h *Handler) handlePut(w http.ResponseWriter, r *http.Request, fs *filesyst
 	}
 	fileName := path.Base(reqPath)
 	filePath := path.Dir(reqPath)
+
+	// 计算hash
+	file, hook := filesystem.NewChecksumFileStreamAndAfterUploadHook(r.Body)
+
 	fileData := fsctx.FileStream{
 		MimeType:    r.Header.Get("Content-Type"),
-		File:        r.Body,
+		File:        file,
 		Size:        fileSize,
 		Name:        fileName,
 		VirtualPath: filePath,
@@ -391,6 +395,7 @@ func (h *Handler) handlePut(w http.ResponseWriter, r *http.Request, fs *filesyst
 
 	// rclone 请求
 	fs.Use("AfterUpload", filesystem.NewWebdavAfterUploadHook(r))
+	fs.Use("AfterUpload", hook)
 
 	// 执行上传
 	err = fs.Upload(ctx, &fileData)
diff --git a/service/explorer/file.go b/service/explorer/file.go
index 177e46d..2ff9b50 100644
--- a/service/explorer/file.go
+++ b/service/explorer/file.go
@@ -411,9 +411,12 @@ func (service *FileIDService) PutContent(ctx context.Context, c *gin.Context) se
 		return serializer.ParamErr("Invalid content-length value", err)
 	}
 
+	// 计算hash
+	file, hook := filesystem.NewChecksumFileStreamAndAfterUploadHook(c.Request.Body)
+
 	fileData := fsctx.FileStream{
 		MimeType: c.Request.Header.Get("Content-Type"),
-		File:     c.Request.Body,
+		File:     file,
 		Size:     fileSize,
 		Mode:     fsctx.Overwrite,
 	}
@@ -453,6 +456,7 @@ func (service *FileIDService) PutContent(ctx context.Context, c *gin.Context) se
 	fs.Use("BeforeUpload", filesystem.HookValidateFile)
 	fs.Use("BeforeUpload", filesystem.HookValidateCapacityDiff)
 	fs.Use("AfterUpload", filesystem.GenericAfterUpdate)
+	fs.Use("AfterUpload", hook)
 
 	// 执行上传
 	uploadCtx = context.WithValue(uploadCtx, fsctx.FileModelCtx, originFile[0])