Open-IM-Server/pkg/common/db/controller/storage.go

package controller

import "C"
import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/OpenIMSDK/Open-IM-Server/pkg/common/config"
	"github.com/OpenIMSDK/Open-IM-Server/pkg/common/db/obj"
	"github.com/OpenIMSDK/Open-IM-Server/pkg/common/db/table/relation"
	"github.com/OpenIMSDK/Open-IM-Server/pkg/common/log"
	"github.com/OpenIMSDK/Open-IM-Server/pkg/errs"
	"github.com/OpenIMSDK/Open-IM-Server/pkg/proto/third"
	"github.com/OpenIMSDK/Open-IM-Server/pkg/utils"
	"github.com/google/uuid"
	"path"
	"strconv"
	"time"
)

const (
	hashPrefix     = "hash"
	tempPrefix     = "temp"
	fragmentPrefix = "fragment_"
)

type S3Database interface {
	ApplyPut(ctx context.Context, req *third.ApplyPutReq) (*third.ApplyPutResp, error)
	GetPut(ctx context.Context, req *third.GetPutReq) (*third.GetPutResp, error)
	ConfirmPut(ctx context.Context, req *third.ConfirmPutReq) (*third.ConfirmPutResp, error)
	GetUrl(ctx context.Context, req *third.GetUrlReq) (*third.GetUrlResp, error)
	GetHashInfo(ctx context.Context, req *third.GetHashInfoReq) (*third.GetHashInfoResp, error)
	CleanExpirationObject(ctx context.Context, t time.Time)
}

func NewS3Database(obj obj.Interface, hash relation.ObjectHashModelInterface, info relation.ObjectInfoModelInterface, put relation.ObjectPutModelInterface) S3Database {
	return &s3Database{
		obj:  obj,
		hash: hash,
		info: info,
		put:  put,
	}
}

type s3Database struct {
	obj  obj.Interface
	hash relation.ObjectHashModelInterface
	info relation.ObjectInfoModelInterface
	put  relation.ObjectPutModelInterface
}

// today 今天的日期
func (c *s3Database) today() string {
	return time.Now().Format("20060102")
}

// fragmentName 根据序号生成文件名
func (c *s3Database) fragmentName(index int) string {
	return fragmentPrefix + strconv.Itoa(index+1)
}

// getFragmentNum 获取分片大小和分片数量
func (c *s3Database) getFragmentNum(fragmentSize int64, objectSize int64) (int64, int) {
	if size := c.obj.MinFragmentSize(); fragmentSize < size {
		fragmentSize = size
	}
	if fragmentSize <= 0 || objectSize <= fragmentSize {
		return objectSize, 1
	} else {
		num := int(objectSize / fragmentSize)
		if objectSize%fragmentSize > 0 {
			num++
		}
		if n := c.obj.MaxFragmentNum(); num > n {
			num = n
		}
		return fragmentSize, num
	}
}

func (c *s3Database) CheckHash(hash string) error {
	val, err := hex.DecodeString(hash)
	if err != nil {
		return err
	}
	if len(val) != md5.Size {
		return errs.ErrArgs.Wrap("invalid hash")
	}
	return nil
}

func (c *s3Database) urlName(name string) string {
	return config.Config.Object.ApiURL + name
}

func (c *s3Database) UUID() string {
	return uuid.New().String()
}

func (c *s3Database) HashName(hash string) string {
	return path.Join(hashPrefix, hash+"_"+c.today()+"_"+c.UUID())
}

func (c *s3Database) isNotFound(err error) bool {
	return relation.IsNotFound(err)
}

