portal/tus/tus.go

package tus

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"errors"
	"git.lumeweb.com/LumeWeb/portal/cid"
	"git.lumeweb.com/LumeWeb/portal/db"
	"git.lumeweb.com/LumeWeb/portal/logger"
	"git.lumeweb.com/LumeWeb/portal/model"
	"git.lumeweb.com/LumeWeb/portal/service/files"
	"git.lumeweb.com/LumeWeb/portal/shared"
	"git.lumeweb.com/LumeWeb/portal/tusstore"
	"github.com/golang-queue/queue"
	tusd "github.com/tus/tusd/pkg/handler"
	"github.com/tus/tusd/pkg/memorylocker"
	"go.uber.org/zap"
	"golang.org/x/exp/slices"
	"gorm.io/gorm"
	"io"
)

const TUS_API_PATH = "/files/tus"

const HASH_META_HEADER = "blake3-hash"

func Init() *tusd.Handler {
	store := &tusstore.DbFileStore{
		Path: "/tmp",
	}

	shared.SetTusStore(store)

	composer := tusd.NewStoreComposer()
	composer.UseCore(store)
	composer.UseConcater(store)
	composer.UseLocker(memorylocker.New())
	composer.UseTerminater(store)
	shared.SetTusComposer(composer)

	handler, err := tusd.NewHandler(tusd.Config{
		BasePath:      "/api/v1" + TUS_API_PATH,
		StoreComposer: composer,
		PreUploadCreateCallback: func(hook tusd.HookEvent) error {
			hash := hook.Upload.MetaData[HASH_META_HEADER]

			if len(hash) == 0 {
				msg := "missing blake3-hash metadata"
				logger.Get().Debug(msg)
				return errors.New(msg)
			}

			var upload model.Upload
			result := db.Get().Where(&model.Upload{Hash: hash}).First(&upload)
			if (result.Error != nil && !errors.Is(result.Error, gorm.ErrRecordNotFound)) || result.RowsAffected > 0 {
				hashBytes, err := hex.DecodeString(hash)
				if err != nil {
					logger.Get().Debug("invalid hash", zap.Error(err))
					return err
				}

				cidString, err := cid.Encode(hashBytes, uint64(hook.Upload.Size))

				if err != nil {
					logger.Get().Debug("failed to create cid", zap.Error(err))
					return err
				}

				resp, err := json.Marshal(UploadResponse{Cid: cidString})

				if err != nil {
					logger.Get().Error("failed to create response", zap.Error(err))
					return err
				}

				return tusd.NewHTTPError(errors.New(string(resp)), 304)
			}

			return nil
		},
	})
	if err != nil {
		panic(err)
	}

	pool := queue.NewPool(5)

	shared.SetTusQueue(pool)
	shared.SetTusWorker(tusWorker)

	go tusStartup()

	return handler
}

func tusStartup() {
	tusQueue := getQueue()
	store := getStore()

	rows, err := db.Get().Model(&model.Tus{}).Rows()

	if err != nil {
		logger.Get().Error("failed to query tus uploads", zap.Error(err))
	}

	defer rows.Close()

	processedHashes := make([]string, 0)

	for rows.Next() {
		var tusUpload model.Tus
		err := db.Get().ScanRows(rows, &tusUpload)
		if err != nil {
			logger.Get().Error("failed to scan tus records", zap.Error(err))
			return
		}

		upload, err := store.GetUpload(nil, tusUpload.UploadID)
		if err != nil {
			logger.Get().Error("failed to query tus upload", zap.Error(err))
			db.Get().Delete(&tusUpload)
			continue
		}

		if slices.Contains(processedHashes, tusUpload.Hash) {
			err := terminateUpload(upload)
			if err != nil {
				logger.Get().Error("failed to terminate tus upload", zap.Error(err))
			}
			continue
		}

		if err := tusQueue.QueueTask(func(ctx context.Context) error {
			return tusWorker(&upload)
		}); err != nil {
			logger.Get().Error("failed to queue tus upload", zap.Error(err))
		} else {
			processedHashes = append(processedHashes, tusUpload.Hash)
		}
	}
}

func tusWorker(upload *tusd.Upload) error {
	info, err := (*upload).GetInfo(context.Background())
	if err != nil {
		logger.Get().Error("failed to query tus upload metadata", zap.Error(err))
		return err
	}
	file, err := (*upload).GetReader(context.Background())
	if err != nil {
		logger.Get().Error("failed reading upload", zap.Error(err))
		return err
	}

	hashHex := info.MetaData[HASH_META_HEADER]

	hashBytes, err := hex.DecodeString(hashHex)

	if err != nil {
		logger.Get().Error("failed decoding hash", zap.Error(err))
		tErr := terminateUpload(*upload)

		if tErr != nil {
			return tErr
		}
		return err
	}

	newUpload, err := files.Upload(file.(io.ReadSeeker), info.Size, hashBytes)
	tErr := terminateUpload(*upload)

	if tErr != nil {
		return tErr
	}

	if err != nil {
		return err
	}

	err = files.Pin(newUpload.Hash, newUpload.AccountID)
	if err != nil {
		return err
	}

	return nil
}

func terminateUpload(upload tusd.Upload) error {
	err := getComposer().Terminater.AsTerminatableUpload(upload).Terminate(context.Background())

	if err != nil {
		logger.Get().Error("failed deleting tus upload", zap.Error(err))
	}

	if err != nil {
		return err
	}

	return nil
}

type UploadResponse struct {
	Cid string `json:"cid"`
}

func getQueue() *queue.Queue {
	ret := shared.GetTusQueue()
	return (*ret).(*queue.Queue)
}

func getStore() *tusstore.DbFileStore {
	ret := shared.GetTusStore()
	return (*ret).(*tusstore.DbFileStore)
}
func getComposer() *tusd.StoreComposer {
	ret := shared.GetTusComposer()
	return (*ret).(*tusd.StoreComposer)
}