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.
214 lines
7.9 KiB
Go
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)
|
|
}
|