repo.me/internal/interfaces/uploads/local.go

package uploads

import (
	"context"
	"errors"
	"fmt"
	"github.com/sirupsen/logrus"
	"io"
	"net/http"
	"nf-repo/internal/interfaces"
	"nf-repo/internal/model"
	"nf-repo/internal/util/rerr"
	"nf-repo/internal/verify"
	"os"
	"path"
	"sync"
	"time"
)

type localUploader struct {
	lock    sync.Mutex
	basedir string
	sizeMap map[string]int
}

func (l *localUploader) UploadId() string {
	id := fmt.Sprintf("%d", time.Now().UnixNano())
	l.lock.Lock()
	l.sizeMap[id] = 0
	l.lock.Unlock()

	return id
}

func (l *localUploader) Write(ctx context.Context, id string, reader io.ReadCloser, start, end int) (int, *rerr.RepositoryError) {
	var (
		err      error
		filename = path.Join(l.basedir, id)
		f        *os.File
		ok       bool
		flag     = os.O_CREATE | os.O_RDWR | os.O_TRUNC
		copied   int64
	)

	if _, ok = l.sizeMap[id]; !ok {
		return 0, &rerr.RepositoryError{
			Status:  http.StatusRequestedRangeNotSatisfiable,
			Code:    "BLOB_UPLOAD_UNKNOWN",
			Message: "Your content range doesn't match what we have",
		}
	}

	if start > 0 {
		flag = os.O_APPEND | os.O_RDWR

	}

	if f, err = os.OpenFile(filename, flag, 0644); err != nil {
		return 0, rerr.ErrInternal(err)
	}

	if copied, err = io.Copy(f, reader); err != nil {
		return 0, rerr.ErrInternal(err)
	}

	reader.Close()

	l.lock.Lock()
	l.sizeMap[id] += int(copied)
	l.lock.Unlock()

	return l.sizeMap[id], nil
}

func (l *localUploader) Done(ctx context.Context, bh interfaces.BlobHandler, id string, reader io.ReadCloser, contentLength int, repo string, hash model.Hash) *rerr.RepositoryError {
	size := verify.SizeUnknown
	if contentLength > 0 {
		size = l.sizeMap[id] + contentLength
	}

	var (
		err      error
		f        *os.File
		filename = path.Join(l.basedir, id)
	)

	if f, err = os.OpenFile(filename, os.O_RDONLY, 0644); err != nil {
		return rerr.ErrInternal(err)
	}

	in := io.NopCloser(io.MultiReader(f, reader))

	vrc, err := verify.ReadCloser(in, int64(size), hash)
	if err != nil {
		return rerr.ErrInternal(err)
	}
	defer vrc.Close()

	logrus.
		WithField("path", "localUploader.Done").
		WithField("id", id).
		WithField("size", size).
		Error()

	if err := bh.Put(ctx, repo, hash, vrc); err != nil {
		if errors.As(err, &verify.Error{}) {
			logrus.
				WithField("path", "handleBlobs.Put").
				WithField("repo", repo).
				WithField("hash", hash.String()).
				WithField("err", fmt.Sprintf("Digest mismatch: %v", err)).Debug()
			return rerr.ErrDigestMismatch
		}
		return rerr.ErrInternal(err)
	}

	delete(l.sizeMap, id)
	f.Close()

	if err = os.Remove(filename); err != nil {
		logrus.
			WithField("path", "localUploader.Done.Remove").
			WithField("filename", filename).
			WithField("err", err.Error()).
			Warn()
	}

	return nil
}

func NewLocalUploader(basedir string) interfaces.UploadHandler {
	var (
		err error
	)

	if err = os.MkdirAll(basedir, 0755); err != nil {
		logrus.
			WithField("path", "uploads.localUploader.NewLocalUploader").
			WithField("basedir", basedir).
			WithField("err", err.Error()).
			Panic()
	}

	return &localUploader{lock: sync.Mutex{}, basedir: basedir, sizeMap: make(map[string]int)}
}