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

package uploads

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"sync"
	"time"

	"nf-repo/internal/interfaces"
	"nf-repo/internal/model"
	"nf-repo/internal/tool/rerr"
	"nf-repo/internal/verify"

	"github.com/loveuer/nf/nft/log"
)

type memUploader struct {
	lock    sync.Mutex
	uploads map[string][]byte
}

func (m *memUploader) UploadId() string {
	m.lock.Lock()
	m.lock.Unlock()

	id := strconv.Itoa(int(time.Now().UnixNano()))

	m.uploads[id] = []byte{}

	return id
}

func (m *memUploader) Write(ctx context.Context, id string, reader io.ReadCloser, start, end int) (int, *rerr.RepositoryError) {
	m.lock.Lock()
	defer m.lock.Unlock()

	if start != len(m.uploads[id]) {
		return 0, &rerr.RepositoryError{
			Status:  http.StatusRequestedRangeNotSatisfiable,
			Code:    "BLOB_UPLOAD_UNKNOWN",
			Message: "Your content range doesn't match what we have",
		}
	}

	l := bytes.NewBuffer(m.uploads[id])
	size, err := io.Copy(l, reader)
	if err != nil {
		return 0, rerr.ErrInternal(err)
	}

	_ = size

	m.uploads[id] = l.Bytes()

	return len(m.uploads[id]), nil
}

func (m *memUploader) 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 = len(m.uploads[id]) + contentLength
	}

	m.lock.Lock()
	defer m.lock.Unlock()

	in := io.NopCloser(io.MultiReader(bytes.NewBuffer(m.uploads[id]), reader))

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

	if err := bh.Put(ctx, repo, hash, vrc); err != nil {
		if errors.As(err, &verify.Error{}) {
			log.Debug("memUploader: handle blob put err, repo = %s, hash = %s, err = %s", repo, hash.String(), fmt.Sprintf("Digest mismatch: %v", err))
			return rerr.ErrDigestMismatch
		}
		return rerr.ErrInternal(err)
	}

	m.uploads[id] = nil
	delete(m.uploads, id)

	return nil
}

func NewMemUploader() interfaces.UploadHandler {
	return &memUploader{
		lock:    sync.Mutex{},
		uploads: make(map[string][]byte),
	}
}