refactor(wecom): simplify image download and storage process in storeWSImage

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.
+// storeWSImage downloads the image at imageURL (with optional AES-CBC decryption) and stores it in the MediaStore.
-// 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.
 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")
+	// Buffer the image in memory, bounded to maxSize.
-	scope := channels.BuildMediaScope("wecom_aibot", chatID, msgID)
+	data, err := io.ReadAll(io.LimitReader(resp.Body, int64(maxSize)+1))
-	meta := media.MediaMeta{Filename: msgID + ".jpg", Source: "wecom_aibot"}
+	if err != nil {
-
+		return "", fmt.Errorf("read image: %w", err)
-	// 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
-
-	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 {
+	if len(data) > maxSize {
 		return "", fmt.Errorf("image too large (> %d MB)", maxSize>>20)
 	}
+
+	// AES-CBC decryption if a key is present.
+	if aesKey != "" {
 		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)
+		data, err = decryptAESCBC(key, data)
-		if decErr != nil {
+		if err != nil {
-			return "", fmt.Errorf("decrypt image: %w", decErr)
+			return "", fmt.Errorf("decrypt image: %w", err)
 		}
-		r = bytes.NewReader(decrypted)
 	}
 
-	// Fast path: FileMediaStore supports StoreFromReader which manages the file
+	// Write to a temp file. The file is owned by the MediaStore and deleted by
-	// lifecycle entirely (temp-rename-fsync pattern, deleted by ReleaseAll).
+	// store.ReleaseAll — no caller-side cleanup needed.
-	if fsStore, ok := store.(*media.FileMediaStore); ok {
+	mediaDir := filepath.Join(os.TempDir(), "picoclaw_media")
-		ref, storeErr := fsStore.StoreFromReader(r, meta, scope, mediaDir)
+	if err = os.MkdirAll(mediaDir, 0o700); err != nil {
-		if storeErr != nil {
+		return "", fmt.Errorf("mkdir: %w", err)
-			return "", storeErr
 	}
-		// For the no-AES path, check whether the response body hit the size limit.
+	tmpFile, err := os.CreateTemp(mediaDir, msgID+"-*.jpg")
-		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
-	}
-
-	// 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)
 	if err != nil {
 		return "", fmt.Errorf("create temp file: %w", err)
 	}
-	cleanup := true
+	tmpPath := tmpFile.Name()
-	defer func() {
+	_, writeErr := tmpFile.Write(data)
-		if cleanup {
-			tmpFile.Close()
-			os.Remove(tmpPath)
-		}
-	}()
-	n, cpErr := io.Copy(tmpFile, r)
-	syncErr := tmpFile.Sync()
 	closeErr := tmpFile.Close()
-	cleanup = false
+	if writeErr != nil {
-	if cpErr != nil {
 		os.Remove(tmpPath)
-		return "", fmt.Errorf("write image: %w", cpErr)
+		return "", fmt.Errorf("write image: %w", writeErr)
-	}
-	if syncErr != nil {
-		os.Remove(tmpPath)
-		return "", fmt.Errorf("sync image: %w", syncErr)
 	}
 	if closeErr != nil {
 		os.Remove(tmpPath)
 		return "", fmt.Errorf("close image: %w", closeErr)
 	}
-	if aesKey == "" && lr.N == 0 {
+
-		os.Remove(tmpPath)
+	scope := channels.BuildMediaScope("wecom_aibot", chatID, msgID)
-		return "", fmt.Errorf("image too large (> %d MB)", maxSize>>20)
+	ref, err := store.Store(tmpPath, media.MediaMeta{
-	}
+		Filename: msgID + ".jpg",
-	_ = n
+		Source:   "wecom_aibot",
-	finalPath := filepath.Join(mediaDir, wsGenerateID()+".jpg")
+	}, scope)
-	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)
 	if err != nil {
-		os.Remove(finalPath)
+		os.Remove(tmpPath)
 		return "", fmt.Errorf("store: %w", err)
 	}
 	return ref, nil

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()