2024-04-14 21:48:27 +08:00

140 lines
2.9 KiB
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()
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)}
}