func (c *s3Database) ApplyPut(ctx context.Context, req *third.ApplyPutReq) (*third.ApplyPutResp, error) {
	if err := c.CheckHash(req.Hash); err != nil {
		return nil, err
	}
	if err := c.obj.CheckName(req.Name); err != nil {
		return nil, err
	}
	if req.ValidTime != 0 && req.ValidTime <= time.Now().UnixMilli() {
		return nil, errors.New("invalid ValidTime")
	}
	var expirationTime *time.Time
	if req.ValidTime != 0 {
		expirationTime = utils.ToPtr(time.UnixMilli(req.ValidTime))
	}
	if hash, err := c.hash.Take(ctx, req.Hash, c.obj.Name()); err == nil {
		o := relation.ObjectInfoModel{
			Name:        req.Name,
			Hash:        hash.Hash,
			ValidTime:   expirationTime,
			ContentType: req.ContentType,
			CreateTime:  time.Now(),
		}
		if err := c.info.SetObject(ctx, &o); err != nil {
			return nil, err
		}
		return &third.ApplyPutResp{Url: c.urlName(o.Name)}, nil // 服务器已存在
	} else if !c.isNotFound(err) {
		return nil, err
	}
	// 新上传
	var pack int
	const effective = time.Hour * 24 * 2
	req.FragmentSize, pack = c.getFragmentNum(req.FragmentSize, req.Size)
	put := relation.ObjectPutModel{
		PutID:         c.UUID(),
		Hash:          req.Hash,
		Name:          req.Name,
		ObjectSize:    req.Size,
		ContentType:   req.ContentType,
		FragmentSize:  req.FragmentSize,
		ValidTime:     expirationTime,
		EffectiveTime: time.Now().Add(effective),
	}
	put.Path = path.Join(tempPrefix, c.today(), req.Hash, put.PutID)
	putURLs := make([]string, 0, pack)
	for i := 0; i < pack; i++ {
		url, err := c.obj.PresignedPutURL(ctx, &obj.ApplyPutArgs{
			Bucket:        c.obj.TempBucket(),
			Name:          path.Join(put.Path, c.fragmentName(i)),
			Effective:     effective,
			MaxObjectSize: req.FragmentSize,
		})
		if err != nil {
			return nil, err
		}
		putURLs = append(putURLs, url)
	}
	urlsJsonData, err := json.Marshal(putURLs)
	if err != nil {
		return nil, err
	}
	t := md5.Sum(urlsJsonData)
	put.PutURLsHash = hex.EncodeToString(t[:])
	put.CreateTime = time.Now()
	if err := c.put.Create(ctx, []*relation.ObjectPutModel{&put}); err != nil {
		return nil, err
	}
	return &third.ApplyPutResp{
		PutID:        put.PutID,
		FragmentSize: put.FragmentSize,
		PutURLs:      putURLs,
		ValidTime:    put.EffectiveTime.UnixMilli(),
	}, nil
}

func (c *s3Database) GetPut(ctx context.Context, req *third.GetPutReq) (*third.GetPutResp, error) {
	up, err := c.put.Take(ctx, req.PutID)
	if err != nil {
		return nil, err
	}
	if up.Complete {
		return nil, errors.New("up completed")
	}
	_, pack := c.getFragmentNum(up.FragmentSize, up.ObjectSize)
	fragments := make([]*third.GetPutFragment, pack)
	for i := 0; i < pack; i++ {
		name := path.Join(up.Path, c.fragmentName(i))
		o, err := c.obj.GetObjectInfo(ctx, &obj.BucketObject{
			Bucket: c.obj.TempBucket(),
			Name:   name,
		})
		if err != nil {
			if c.obj.IsNotFound(err) {
				fragments[i] = &third.GetPutFragment{}
				continue
			}
			return nil, err
		}
		fragments[i] = &third.GetPutFragment{Size: o.Size, Hash: o.Hash}
	}
	var validTime int64
	if up.ValidTime != nil {
		validTime = up.ValidTime.UnixMilli()
	}
	return &third.GetPutResp{
		FragmentSize: up.FragmentSize,
		Size:         up.ObjectSize,
		Name:         up.Name,
		Hash:         up.Hash,
		Fragments:    fragments,
		PutURLsHash:  up.PutURLsHash,
		ContentType:  up.ContentType,
		ValidTime:    validTime,
	}, nil
}

