From d8537ef319babc6d6583dee3186463ed357dbbd9 Mon Sep 17 00:00:00 2001 From: Zhang Rui Date: Thu, 12 Mar 2026 09:14:51 +0800 Subject: [PATCH] refactor(wecom): simplify image download and storage process in storeWSImage --- pkg/channels/wecom/aibot_ws.go | 110 +++++++++------------------------ pkg/media/store.go | 60 ------------------ 2 files changed, 30 insertions(+), 140 deletions(-) diff --git a/pkg/channels/wecom/aibot_ws.go b/pkg/channels/wecom/aibot_ws.go index eddea1cb4..03cdea6d4 100644 --- a/pkg/channels/wecom/aibot_ws.go +++ b/pkg/channels/wecom/aibot_ws.go @@ -1,7 +1,6 @@ package wecom import ( - "bytes" "context" "encoding/base64" "encoding/json" @@ -1045,10 +1044,7 @@ func wsGenerateID() string { // ---- Inbound image download helpers ---- -// storeWSImage downloads, optionally decrypts, and stores an inbound image. -// It streams the HTTP response body directly into the MediaStore, avoiding any -// caller-managed temp files. File lifecycle (creation and deletion) is owned -// entirely by the store and cleaned up via store.ReleaseAll. +// storeWSImage downloads the image at imageURL (with optional AES-CBC decryption) and stores it in the MediaStore. func (c *WeComAIBotWSChannel) storeWSImage( ctx context.Context, chatID, msgID, imageURL, aesKey string, @@ -1073,25 +1069,17 @@ func (c *WeComAIBotWSChannel) storeWSImage( return "", fmt.Errorf("download HTTP %d", resp.StatusCode) } - mediaDir := filepath.Join(os.TempDir(), "picoclaw_media") - scope := channels.BuildMediaScope("wecom_aibot", chatID, msgID) - meta := media.MediaMeta{Filename: msgID + ".jpg", Source: "wecom_aibot"} - - // lr wraps the response body with a +1 limit: if lr.N reaches 0 after the - // transfer, the body was at least maxSize+1 bytes and must be rejected. - lr := &io.LimitedReader{R: resp.Body, N: int64(maxSize) + 1} - var r io.Reader = lr + // Buffer the image in memory, bounded to maxSize. + data, err := io.ReadAll(io.LimitReader(resp.Body, int64(maxSize)+1)) + if err != nil { + return "", fmt.Errorf("read image: %w", err) + } + if len(data) > maxSize { + return "", fmt.Errorf("image too large (> %d MB)", maxSize>>20) + } + // AES-CBC decryption if a key is present. if aesKey != "" { - // AES-CBC decryption requires full ciphertext. Buffer the bounded encrypted - // bytes, decrypt, then pass the plaintext as a reader. - encrypted, readErr := io.ReadAll(lr) - if readErr != nil { - return "", fmt.Errorf("read for decrypt: %w", readErr) - } - if lr.N == 0 { - return "", fmt.Errorf("image too large (> %d MB)", maxSize>>20) - } key, decErr := base64.StdEncoding.DecodeString(aesKey) if decErr != nil || len(key) != 32 { key, decErr = decodeWeComAESKey(aesKey) @@ -1099,79 +1087,41 @@ func (c *WeComAIBotWSChannel) storeWSImage( return "", fmt.Errorf("decode image AES key: %w", decErr) } } - decrypted, decErr := decryptAESCBC(key, encrypted) - if decErr != nil { - return "", fmt.Errorf("decrypt image: %w", decErr) + data, err = decryptAESCBC(key, data) + if err != nil { + return "", fmt.Errorf("decrypt image: %w", err) } - r = bytes.NewReader(decrypted) } - // Fast path: FileMediaStore supports StoreFromReader which manages the file - // lifecycle entirely (temp-rename-fsync pattern, deleted by ReleaseAll). - if fsStore, ok := store.(*media.FileMediaStore); ok { - ref, storeErr := fsStore.StoreFromReader(r, meta, scope, mediaDir) - if storeErr != nil { - return "", storeErr - } - // For the no-AES path, check whether the response body hit the size limit. - if aesKey == "" && lr.N == 0 { - _ = store.ReleaseAll(scope) // remove the oversized file - return "", fmt.Errorf("image too large (> %d MB)", maxSize>>20) - } - return ref, nil + // Write to a temp file. The file is owned by the MediaStore and deleted by + // store.ReleaseAll — no caller-side cleanup needed. + mediaDir := filepath.Join(os.TempDir(), "picoclaw_media") + if err = os.MkdirAll(mediaDir, 0o700); err != nil { + return "", fmt.Errorf("mkdir: %w", err) } - - // Fallback path for non-FileMediaStore implementations: stream into a temp - // file using the temp-rename-fsync pattern, then register the path. - // The file lifecycle is managed by store.ReleaseAll(). - if mkErr := os.MkdirAll(mediaDir, 0o700); mkErr != nil { - return "", fmt.Errorf("mkdir: %w", mkErr) - } - tmpPath := filepath.Join(mediaDir, ".tmp-"+wsGenerateID()) - tmpFile, err := os.OpenFile(tmpPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + tmpFile, err := os.CreateTemp(mediaDir, msgID+"-*.jpg") if err != nil { return "", fmt.Errorf("create temp file: %w", err) } - cleanup := true - defer func() { - if cleanup { - tmpFile.Close() - os.Remove(tmpPath) - } - }() - n, cpErr := io.Copy(tmpFile, r) - syncErr := tmpFile.Sync() + tmpPath := tmpFile.Name() + _, writeErr := tmpFile.Write(data) closeErr := tmpFile.Close() - cleanup = false - if cpErr != nil { + if writeErr != nil { os.Remove(tmpPath) - return "", fmt.Errorf("write image: %w", cpErr) - } - if syncErr != nil { - os.Remove(tmpPath) - return "", fmt.Errorf("sync image: %w", syncErr) + return "", fmt.Errorf("write image: %w", writeErr) } if closeErr != nil { os.Remove(tmpPath) return "", fmt.Errorf("close image: %w", closeErr) } - if aesKey == "" && lr.N == 0 { - os.Remove(tmpPath) - return "", fmt.Errorf("image too large (> %d MB)", maxSize>>20) - } - _ = n - finalPath := filepath.Join(mediaDir, wsGenerateID()+".jpg") - if err = os.Rename(tmpPath, finalPath); err != nil { - os.Remove(tmpPath) - return "", fmt.Errorf("rename: %w", err) - } - if d, openErr := os.Open(mediaDir); openErr == nil { - _ = d.Sync() - d.Close() - } - ref, err := store.Store(finalPath, meta, scope) + + scope := channels.BuildMediaScope("wecom_aibot", chatID, msgID) + ref, err := store.Store(tmpPath, media.MediaMeta{ + Filename: msgID + ".jpg", + Source: "wecom_aibot", + }, scope) if err != nil { - os.Remove(finalPath) + os.Remove(tmpPath) return "", fmt.Errorf("store: %w", err) } return ref, nil diff --git a/pkg/media/store.go b/pkg/media/store.go index 7eb3cbf72..30220986c 100644 --- a/pkg/media/store.go +++ b/pkg/media/store.go @@ -2,9 +2,7 @@ package media import ( "fmt" - "io" "os" - "path/filepath" "sync" "time" @@ -110,64 +108,6 @@ func (s *FileMediaStore) Store(localPath string, meta MediaMeta, scope string) ( return ref, nil } -// StoreFromReader creates a temp file in dir by streaming r into it, then -// atomically renames it to a UUID-based final filename and registers it under -// scope. The file is owned by the store and deleted by ReleaseAll. -// This uses the same temp-write + sync + rename pattern as fileutil.WriteFileAtomic -// to guarantee flash-storage durability. -func (s *FileMediaStore) StoreFromReader(r io.Reader, meta MediaMeta, scope, dir string) (string, error) { - if err := os.MkdirAll(dir, 0o755); err != nil { - return "", fmt.Errorf("media store: mkdir: %w", err) - } - tmpPath := filepath.Join(dir, ".tmp-picoclaw-"+uuid.New().String()) - tmpFile, err := os.OpenFile(tmpPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) - if err != nil { - return "", fmt.Errorf("media store: create temp: %w", err) - } - cleanup := true - defer func() { - if cleanup { - tmpFile.Close() - os.Remove(tmpPath) - } - }() - - if _, err = io.Copy(tmpFile, r); err != nil { - return "", fmt.Errorf("media store: write: %w", err) - } - if err = tmpFile.Sync(); err != nil { - return "", fmt.Errorf("media store: sync: %w", err) - } - if err = tmpFile.Close(); err != nil { - return "", fmt.Errorf("media store: close: %w", err) - } - cleanup = false - - // Atomic rename to a UUID-based final name (preserving the original extension). - ext := filepath.Ext(meta.Filename) - finalPath := filepath.Join(dir, uuid.New().String()+ext) - if err = os.Rename(tmpPath, finalPath); err != nil { - os.Remove(tmpPath) - return "", fmt.Errorf("media store: rename: %w", err) - } - // Sync the directory so the rename is durable. - if d, openErr := os.Open(dir); openErr == nil { - _ = d.Sync() - d.Close() - } - - ref := "media://" + uuid.New().String() - s.mu.Lock() - defer s.mu.Unlock() - s.refs[ref] = mediaEntry{path: finalPath, meta: meta, storedAt: s.nowFunc()} - if s.scopeToRefs[scope] == nil { - s.scopeToRefs[scope] = make(map[string]struct{}) - } - s.scopeToRefs[scope][ref] = struct{}{} - s.refToScope[ref] = scope - return ref, nil -} - // Resolve returns the local path for the given ref. func (s *FileMediaStore) Resolve(ref string) (string, error) { s.mu.RLock()