base-golang/pkg/upload/tus.go

package upload

import (
	"context"
	"crypto/sha256"
	"errors"
	"fmt"
	"git.bvbej.com/bvbej/base-golang/pkg/color"
	"git.bvbej.com/bvbej/base-golang/pkg/token"
	"github.com/rs/cors"
	"github.com/tus/tusd/pkg/filestore"
	tus "github.com/tus/tusd/pkg/handler"
	"go.uber.org/zap"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"
)

var _ Server = (*server)(nil)

type Server interface {
	GetUploadToken(string, string, time.Duration) string
	GetFileInfo(id string) (*tus.FileInfo, error)

	Start(func(string, string, tus.FileInfo)) error
	Stop() error
}

type server struct {
	headerTokenKey string
	uploading      sync.Map
	config         Config
	token          token.Token
	store          filestore.FileStore
	logger         *zap.Logger
	httpServer     *http.Server
	completedEvent func(user, param string, info tus.FileInfo)
}

type Config struct {
	ListenAddr      string
	Path            string
	Dir             string
	Secret          string
	DisableDownload bool
	Debug           bool
}

func New(conf Config, logger *zap.Logger) Server {
	return &server{
		config:         conf,
		uploading:      sync.Map{},
		headerTokenKey: "Authorization",
		logger:         logger,
		token:          token.New(conf.Secret),
	}
}

func (s *server) GetUploadToken(user, param string, ttl time.Duration) string {
	sign, _ := s.token.JwtSign(user, param, ttl)
	return sign
}

func (s *server) GetFileInfo(id string) (*tus.FileInfo, error) {
	upload, err := s.store.GetUpload(context.Background(), id)
	if err != nil {
		return nil, err
	}
	info, err := upload.GetInfo(context.Background())
	if err != nil {
		return nil, err
	}
	return &info, nil
}

func (s *server) Start(completedEvent func(user, param string, info tus.FileInfo)) error {
	s.completedEvent = completedEvent

	composer := tus.NewStoreComposer()
	if err := os.MkdirAll(s.config.Dir, os.ModePerm); err != nil {
		return err
	}
	s.store = filestore.New(s.config.Dir)
	s.store.UseIn(composer)

	handler, err := tus.NewHandler(tus.Config{
		StoreComposer:           composer,
		BasePath:                s.config.Path,
		Logger:                  zap.NewStdLog(s.logger),
		NotifyCompleteUploads:   true,
		DisableTermination:      true,
		DisableDownload:         s.config.DisableDownload,
		RespectForwardedHeaders: strings.Contains(s.config.ListenAddr, "127.0.0.1"),
		PreUploadCreateCallback: func(hook tus.HookEvent) error {
			authStr := hook.HTTPRequest.Header.Get(s.headerTokenKey)
			jwtParse, err := s.token.JwtParse(authStr)
			if err == nil {
				_, ok := s.uploading.Load(jwtParse.ID)
				if !ok {
					s.uploading.Store(jwtParse.ID, time.Now())
					return nil
				}
			}
			return errors.New("unauthorized")
		},
		PreFinishResponseCallback: func(hook tus.HookEvent) error {
			authStr := hook.HTTPRequest.Header.Get(s.headerTokenKey)
			jwtParse, err := s.token.JwtParse(authStr)
			if err != nil {
				return errors.New("token expired")
			}
			_, ok := s.uploading.Load(jwtParse.ID)
			if ok {
				s.uploading.Delete(jwtParse.ID)
			}

			upload, err := s.store.GetUpload(context.Background(), hook.Upload.ID)
			if err != nil {
				return err
			}
			reader, err := upload.GetReader(context.Background())
			if err != nil {
				return err
			}
			var fileBytes []byte
			_, err = reader.Read(fileBytes)
			if err != nil {
				return err
			}

			hash := sha256.New()
			hash.Write(fileBytes)
			sha256Byte := hash.Sum(nil)
			sha256String := fmt.Sprintf("%x", sha256Byte)
			if !s.config.Debug && sha256String != strings.ToLower(jwtParse.ID) {
				info, _ := upload.GetInfo(context.Background())
				path, exist := info.Storage["Path"]
				if exist {
					_ = os.Remove(path)
					_ = os.Remove(path + ".info")
				}
				return errors.New("file check error")
			}
			return nil
		},
	})
	if err != nil {
		return err
	}

	go func() {
		for {
			event := <-handler.CompleteUploads
			authStr := event.HTTPRequest.Header.Get(s.headerTokenKey)
			jwtParse, _ := s.token.JwtParse(authStr)
			if s.completedEvent != nil {
				go func() {
					s.completedEvent(jwtParse.ID, jwtParse.Subject, event.Upload)
				}()
			}
		}
	}()

	//监听服务
	addr := s.config.ListenAddr
	mux := http.NewServeMux()
	mux.Handle(s.config.Path, http.StripPrefix(s.config.Path, handler))
	s.httpServer = &http.Server{
		Addr:    addr,
		Handler: cors.AllowAll().Handler(mux),
	}
	go func() {
		if err = s.httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			s.logger.Sugar().Fatal("upload server startup err", zap.Error(err))
		}
	}()
	fmt.Println(color.Green(fmt.Sprintf("* [register tusd listen %s]", addr)))

	return nil
}

func (s *server) Stop() error {
	return s.httpServer.Close()
}