func (c *s3Database) ConfirmPut(ctx context.Context, req *third.ConfirmPutReq) (_ *third.ConfirmPutResp, _err error) {
	put, err := c.put.Take(ctx, req.PutID)
	if err != nil {
		return nil, err
	}
	_, pack := c.getFragmentNum(put.FragmentSize, put.ObjectSize)
	defer func() {
		if _err == nil {
			// 清理上传的碎片
			err := c.obj.DeleteObject(ctx, &obj.BucketObject{Bucket: c.obj.TempBucket(), Name: put.Path})
			if err != nil {
				log.ZError(ctx, "deleteObject failed", err, "Bucket", c.obj.TempBucket(), "Path", put.Path)
			}
		}
	}()
	if put.Complete {
		return nil, errors.New("put completed")
	}
	now := time.Now().UnixMilli()
	if put.EffectiveTime.UnixMilli() < now {
		return nil, errors.New("upload expired")
	}
	if put.ValidTime != nil && put.ValidTime.UnixMilli() < now {
		return nil, errors.New("object expired")
	}
	if hash, err := c.hash.Take(ctx, put.Hash, c.obj.Name()); err == nil {
		o := relation.ObjectInfoModel{
			Name:        put.Name,
			Hash:        hash.Hash,
			ValidTime:   put.ValidTime,
			ContentType: put.ContentType,
			CreateTime:  time.Now(),
		}
		if err := c.info.SetObject(ctx, &o); err != nil {
			return nil, err
		}
		defer func() {
			err := c.obj.DeleteObject(ctx, &obj.BucketObject{
				Bucket: c.obj.TempBucket(),
				Name:   put.Path,
			})
			if err != nil {
				log.ZError(ctx, "DeleteObject", err, "Bucket", c.obj.TempBucket(), "Path", put.Path)
			}
		}()
		// 服务端已存在
		return &third.ConfirmPutResp{
			Url: c.urlName(o.Name),
		}, nil
	} else if !c.isNotFound(err) {
		return nil, err
	}
	src := make([]obj.BucketObject, pack)
	for i := 0; i < pack; i++ {
		name := path.Join(put.Path, c.fragmentName(i))
		o, err := c.obj.GetObjectInfo(ctx, &obj.BucketObject{
			Bucket: c.obj.TempBucket(),
			Name:   name,
		})
		if err != nil {
			return nil, err
		}
		if i+1 == pack { // 最后一个
			size := put.ObjectSize - put.FragmentSize*int64(i)
			if size != o.Size {
				return nil, fmt.Errorf("last fragment %d size %d not equal to %d hash %s", i, o.Size, size, o.Hash)
			}
		} else {
			if o.Size != put.FragmentSize {
				return nil, fmt.Errorf("fragment %d size %d not equal to %d hash %s", i, o.Size, put.FragmentSize, o.Hash)
			}
		}
		src[i] = obj.BucketObject{
			Bucket: c.obj.TempBucket(),
			Name:   name,
		}
	}
	dst := &obj.BucketObject{
		Bucket: c.obj.DataBucket(),
		Name:   c.HashName(put.Hash),
	}
	if len(src) == 1 { // 未分片直接触发copy
		// 检查数据完整性,避免脏数据
		o, err := c.obj.GetObjectInfo(ctx, &src[0])
		if err != nil {
			return nil, err
		}
		if put.ObjectSize != o.Size {
			return nil, fmt.Errorf("size mismatching should %d reality %d", put.ObjectSize, o.Size)
		}
		if put.Hash != o.Hash {
			return nil, fmt.Errorf("hash mismatching should %s reality %s", put.Hash, o.Hash)
		}
		if err := c.obj.CopyObject(ctx, &src[0], dst); err != nil {
			return nil, err
		}
	} else {
		tempBucket := &obj.BucketObject{
			Bucket: c.obj.TempBucket(),
			Name:   path.Join(put.Path, "merge_"+c.UUID()),
		}
		defer func() { // 清理合成的文件
			if err := c.obj.DeleteObject(ctx, tempBucket); err != nil {
				log.ZError(ctx, "DeleteObject", err, "Bucket", tempBucket.Bucket, "Path", tempBucket.Name)
			}
		}()
		err := c.obj.ComposeObject(ctx, src, tempBucket)
		if err != nil {
			return nil, err
		}
		info, err := c.obj.GetObjectInfo(ctx, tempBucket)
		if err != nil {
			return nil, err
		}
		if put.ObjectSize != info.Size {
			return nil, fmt.Errorf("size mismatch should %d reality %d", put.ObjectSize, info.Size)
		}
		if put.Hash != info.Hash {
			return nil, fmt.Errorf("hash mismatch should %s reality %s", put.Hash, info.Hash)
		}
		if err := c.obj.CopyObject(ctx, tempBucket, dst); err != nil {
			return nil, err
		}
	}
	h := &relation.ObjectHashModel{
		Hash:       put.Hash,
		Engine:     c.obj.Name(),
		Size:       put.ObjectSize,
		Bucket:     c.obj.DataBucket(),
		Name:       dst.Name,
		CreateTime: time.Now(),
	}
	if err := c.hash.Create(ctx, []*relation.ObjectHashModel{h}); err != nil {
		return nil, err
	}
	o := &relation.ObjectInfoModel{
		Name:        put.Name,
		Hash:        put.Hash,
		ContentType: put.ContentType,
		ValidTime:   put.ValidTime,
		CreateTime:  time.Now(),
	}
	if err := c.info.SetObject(ctx, o); err != nil {
		return nil, err
	}
	if err := c.put.SetCompleted(ctx, put.PutID); err != nil {
		log.ZError(ctx, "SetCompleted", err, "PutID", put.PutID)
	}
	return &third.ConfirmPutResp{
		Url: c.urlName(o.Name),
	}, nil
}

