Files
cluster/handler/upload_handler.go
2025-11-09 15:19:11 +08:00

100 lines
2.5 KiB
Go

package handler
import (
"context"
"io"
"os"
"gitea.loveuer.com/loveuer/cluster/internal/database"
"gitea.loveuer.com/loveuer/cluster/internal/model"
"gitea.loveuer.com/loveuer/cluster/internal/rerr"
"github.com/google/uuid"
"gorm.io/gorm"
)
type uploadHandlerImpl struct{}
// NewUploadHandler creates a new upload handler
func NewUploadHandler() UploadHandler {
return &uploadHandlerImpl{}
}
func (u *uploadHandlerImpl) UploadId() string {
return uuid.New().String()
}
func (u *uploadHandlerImpl) Write(ctx context.Context, sessionID string, r io.Reader, start, end int) (int, *rerr.RepositoryError) {
// Get upload session
var session model.RegistryUploadSession
if err := database.DB.Where("session_id = ?", sessionID).First(&session).Error; err != nil {
if err == gorm.ErrRecordNotFound {
return 0, &rerr.RepositoryError{
Status: 404,
Code: "BLOB_UPLOAD_UNKNOWN",
Message: "upload session not found",
}
}
return 0, rerr.ErrInternal(err)
}
// Open file for writing
file, err := os.OpenFile(session.Path, os.O_WRONLY|os.O_CREATE, 0644)
if err != nil {
return 0, rerr.ErrInternal(err)
}
defer file.Close()
// Seek to start position
if _, err := file.Seek(int64(start), 0); err != nil {
return 0, rerr.ErrInternal(err)
}
// Write data
n, err := io.CopyN(file, r, int64(end-start+1))
if err != nil && err != io.EOF {
return 0, rerr.ErrInternal(err)
}
// Update session size
session.Size = int64(start) + n
if err := database.DB.Save(&session).Error; err != nil {
return 0, rerr.ErrInternal(err)
}
return int(n), nil
}
func (u *uploadHandlerImpl) Done(ctx context.Context, blobHandler BlobHandler, sessionID string, r io.Reader, contentLength int, repo string, h model.Hash) *rerr.RepositoryError {
// Get upload session
var session model.RegistryUploadSession
if err := database.DB.Where("session_id = ?", sessionID).First(&session).Error; err != nil {
if err == gorm.ErrRecordNotFound {
return &rerr.RepositoryError{
Status: 404,
Code: "BLOB_UPLOAD_UNKNOWN",
Message: "upload session not found",
}
}
return rerr.ErrInternal(err)
}
// Read from the uploaded file instead of the request body
file, err := os.Open(session.Path)
if err != nil {
return rerr.ErrInternal(err)
}
defer file.Close()
// Store blob
if err := blobHandler.Put(ctx, repo, h, file); err != nil {
return rerr.ErrInternal(err)
}
// Clean up upload session
os.Remove(session.Path)
database.DB.Delete(&session)
return nil
}