Improve error handling and validation in acquisition module also handle

apk verification in memory instead of using temporary files

Author: besendorf

acquisition/acquisition.go

@@ -103,6 +103,7 @@ func New(path string) (*Acquisition, error) {
 
 		// Initialize streaming puller for direct operations
 		acq.StreamingPuller = NewStreamingPuller(adb.Client.ExePath, adb.Client.Serial, 100) // 100MB max memory
+
 	}
 
 	return &acq, nil
@@ -264,37 +265,51 @@ func (a *Acquisition) StoreInfo() error {
 
 // StreamAPKToZip streams an APK file directly to encrypted zip with certificate processing
 func (a *Acquisition) StreamAPKToZip(remotePath, zipPath string, processFunc func(io.Reader) error) error {
-	if !a.StreamingMode || a.EncryptedWriter == nil {
-		return fmt.Errorf("streaming mode not available")
+	if err := a.validateStreamingMode(); err != nil {
+		return err
+	}
+
+	if remotePath == "" {
+		return fmt.Errorf("remote path cannot be empty")
+	}
+	if zipPath == "" {
+		return fmt.Errorf("zip path cannot be empty")
 	}
 
 	// Pull APK data to memory buffer
 	buffer, err := a.StreamingPuller.PullToBuffer(remotePath)
 	if err != nil {
-		return fmt.Errorf("failed to pull APK %s: %v", remotePath, err)
+		return fmt.Errorf("failed to pull APK %q: %v", remotePath, err)
 	}
 
 	// Process APK if processor provided (e.g., certificate verification)
 	if processFunc != nil {
 		err = processFunc(buffer.Reader())
 		if err != nil {
-			return fmt.Errorf("failed to process APK %s: %v", remotePath, err)
+			return fmt.Errorf("failed to process APK %q: %v", remotePath, err)
 		}
 	}
 
 	// Stream to encrypted zip
 	err = a.EncryptedWriter.CreateFileFromReader(zipPath, buffer.Reader())
 	if err != nil {
-		return fmt.Errorf("failed to add APK %s to encrypted zip: %v", remotePath, err)
+		return fmt.Errorf("failed to add APK %q to encrypted zip: %v", remotePath, err)
 	}
 
 	return nil
 }
 
 // StreamBackupToZip streams a backup directly to encrypted zip
 func (a *Acquisition) StreamBackupToZip(arg, zipPath string) error {
-	if !a.StreamingMode || a.EncryptedWriter == nil {
-		return fmt.Errorf("streaming mode not available")
+	if err := a.validateStreamingMode(); err != nil {
+		return err
+	}
+
+	if arg == "" {
+		return fmt.Errorf("backup argument cannot be empty")
+	}
+	if zipPath == "" {
+		return fmt.Errorf("zip path cannot be empty")
 	}
 
 	// Create zip entry writer
@@ -306,16 +321,20 @@ func (a *Acquisition) StreamBackupToZip(arg, zipPath string) error {
 	// Stream backup directly to zip
 	err = a.StreamingPuller.BackupToWriter(arg, writer)
 	if err != nil {
-		return fmt.Errorf("failed to stream backup to zip: %v", err)
+		return fmt.Errorf("failed to stream backup %q to zip: %v", arg, err)
 	}
 
 	return nil
 }
 
 // StreamBugreportToZip streams a bugreport directly to encrypted zip
 func (a *Acquisition) StreamBugreportToZip(zipPath string) error {
-	if !a.StreamingMode || a.EncryptedWriter == nil {
-		return fmt.Errorf("streaming mode not available")
+	if err := a.validateStreamingMode(); err != nil {
+		return err
+	}
+
+	if zipPath == "" {
+		return fmt.Errorf("zip path cannot be empty")
 	}
 
 	// Create zip entry writer
@@ -332,3 +351,17 @@ func (a *Acquisition) StreamBugreportToZip(zipPath string) error {
 
 	return nil
 }
+
+// validateStreamingMode checks if streaming mode is properly initialized
+func (a *Acquisition) validateStreamingMode() error {
+	if !a.StreamingMode {
+		return fmt.Errorf("streaming mode not enabled")
+	}
+	if a.EncryptedWriter == nil {
+		return fmt.Errorf("encrypted writer not initialized")
+	}
+	if a.StreamingPuller == nil {
+		return fmt.Errorf("streaming puller not initialized")
+	}
+	return nil
+}

acquisition/encrypted_stream.go

@@ -7,6 +7,7 @@ package acquisition
 
 import (
 	"archive/zip"
+	"bytes"
 	"fmt"
 	"io"
 	"os"
@@ -65,6 +66,7 @@ func NewEncryptedZipWriter(uuid string) (*EncryptedZipWriter, error) {
 	encWriter, err := age.Encrypt(file, recipient)
 	if err != nil {
 		file.Close()
+		os.Remove(outputPath) // Clean up the created file
 		return nil, fmt.Errorf("failed to create encrypted writer: %v", err)
 	}
 
@@ -84,8 +86,12 @@ func NewEncryptedZipWriter(uuid string) (*EncryptedZipWriter, error) {
 
 // CreateFile creates a new file in the encrypted zip and returns a writer
 func (ezw *EncryptedZipWriter) CreateFile(name string) (io.Writer, error) {
-	if ezw.closed {
-		return nil, fmt.Errorf("encrypted zip writer is closed")
+	if err := ezw.checkClosed(); err != nil {
+		return nil, err
+	}
+
+	if name == "" {
+		return nil, fmt.Errorf("file name cannot be empty")
 	}
 
 	header := &zip.FileHeader{
@@ -99,8 +105,8 @@ func (ezw *EncryptedZipWriter) CreateFile(name string) (io.Writer, error) {
 
 // CreateFileFromReader copies data from a reader to a file in the encrypted zip
 func (ezw *EncryptedZipWriter) CreateFileFromReader(name string, src io.Reader) error {
-	if ezw.closed {
-		return fmt.Errorf("encrypted zip writer is closed")
+	if src == nil {
+		return fmt.Errorf("source reader cannot be nil")
 	}
 
 	writer, err := ezw.CreateFile(name)
@@ -118,112 +124,57 @@ func (ezw *EncryptedZipWriter) CreateFileFromReader(name string, src io.Reader)
 
 // CreateFileFromString creates a file with string content in the encrypted zip
 func (ezw *EncryptedZipWriter) CreateFileFromString(name, content string) error {
-	if ezw.closed {
-		return fmt.Errorf("encrypted zip writer is closed")
-	}
-
-	writer, err := ezw.CreateFile(name)
-	if err != nil {
-		return fmt.Errorf("failed to create file in zip: %v", err)
-	}
-
-	_, err = writer.Write([]byte(content))
-	if err != nil {
-		return fmt.Errorf("failed to write content to zip file: %v", err)
-	}
-
-	return nil
+	return ezw.CreateFileFromReader(name, strings.NewReader(content))
 }
 
 // CreateFileFromBytes creates a file with byte content in the encrypted zip
 func (ezw *EncryptedZipWriter) CreateFileFromBytes(name string, content []byte) error {
-	if ezw.closed {
-		return fmt.Errorf("encrypted zip writer is closed")
-	}
-
-	writer, err := ezw.CreateFile(name)
-	if err != nil {
-		return fmt.Errorf("failed to create file in zip: %v", err)
-	}
-
-	_, err = writer.Write(content)
-	if err != nil {
-		return fmt.Errorf("failed to write content to zip file: %v", err)
-	}
-
-	return nil
+	return ezw.CreateFileFromReader(name, bytes.NewReader(content))
 }
 
 // CreateFileFromPath reads a file from disk and adds it to the encrypted zip
 func (ezw *EncryptedZipWriter) CreateFileFromPath(name, filePath string) error {
-	if ezw.closed {
-		return fmt.Errorf("encrypted zip writer is closed")
-	}
-
 	file, err := os.Open(filePath)
 	if err != nil {
-		return fmt.Errorf("failed to open source file: %v", err)
+		return fmt.Errorf("failed to open source file %q: %v", filePath, err)
 	}
 	defer file.Close()
 
 	return ezw.CreateFileFromReader(name, file)
 }
 
-// AddDirectory recursively adds all files from a directory to the encrypted zip
-func (ezw *EncryptedZipWriter) AddDirectory(dirPath, basePath string) error {
-	if ezw.closed {
-		return fmt.Errorf("encrypted zip writer is closed")
-	}
-
-	return filepath.Walk(dirPath, func(path string, info os.FileInfo, err error) error {
-		if err != nil {
-			return err
-		}
-
-		// Skip directories
-		if info.IsDir() {
-			return nil
-		}
-
-		// Calculate relative path for zip entry
-		relPath, err := filepath.Rel(basePath, path)
-		if err != nil {
-			return fmt.Errorf("failed to get relative path: %v", err)
-		}
-
-		// Use forward slashes for zip paths
-		zipPath := filepath.ToSlash(relPath)
-
-		// Add file to zip
-		return ezw.CreateFileFromPath(zipPath, path)
-	})
-}
-
 // Close finalizes and closes the encrypted zip
 func (ezw *EncryptedZipWriter) Close() error {
 	if ezw.closed {
 		return nil
 	}
 
 	ezw.closed = true
+	var lastErr error
 
 	// Close zip writer first
 	if err := ezw.zipWriter.Close(); err != nil {
-		return fmt.Errorf("failed to close zip writer: %v", err)
+		lastErr = fmt.Errorf("failed to close zip writer: %v", err)
 	}
 
 	// Close encryption writer
 	if err := ezw.encWriter.Close(); err != nil {
-		return fmt.Errorf("failed to close encryption writer: %v", err)
+		if lastErr == nil {
+			lastErr = fmt.Errorf("failed to close encryption writer: %v", err)
+		}
 	}
 
 	// Close file
 	if err := ezw.file.Close(); err != nil {
-		return fmt.Errorf("failed to close output file: %v", err)
+		if lastErr == nil {
+			lastErr = fmt.Errorf("failed to close output file: %v", err)
+		}
 	}
 
-	log.Infof("Encrypted archive created successfully at %s", ezw.outputPath)
-	return nil
+	if lastErr == nil {
+		log.Infof("Encrypted archive created successfully at %s", ezw.outputPath)
+	}
+	return lastErr
 }
 
 // GetOutputPath returns the path to the encrypted zip file
@@ -235,3 +186,11 @@ func (ezw *EncryptedZipWriter) GetOutputPath() string {
 func (ezw *EncryptedZipWriter) IsClosed() bool {
 	return ezw.closed
 }
+
+// checkClosed is a helper method to check if the writer is closed
+func (ezw *EncryptedZipWriter) checkClosed() error {
+	if ezw.closed {
+		return fmt.Errorf("encrypted zip writer is closed")
+	}
+	return nil
+}