func (c *s3Database) GetUrl(ctx context.Context, req *third.GetUrlReq) (*third.GetUrlResp, error) {
	info, err := c.info.Take(ctx, req.Name)
	if err != nil {
		return nil, err
	}
	if info.ValidTime != nil && info.ValidTime.Before(time.Now()) {
		return nil, errs.ErrRecordNotFound.Wrap("object expired")
	}
	hash, err := c.hash.Take(ctx, info.Hash, c.obj.Name())
	if err != nil {
		return nil, err
	}
	opt := obj.HeaderOption{ContentType: info.ContentType}
	if req.Attachment {
		opt.Filename = info.Name
	}
	u, err := c.obj.PresignedGetURL(ctx, hash.Bucket, hash.Name, time.Duration(req.Expires)*time.Millisecond, &opt)
	if err != nil {
		return nil, err
	}
	return &third.GetUrlResp{
		Url:  u,
		Size: hash.Size,
		Hash: hash.Hash,
	}, nil
}

func (c *s3Database) CleanExpirationObject(ctx context.Context, t time.Time) {
	// 清理上传产生的临时文件
	c.cleanPutTemp(ctx, t, 10)
	// 清理hash引用全过期的文件
	c.cleanExpirationObject(ctx, t)
	// 清理没有引用的hash对象
	c.clearNoCitation(ctx, c.obj.Name(), 10)
}

func (c *s3Database) cleanPutTemp(ctx context.Context, t time.Time, num int) {
	for {
		puts, err := c.put.FindExpirationPut(ctx, t, num)
		if err != nil {
			log.ZError(ctx, "FindExpirationPut", err, "Time", t, "Num", num)
			return
		}
		if len(puts) == 0 {
			return
		}
		for _, put := range puts {
			err := c.obj.DeleteObject(ctx, &obj.BucketObject{Bucket: c.obj.TempBucket(), Name: put.Path})
			if err != nil {
				log.ZError(ctx, "DeleteObject", err, "Bucket", c.obj.TempBucket(), "Path", put.Path)
				return
			}
		}
		ids := utils.Slice(puts, func(e *relation.ObjectPutModel) string { return e.PutID })
		err = c.put.DelPut(ctx, ids)
		if err != nil {
			log.ZError(ctx, "DelPut", err, "PutID", ids)
			return
		}
	}
}

func (c *s3Database) cleanExpirationObject(ctx context.Context, t time.Time) {
	err := c.info.DeleteExpiration(ctx, t)
	if err != nil {
		log.ZError(ctx, "DeleteExpiration", err, "Time", t)
	}
}

func (c *s3Database) clearNoCitation(ctx context.Context, engine string, limit int) {
	for {
		list, err := c.hash.DeleteNoCitation(ctx, engine, limit)
		if err != nil {
			log.ZError(ctx, "DeleteNoCitation", err, "Engine", engine, "Limit", limit)
			return
		}
		if len(list) == 0 {
			return
		}
		var hasErr bool
		for _, h := range list {
			err := c.obj.DeleteObject(ctx, &obj.BucketObject{Bucket: h.Bucket, Name: h.Name})
			if err != nil {
				hasErr = true
				log.ZError(ctx, "DeleteObject", err, "Bucket", h.Bucket, "Path", h.Name)
				continue
			}
		}
		if hasErr {
			return
		}
	}
}

func (c *s3Database) GetHashInfo(ctx context.Context, req *third.GetHashInfoReq) (*third.GetHashInfoResp, error) {
	if err := c.CheckHash(req.Hash); err != nil {
		return nil, err
	}
	o, err := c.hash.Take(ctx, req.Hash, c.obj.Name())
	if err != nil {
		return nil, err
	}
	return &third.GetHashInfoResp{
		Hash: o.Hash,
		Size: o.Size,
	}, nil
}