convert to stream

Author: NullLatency

internal/storage/backend.go

@@ -2,6 +2,7 @@ package storage
 
 import (
 	"context"
+	"io"
 )
 
 // Backend defines the interface for our pluggable storage mechanism that acts as the
@@ -11,15 +12,15 @@ type Backend interface {
 	Login(ctx context.Context) error
 
 	// Upload writes a new file to the storage backend.
-	// filename is typically of the format request-<session>-<seq>-<timestamp>.json
-	Upload(ctx context.Context, filename string, data []byte) error
+	// filename is typically of the format request-<session>-<seq>-<timestamp>.bin
+	Upload(ctx context.Context, filename string, data io.Reader) error
 
 	// ListQuery searches the backend for files matching a specific prefix or criteria.
 	// We use this to discover new request or response payloads.
 	ListQuery(ctx context.Context, prefix string) ([]string, error)
 
-	// Download reads the file content from the backend.
-	Download(ctx context.Context, filename string) ([]byte, error)
+	// Download returns an io.ReadCloser for the file content from the backend.
+	Download(ctx context.Context, filename string) (io.ReadCloser, error)
 
 	// Delete removes a file from the backend after it has been read or expired.
 	Delete(ctx context.Context, filename string) error

internal/storage/google.go

@@ -217,35 +217,39 @@ func (b *GoogleBackend) getValidToken(ctx context.Context) (string, error) {
 	return b.token, nil
 }
 
-func (b *GoogleBackend) Upload(ctx context.Context, filename string, data []byte) error {
+func (b *GoogleBackend) Upload(ctx context.Context, filename string, data io.Reader) error {
 	tok, err := b.getValidToken(ctx)
 	if err != nil {
 		return err
 	}
 
-	var metaBuf bytes.Buffer
-	metaWriter := multipart.NewWriter(&metaBuf)
+	pr, pw := io.Pipe()
+	metaWriter := multipart.NewWriter(pw)
 
-	// Part 1: Metadata
-	h := make(textproto.MIMEHeader)
-	h.Set("Content-Type", "application/json; charset=UTF-8")
-	part1, _ := metaWriter.CreatePart(h)
-	meta := map[string]interface{}{
-		"name": filename,
-	}
-	if b.folderID != "" {
-		meta["parents"] = []string{b.folderID}
-	}
-	json.NewEncoder(part1).Encode(meta)
+	go func() {
+		defer pw.Close()
+		defer metaWriter.Close()
+
+		// Part 1: Metadata
+		h := make(textproto.MIMEHeader)
+		h.Set("Content-Type", "application/json; charset=UTF-8")
+		part1, _ := metaWriter.CreatePart(h)
+		meta := map[string]interface{}{
+			"name": filename,
+		}
+		if b.folderID != "" {
+			meta["parents"] = []string{b.folderID}
+		}
+		json.NewEncoder(part1).Encode(meta)
 
-	// Part 2: Content
-	h = make(textproto.MIMEHeader)
-	h.Set("Content-Type", "application/octet-stream")
-	part2, _ := metaWriter.CreatePart(h)
-	part2.Write(data)
-	metaWriter.Close()
+		// Part 2: Content
+		h = make(textproto.MIMEHeader)
+		h.Set("Content-Type", "application/octet-stream")
+		part2, _ := metaWriter.CreatePart(h)
+		io.Copy(part2, data)
+	}()
 
-	req, err := http.NewRequestWithContext(ctx, "POST", "https://www.googleapis.com/upload/drive/v3/files?uploadType=multipart", &metaBuf)
+	req, err := http.NewRequestWithContext(ctx, "POST", "https://www.googleapis.com/upload/drive/v3/files?uploadType=multipart", pr)
 	if err != nil {
 		return err
 	}
@@ -310,6 +314,11 @@ func (b *GoogleBackend) ListQuery(ctx context.Context, prefix string) ([]string,
 	}
 
 	b.fileIdsMu.Lock()
+	// SAFETY: Prevent fileIDs map from infinite growth
+	if len(b.fileIDs) > 2000 {
+		b.fileIDs = make(map[string]string)
+	}
+
 	var names []string
 	for _, f := range resData.Files {
 		// Only collect exact prefix matches client-side just in case
@@ -323,7 +332,7 @@ func (b *GoogleBackend) ListQuery(ctx context.Context, prefix string) ([]string,
 	return names, nil
 }
 
-func (b *GoogleBackend) Download(ctx context.Context, filename string) ([]byte, error) {
+func (b *GoogleBackend) Download(ctx context.Context, filename string) (io.ReadCloser, error) {
 	b.fileIdsMu.RLock()
 	fileID, ok := b.fileIDs[filename]
 	b.fileIdsMu.RUnlock()
@@ -347,14 +356,14 @@ func (b *GoogleBackend) Download(ctx context.Context, filename string) ([]byte,
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusOK {
 		body, _ := io.ReadAll(resp.Body)
+		resp.Body.Close()
 		return nil, fmt.Errorf("download returned %d: %s", resp.StatusCode, string(body))
 	}
 
-	return io.ReadAll(resp.Body)
+	return resp.Body, nil
 }
 
 func (b *GoogleBackend) Delete(ctx context.Context, filename string) error {

internal/storage/local.go

@@ -34,13 +34,22 @@ func (b *LocalBackend) Login(ctx context.Context) error {
 	return nil
 }
 
-func (b *LocalBackend) Upload(ctx context.Context, filename string, data []byte) error {
+func (b *LocalBackend) Upload(ctx context.Context, filename string, data io.Reader) error {
 	path := filepath.Join(b.baseDir, filename)
 	// Write to a temporary file first, then rename to avoid partial reads by the polling server
 	tmpPath := path + ".tmp"
-	if err := os.WriteFile(tmpPath, data, 0644); err != nil {
-		return fmt.Errorf("failed to write temp file: %w", err)
+	f, err := os.Create(tmpPath)
+	if err != nil {
+		return fmt.Errorf("failed to create temp file: %w", err)
+	}
+	
+	if _, err := io.Copy(f, data); err != nil {
+		f.Close()
+		os.Remove(tmpPath)
+		return fmt.Errorf("failed to write data: %w", err)
 	}
+	f.Close()
+
 	if err := os.Rename(tmpPath, path); err != nil {
 		os.Remove(tmpPath)
 		return fmt.Errorf("failed to rename temp file: %w", err)
@@ -66,7 +75,7 @@ func (b *LocalBackend) ListQuery(ctx context.Context, prefix string) ([]string,
 	return results, nil
 }
 
-func (b *LocalBackend) Download(ctx context.Context, filename string) ([]byte, error) {
+func (b *LocalBackend) Download(ctx context.Context, filename string) (io.ReadCloser, error) {
 	path := filepath.Join(b.baseDir, filename)
 	file, err := os.Open(path)
 	if err != nil {
@@ -76,13 +85,7 @@ func (b *LocalBackend) Download(ctx context.Context, filename string) ([]byte, e
 		}
 		return nil, fmt.Errorf("failed to open file: %w", err)
 	}
-	defer file.Close()
-
-	data, err := io.ReadAll(file)
-	if err != nil {
-		return nil, fmt.Errorf("failed to read file: %w", err)
-	}
-	return data, nil
+	return file, nil
 }
 
 func (b *LocalBackend) Delete(ctx context.Context, filename string) error {

internal/transport/conn.go

@@ -72,6 +72,7 @@ func (v *VirtualConn) Write(b []byte) (n int, err error) {
 func (v *VirtualConn) Close() error {
 	v.session.mu.Lock()
 	v.session.closed = true
+	v.session.txCond.Broadcast() // Wake up any writers blocked on backpressure
 	v.session.mu.Unlock()
 
 	// A closed connection no longer accepts writes efficiently