memos/server/router/api/v1/attachment_upload.go
boojack b3e67399b8 feat(attachment): upload files in resumable chunks
Add UploadAttachment, a unary RPC that accepts a file in bounded chunks
and streams it to a temp file instead of holding the whole blob in one
request body and in memory. A call carrying a spec starts an upload and
returns an opaque upload_id; later calls carry that id with
write_offset,
data, and finish_write, following the Cloud Storage WriteObject shape.
Retrying the last chunk after a lost response is idempotent, uploads are
bound to their owner, and they expire after 30 minutes of inactivity.

CreateAttachment and the chunked finalize now share one processing
pipeline for motion-photo detection, EXIF stripping, and storage. The
motion-photo detector works over an io.ReaderAt so large JPEGs are never
loaded whole, and the request body cap is a single per-procedure lookup.

The web editor uploads through the new RPC in 2 MiB chunks with bounded
concurrency under the server's per-user limit.
2026-09-08 21:07:06 +08:00

214 lines
7.9 KiB
Go

package v1
import (
"context"
"crypto/sha256"
"io"
"os"
"time"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"google.golang.org/protobuf/proto"
v1pb "github.com/usememos/memos/proto/gen/api/v1"
"github.com/usememos/memos/store"
)
// UploadAttachment accepts bounded unary chunks. A call carrying a spec opens
// a new upload; a call carrying an upload ID continues one. Either kind may
// write data and finalize.
func (s *APIV1Service) UploadAttachment(ctx context.Context, request *v1pb.UploadAttachmentRequest) (*v1pb.UploadAttachmentResponse, error) {
user, err := s.requireCurrentSpaceUser(ctx)
if err != nil {
return nil, err
}
if len(request.Data) > attachmentUploadChunkSize {
return nil, status.Errorf(codes.ResourceExhausted, "upload chunk exceeds the limit")
}
var id string
var upload *attachmentUpload
switch u := request.Upload.(type) {
case *v1pb.UploadAttachmentRequest_Spec:
id, upload, err = s.startAttachmentUpload(ctx, request, u.Spec, user.ID)
if err != nil {
return nil, err
}
upload.mu.Lock()
case *v1pb.UploadAttachmentRequest_UploadId:
id = u.UploadId
upload, err = s.attachmentUploads.get(id, user.ID)
if err != nil {
return nil, err
}
upload.mu.Lock()
if !time.Now().Before(upload.expireTime) {
upload.mu.Unlock()
return nil, status.Errorf(codes.NotFound, "upload not found or expired")
}
default:
return nil, status.Errorf(codes.InvalidArgument, "spec or upload_id is required")
}
defer upload.mu.Unlock()
if err := ctx.Err(); err != nil {
return nil, status.FromContextError(err).Err()
}
if err := upload.write(ctx, request); err != nil {
return nil, err
}
var attachment *v1pb.Attachment
if upload.complete {
// Recheck access and existence; a completed upload must not resurrect a
// deleted attachment or disclose metadata after access is revoked.
attachment, err = s.GetAttachment(ctx, &v1pb.GetAttachmentRequest{Name: "attachments/" + upload.uid})
} else if request.FinishWrite {
if upload.committedSize != upload.totalSize {
return nil, status.Errorf(codes.FailedPrecondition, "upload is incomplete")
}
attachment, err = s.finishAttachmentUpload(ctx, upload)
if err == nil {
upload.complete = true
os.Remove(upload.path)
}
}
if err != nil {
return nil, err
}
upload.expireTime = time.Now().Add(attachmentUploadTTL)
return &v1pb.UploadAttachmentResponse{
UploadId: id, CommittedSize: upload.committedSize, Attachment: attachment, MaxChunkSize: attachmentUploadChunkSize,
}, nil
}
// startAttachmentUpload validates the spec and registers a new upload. Every
// check that can fail on the accompanying data runs first, so a rejected call
// never leaves an orphaned upload behind.
func (s *APIV1Service) startAttachmentUpload(ctx context.Context, request *v1pb.UploadAttachmentRequest, spec *v1pb.UploadAttachmentSpec, ownerID int32) (string, *attachmentUpload, error) {
if spec.Attachment == nil {
return "", nil, status.Errorf(codes.InvalidArgument, "spec.attachment is required")
}
if len(spec.Attachment.Content) != 0 {
return "", nil, status.Errorf(codes.InvalidArgument, "spec.attachment.content must be empty; send file bytes in data")
}
if proto.Size(spec) > attachmentUploadMetadataLimit || spec.TotalSize < 0 {
return "", nil, status.Errorf(codes.InvalidArgument, "invalid attachment metadata or total_size")
}
if request.WriteOffset != 0 {
return "", nil, status.Errorf(codes.OutOfRange, "write_offset must be 0 for a new upload")
}
if int64(len(request.Data)) > spec.TotalSize {
return "", nil, status.Errorf(codes.InvalidArgument, "data exceeds total_size")
}
setting, err := s.Store.GetInstanceStorageSetting(ctx)
if err != nil {
return "", nil, status.Errorf(codes.Internal, "failed to get storage setting: %v", err)
}
if err := checkUploadSize(setting, spec.TotalSize); err != nil {
return "", nil, err
}
// Validate before allocating a temporary file. Keep the original MIME type
// so an omitted type can be sniffed from real bytes at finalization.
metadata := proto.CloneOf(spec.Attachment)
validationMetadata := proto.CloneOf(metadata)
if validationMetadata.Type == "" {
// Type-dependent metadata validation must wait for content sniffing.
validationMetadata.MediaMetadata = nil
}
create, err := s.prepareAttachment(ctx, &v1pb.CreateAttachmentRequest{Attachment: validationMetadata, AttachmentId: spec.AttachmentId})
if err != nil {
return "", nil, err
}
upload := &attachmentUpload{ownerID: ownerID, metadata: metadata, uid: create.UID, totalSize: spec.TotalSize}
id, err := s.attachmentUploads.create(s.Profile.Data, upload)
if err != nil {
return "", nil, err
}
return id, upload, nil
}
func (u *attachmentUpload) write(ctx context.Context, request *v1pb.UploadAttachmentRequest) error {
if len(request.Data) == 0 {
if request.FinishWrite && request.WriteOffset != u.committedSize {
return status.Errorf(codes.OutOfRange, "write_offset must equal committed_size")
}
return nil
}
digest := sha256.Sum256(request.Data)
if u.committedSize > 0 && request.WriteOffset+int64(len(request.Data)) == u.committedSize && digest == u.lastDigest {
return nil // A lost response can be retried without appending bytes twice.
}
if u.complete {
return status.Errorf(codes.FailedPrecondition, "upload is already complete")
}
if request.WriteOffset != u.committedSize {
return status.Errorf(codes.OutOfRange, "write_offset must equal committed_size")
}
if int64(len(request.Data)) > u.totalSize-u.committedSize {
return status.Errorf(codes.InvalidArgument, "data exceeds total_size")
}
file, err := os.OpenFile(u.path, os.O_WRONLY, 0600)
if err != nil {
return status.Errorf(codes.Internal, "failed to open upload file: %v", err)
}
defer file.Close()
n, err := file.WriteAt(request.Data, u.committedSize)
if err == nil && n != len(request.Data) {
err = io.ErrShortWrite
}
if err == nil {
err = ctx.Err()
}
if err != nil {
if truncateErr := file.Truncate(u.committedSize); truncateErr != nil {
// The file no longer matches committedSize; revoke the upload.
u.expireTime = time.Time{}
os.Remove(u.path)
}
return status.Errorf(codes.Internal, "failed to write upload chunk: %v", err)
}
u.lastDigest = digest
u.committedSize += int64(n)
return nil
}
func (s *APIV1Service) finishAttachmentUpload(ctx context.Context, upload *attachmentUpload) (*v1pb.Attachment, error) {
if upload.finalizeAttempted {
// A database operation may commit before returning an error. Resolve that
// outcome before saving another object under the upload's stable UID.
persisted, err := s.Store.GetAttachment(ctx, &store.FindAttachment{UID: &upload.uid})
if err != nil {
return nil, status.Errorf(codes.Internal, "failed to resolve previous finalization: %v", err)
}
if persisted != nil {
if persisted.CreatorID != upload.ownerID {
return nil, status.Errorf(codes.AlreadyExists, "attachment ID already exists")
}
return s.GetAttachment(ctx, &v1pb.GetAttachmentRequest{Name: "attachments/" + upload.uid})
}
}
setting, err := s.Store.GetInstanceStorageSetting(ctx)
if err != nil {
return nil, status.Errorf(codes.Internal, "failed to get storage setting: %v", err)
}
if err := checkUploadSize(setting, upload.totalSize); err != nil {
return nil, err
}
file, err := os.Open(upload.path)
if err != nil {
return nil, status.Errorf(codes.Internal, "failed to open upload file: %v", err)
}
defer file.Close()
metadata := proto.CloneOf(upload.metadata)
metadata.Content = make([]byte, min(upload.totalSize, 512))
if _, err := io.ReadFull(file, metadata.Content); err != nil {
return nil, status.Errorf(codes.Internal, "failed to read attachment header: %v", err)
}
// Revalidate memo permissions and media metadata with the actual MIME type.
create, err := s.prepareAttachment(ctx, &v1pb.CreateAttachmentRequest{Attachment: metadata, AttachmentId: upload.uid})
if err != nil {
return nil, err
}
create.Size = upload.totalSize
upload.finalizeAttempted = true
return s.processAndSaveAttachment(ctx, create, setting, file)
}