perf: sync protocol and core hardening updates

This commit is contained in:
A 2026-07-11 19:48:26 +08:00
parent 152fed3b87
commit 4390ebf5a9
283 changed files with 29231 additions and 2295 deletions

View file

@ -8,10 +8,10 @@ import (
"fmt"
"image"
"image/color"
"io"
stddraw "image/draw"
_ "image/jpeg" // 注册 jpeg DecodeConfig,用于读取上传头像/图片尺寸
"image/png"
"io"
"math"
"strings"
"time"
@ -48,14 +48,40 @@ func (s *Service) UploadProfilePhotoKind(ctx context.Context, ownerType domain.P
// CreatePhotoFromUpload 把已上传文件组装成 Photo(不绑定 profile_photos),用于频道头像 / 图片消息。
func (s *Service) CreatePhotoFromUpload(ctx context.Context, file domain.UploadedFileRef) (domain.Photo, error) {
data, err := s.assembleUpload(ctx, file.OwnerUserID, file.FileID, file.Parts)
intentHash, err := uploadedMediaIntentHash(domain.UploadedMediaPhoto, file, nil)
if err != nil {
return domain.Photo{}, err
}
if photo, found, err := s.replayUploadedPhoto(ctx, file, intentHash); err != nil || found {
return photo, err
}
data, err := s.readUploadBytes(ctx, file.OwnerUserID, file.FileID, file.Parts)
if err != nil {
return domain.Photo{}, err
}
if len(data) == 0 {
return domain.Photo{}, domain.ErrPhotoInvalid
}
return s.createPhoto(ctx, data, photoSizeSpecsForMessage(data))
photo, err := s.createPhoto(ctx, data, photoSizeSpecsForMessage(data))
if err != nil {
return domain.Photo{}, err
}
receipt, err := s.commitUploadedMediaReceipt(ctx, file, domain.UploadedMediaPhoto, intentHash, photo.ID)
if err != nil {
return domain.Photo{}, err
}
if receipt.MediaID != photo.ID {
winner, found, err := s.media.GetPhoto(ctx, receipt.MediaID)
if err != nil {
return domain.Photo{}, err
}
if !found {
return domain.Photo{}, fmt.Errorf("concurrent upload receipt references missing photo %d", receipt.MediaID)
}
photo = winner
}
s.cleanupMaterializedUpload(ctx, file, "photo materialized")
return photo, nil
}
// CreatePhotoFromBytes stores already-fetched image bytes as a message Photo.
@ -202,6 +228,13 @@ func validateAvatarMarkupSize(size domain.PhotoSize) error {
// CreateDocumentFromUpload 把已上传文件组装成 Document(文件/视频/音频/gif/贴纸消息),落 blob + documents。
func (s *Service) CreateDocumentFromUpload(ctx context.Context, file domain.UploadedFileRef, spec domain.DocumentSpec) (domain.Document, error) {
intentHash, err := uploadedMediaIntentHash(domain.UploadedMediaDocument, file, &spec)
if err != nil {
return domain.Document{}, err
}
if doc, found, err := s.replayUploadedDocument(ctx, file, intentHash); err != nil || found {
return doc, err
}
body, err := s.assembleUploadBlob(ctx, file.OwnerUserID, file.FileID, file.Parts)
if err != nil {
return domain.Document{}, err
@ -234,11 +267,13 @@ func (s *Service) CreateDocumentFromUpload(ctx context.Context, file domain.Uplo
DCID: s.dc,
Attributes: spec.Attributes,
}
thumbMaterialized := false
if spec.Thumb != nil {
thumbData, err := s.assembleUpload(ctx, spec.Thumb.OwnerUserID, spec.Thumb.FileID, spec.Thumb.Parts)
thumbData, err := s.readUploadBytes(ctx, spec.Thumb.OwnerUserID, spec.Thumb.FileID, spec.Thumb.Parts)
if err == nil && len(thumbData) > 0 {
if thumb, err := s.putDocumentThumb(ctx, docID, thumbData); err == nil {
doc.Thumbs = []domain.PhotoSize{thumb}
thumbMaterialized = true
}
}
}
@ -250,12 +285,23 @@ func (s *Service) CreateDocumentFromUpload(ctx context.Context, file domain.Uplo
if err := s.media.PutDocument(ctx, doc); err != nil {
return domain.Document{}, err
}
if err := s.cleanupUploadParts(ctx, file.OwnerUserID, file.FileID); err != nil {
s.log.Warn("cleanup assembled document upload parts failed",
zap.Int64("owner_user_id", file.OwnerUserID),
zap.Int64("file_id", file.FileID),
zap.Int64("document_id", docID),
zap.Error(err))
receipt, err := s.commitUploadedMediaReceipt(ctx, file, domain.UploadedMediaDocument, intentHash, doc.ID)
if err != nil {
return domain.Document{}, err
}
if receipt.MediaID != doc.ID {
winner, found, err := s.media.GetDocument(ctx, receipt.MediaID)
if err != nil {
return domain.Document{}, err
}
if !found {
return domain.Document{}, fmt.Errorf("concurrent upload receipt references missing document %d", receipt.MediaID)
}
doc = winner
}
s.cleanupMaterializedUpload(ctx, file, "document materialized")
if spec.Thumb != nil && thumbMaterialized {
s.cleanupMaterializedUpload(ctx, *spec.Thumb, "document thumbnail materialized")
}
return doc, nil
}
@ -277,6 +323,7 @@ var faststartVideoMimes = map[string]bool{
// 此时只发生几次 16 字节读,不读整段媒体。
// 2. 仅 moov 在末尾时才重写;且优先走流式(仅 ftyp+moov 进内存,mdat 大块分块流式拼接),
// 不把整段视频 2× 驻留内存。moov 非末尾的罕见排布回退到全量重排。
//
// 任何不适用/失败都返回原 body,绝不让上传失败或损坏数据。
func (s *Service) maybeFaststartVideoBlob(ctx context.Context, mimeType string, body assembledUploadBlob) assembledUploadBlob {
if !faststartVideoMimes[strings.ToLower(strings.TrimSpace(mimeType))] {

View file

@ -26,6 +26,7 @@ type fakeMediaStore struct {
parts map[string][]domain.UploadPart
webPages map[int64]domain.MessageWebPage
seedState map[string]string
receipts map[string]domain.UploadedMediaReceipt
}
func newFakeMediaStore() *fakeMediaStore {
@ -36,9 +37,35 @@ func newFakeMediaStore() *fakeMediaStore {
sets: map[int64]domain.StickerSet{},
parts: map[string][]domain.UploadPart{},
seedState: map[string]string{},
receipts: map[string]domain.UploadedMediaReceipt{},
}
}
func fakeUploadReceiptKey(ownerUserID, fileID int64) string {
return fmt.Sprintf("%d/%d", ownerUserID, fileID)
}
func (f *fakeMediaStore) GetUploadedMediaReceipt(_ context.Context, ownerUserID, fileID int64) (domain.UploadedMediaReceipt, bool, error) {
f.mu.Lock()
defer f.mu.Unlock()
receipt, ok := f.receipts[fakeUploadReceiptKey(ownerUserID, fileID)]
receipt.IntentHash = append([]byte(nil), receipt.IntentHash...)
return receipt, ok, nil
}
func (f *fakeMediaStore) PutUploadedMediaReceipt(_ context.Context, receipt domain.UploadedMediaReceipt) (domain.UploadedMediaReceipt, bool, error) {
f.mu.Lock()
defer f.mu.Unlock()
key := fakeUploadReceiptKey(receipt.OwnerUserID, receipt.FileID)
if stored, ok := f.receipts[key]; ok {
stored.IntentHash = append([]byte(nil), stored.IntentHash...)
return stored, false, nil
}
receipt.IntentHash = append([]byte(nil), receipt.IntentHash...)
f.receipts[key] = receipt
return receipt, true, nil
}
func (f *fakeMediaStore) SaveFilePart(_ context.Context, part domain.UploadPart) error {
f.mu.Lock()
defer f.mu.Unlock()

View file

@ -514,6 +514,20 @@ func orderDocuments(docs []domain.Document, ids []int64) []domain.Document {
// assembleUpload 把已上传分片按 part 顺序拼成完整字节,并清理分片。
// expectedParts>0 时校验分片连续且齐全。
func (s *Service) assembleUpload(ctx context.Context, ownerUserID, fileID int64, expectedParts int) ([]byte, error) {
buf, err := s.readUploadBytes(ctx, ownerUserID, fileID, expectedParts)
if err != nil {
return nil, err
}
if err := s.cleanupUploadParts(ctx, ownerUserID, fileID); err != nil {
return nil, err
}
return buf, nil
}
// readUploadBytes validates and reads all parts without consuming them. Message-media
// materialization persists an upload receipt before cleanup; callers that do not need replayability
// continue to use assembleUpload.
func (s *Service) readUploadBytes(ctx context.Context, ownerUserID, fileID int64, expectedParts int) ([]byte, error) {
parts, _, err := s.loadAndValidateUploadParts(ctx, ownerUserID, fileID, expectedParts)
if err != nil {
return nil, err
@ -532,9 +546,6 @@ func (s *Service) assembleUpload(ctx context.Context, ownerUserID, fileID int64,
}
buf = append(buf, data...)
}
if err := s.cleanupUploadParts(ctx, ownerUserID, fileID); err != nil {
return nil, err
}
return buf, nil
}

View file

@ -124,6 +124,51 @@ func TestCreateDocumentFromUploadStreamsBodyAndCleansParts(t *testing.T) {
if string(body) != strings.Join(parts, "") {
t.Fatalf("body blob mismatch")
}
replayed, err := svc.CreateDocumentFromUpload(ctx,
domain.UploadedFileRef{OwnerUserID: 10, FileID: 200, Parts: len(parts), Name: "large.bin", Big: true},
domain.DocumentSpec{MimeType: "application/octet-stream"},
)
if err != nil {
t.Fatalf("replay CreateDocumentFromUpload after part cleanup: %v", err)
}
if replayed.ID != doc.ID || replayed.AccessHash != doc.AccessHash {
t.Fatalf("replayed document = %d/%d, want original %d/%d", replayed.ID, replayed.AccessHash, doc.ID, doc.AccessHash)
}
if _, err := svc.CreateDocumentFromUpload(ctx,
domain.UploadedFileRef{OwnerUserID: 10, FileID: 200, Parts: len(parts), Name: "large.bin", Big: true},
domain.DocumentSpec{MimeType: "text/plain"},
); !errors.Is(err, domain.ErrFilePartsInvalid) {
t.Fatalf("changed materialization intent err = %v, want ErrFilePartsInvalid", err)
}
}
func TestCreatePhotoFromUploadReceiptReplaysAfterPartCleanup(t *testing.T) {
ctx := context.Background()
media := newFakeMediaStore()
svc, _ := newUploadPartTestService(t, media, domain.UploadPartQuota{})
file := domain.UploadedFileRef{OwnerUserID: 10, FileID: 201, Parts: 1, Name: "photo.jpg"}
if _, err := svc.SaveFilePart(ctx, file.OwnerUserID, file.FileID, 0, []byte("image-bytes")); err != nil {
t.Fatalf("SaveFilePart: %v", err)
}
first, err := svc.CreatePhotoFromUpload(ctx, file)
if err != nil {
t.Fatalf("CreatePhotoFromUpload: %v", err)
}
if remaining, err := media.LoadFileParts(ctx, file.OwnerUserID, file.FileID); err != nil || len(remaining) != 0 {
t.Fatalf("upload parts after photo materialization = %+v err=%v", remaining, err)
}
replayed, err := svc.CreatePhotoFromUpload(ctx, file)
if err != nil {
t.Fatalf("replay CreatePhotoFromUpload: %v", err)
}
if replayed.ID != first.ID || replayed.AccessHash != first.AccessHash {
t.Fatalf("replayed photo = %d/%d, want original %d/%d", replayed.ID, replayed.AccessHash, first.ID, first.AccessHash)
}
changed := file
changed.Name = "different.jpg"
if _, err := svc.CreatePhotoFromUpload(ctx, changed); !errors.Is(err, domain.ErrFilePartsInvalid) {
t.Fatalf("changed photo intent err = %v, want ErrFilePartsInvalid", err)
}
}
type countingUploadPartBackend struct {

View file

@ -0,0 +1,106 @@
package files
import (
"bytes"
"context"
"crypto/sha256"
"encoding/json"
"fmt"
"go.uber.org/zap"
"telesrv/internal/domain"
)
const uploadedMediaIntentVersion = 1
type uploadedMediaIntent struct {
Version int `json:"version"`
Kind domain.UploadedMediaKind `json:"kind"`
File domain.UploadedFileRef `json:"file"`
Spec *domain.DocumentSpec `json:"spec,omitempty"`
}
func uploadedMediaIntentHash(kind domain.UploadedMediaKind, file domain.UploadedFileRef, spec *domain.DocumentSpec) ([]byte, error) {
payload, err := json.Marshal(uploadedMediaIntent{
Version: uploadedMediaIntentVersion,
Kind: kind,
File: file,
Spec: spec,
})
if err != nil {
return nil, fmt.Errorf("marshal uploaded media intent: %w", err)
}
sum := sha256.Sum256(payload)
return sum[:], nil
}
func sameUploadedMediaReceipt(receipt domain.UploadedMediaReceipt, kind domain.UploadedMediaKind, intentHash []byte) bool {
return receipt.Kind == kind && len(intentHash) == sha256.Size && bytes.Equal(receipt.IntentHash, intentHash)
}
func (s *Service) replayUploadedPhoto(ctx context.Context, file domain.UploadedFileRef, intentHash []byte) (domain.Photo, bool, error) {
receipt, found, err := s.media.GetUploadedMediaReceipt(ctx, file.OwnerUserID, file.FileID)
if err != nil || !found {
return domain.Photo{}, false, err
}
if !sameUploadedMediaReceipt(receipt, domain.UploadedMediaPhoto, intentHash) {
return domain.Photo{}, false, domain.ErrFilePartsInvalid
}
photo, found, err := s.media.GetPhoto(ctx, receipt.MediaID)
if err != nil {
return domain.Photo{}, false, err
}
if !found {
return domain.Photo{}, false, fmt.Errorf("uploaded photo receipt %d/%d references missing photo %d", file.OwnerUserID, file.FileID, receipt.MediaID)
}
s.cleanupMaterializedUpload(ctx, file, "photo replay")
return photo, true, nil
}
func (s *Service) replayUploadedDocument(ctx context.Context, file domain.UploadedFileRef, intentHash []byte) (domain.Document, bool, error) {
receipt, found, err := s.media.GetUploadedMediaReceipt(ctx, file.OwnerUserID, file.FileID)
if err != nil || !found {
return domain.Document{}, false, err
}
if !sameUploadedMediaReceipt(receipt, domain.UploadedMediaDocument, intentHash) {
return domain.Document{}, false, domain.ErrFilePartsInvalid
}
doc, found, err := s.media.GetDocument(ctx, receipt.MediaID)
if err != nil {
return domain.Document{}, false, err
}
if !found {
return domain.Document{}, false, fmt.Errorf("uploaded document receipt %d/%d references missing document %d", file.OwnerUserID, file.FileID, receipt.MediaID)
}
s.cleanupMaterializedUpload(ctx, file, "document replay")
return doc, true, nil
}
func (s *Service) commitUploadedMediaReceipt(ctx context.Context, file domain.UploadedFileRef, kind domain.UploadedMediaKind, intentHash []byte, mediaID int64) (domain.UploadedMediaReceipt, error) {
receipt, _, err := s.media.PutUploadedMediaReceipt(ctx, domain.UploadedMediaReceipt{
OwnerUserID: file.OwnerUserID,
FileID: file.FileID,
IntentHash: intentHash,
Kind: kind,
MediaID: mediaID,
})
if err != nil {
return domain.UploadedMediaReceipt{}, err
}
if !sameUploadedMediaReceipt(receipt, kind, intentHash) {
return domain.UploadedMediaReceipt{}, domain.ErrFilePartsInvalid
}
return receipt, nil
}
func (s *Service) cleanupMaterializedUpload(ctx context.Context, file domain.UploadedFileRef, reason string) {
if err := s.cleanupUploadParts(ctx, file.OwnerUserID, file.FileID); err != nil {
s.log.Warn("cleanup materialized upload parts failed",
zap.String("reason", reason),
zap.Int64("owner_user_id", file.OwnerUserID),
zap.Int64("file_id", file.FileID),
zap.Error(err),
)
}
}