Secure prompt debug filesystem writes

internal/promptdebug/debug_writer.go

@@ -6,7 +6,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/url"
-	"os"
 	"path/filepath"
 	"strings"
 	"time"
@@ -25,7 +24,8 @@ const (
 // PromptDebugWriter stores explicitly requested content-rich diagnostics. A
 // writer created without a root is disabled.
 type PromptDebugWriter struct {
-	root string
+	root      string
+	directory *secureDirectory
 }
 
 // PromptDebugRef identifies one debug capture directory.
@@ -140,14 +140,27 @@ func NewPromptDebugWriter(root string) (*PromptDebugWriter, error) {
 	if cleaned == string(filepath.Separator) {
 		return nil, fmt.Errorf("prompt debug root must not be the filesystem root")
 	}
-	if err := ensureSecureDirectory(cleaned); err != nil {
+	directory, err := openSecureDirectory(cleaned)
+	if err != nil {
 		return nil, fmt.Errorf("initialize prompt debug root %q: %w", cleaned, err)
 	}
-	return &PromptDebugWriter{root: cleaned}, nil
+	return &PromptDebugWriter{root: cleaned, directory: directory}, nil
 }
 
 func (w *PromptDebugWriter) Enabled() bool {
-	return w != nil && w.root != ""
+	return w != nil && w.root != "" && w.directory != nil
+}
+
+// Close releases the secure directory handle retained for an enabled writer.
+// It is safe to call on disabled writers.
+func (w *PromptDebugWriter) Close() error {
+	if w == nil || w.directory == nil {
+		return nil
+	}
+	directory := w.directory
+	w.directory = nil
+	w.root = ""
+	return directory.Close()
 }
 
 // WritePreparation stores the explicitly captured preparation details and
@@ -156,10 +169,11 @@ func (w *PromptDebugWriter) WritePreparation(ref PromptDebugRef, preparation pro
 	if !w.Enabled() {
 		return "", nil
 	}
-	directory, err := w.runDirectory(ref)
+	directory, secureDirectory, err := w.runDirectory(ref)
 	if err != nil {
 		return "", err
 	}
+	defer secureDirectory.Close()
 	artifact := PromptPreparationDebugArtifact{
 		SchemaVersion: promptPreparationDebugSchemaVersion,
 		ReportID:      ref.ReportID,
@@ -177,7 +191,7 @@ func (w *PromptDebugWriter) WritePreparation(ref PromptDebugRef, preparation pro
 		}
 		artifact.Parameters = parameters
 	}
-	if err := writeSecureJSON(filepath.Join(directory, "preparation.json"), artifact); err != nil {
+	if err := secureDirectory.writeJSON("preparation.json", artifact); err != nil {
 		return "", err
 	}
 	return directory, nil
@@ -189,10 +203,11 @@ func (w *PromptDebugWriter) WriteExecution(ref PromptDebugRef, execution prompte
 	if !w.Enabled() {
 		return "", nil
 	}
-	directory, err := w.runDirectory(ref)
+	directory, secureDirectory, err := w.runDirectory(ref)
 	if err != nil {
 		return "", err
 	}
+	defer secureDirectory.Close()
 	artifact := PromptExecutionDebugArtifact{
 		SchemaVersion: promptExecutionDebugSchemaVersion,
 		ReportID:      ref.ReportID,
@@ -208,27 +223,25 @@ func (w *PromptDebugWriter) WriteExecution(ref PromptDebugRef, execution prompte
 		}
 		artifact.DebugValidationDetails = append([]string(nil), execution.Debug.ValidationDiagnostics...)
 	}
-	if err := writeSecureJSON(filepath.Join(directory, "execution.json"), artifact); err != nil {
+	if err := secureDirectory.writeJSON("execution.json", artifact); err != nil {
 		return "", err
 	}
 	return directory, nil
 }
 
-func (w *PromptDebugWriter) runDirectory(ref PromptDebugRef) (string, error) {
+func (w *PromptDebugWriter) runDirectory(ref PromptDebugRef) (string, *secureDirectory, error) {
 	if err := validatePromptDebugRef(ref); err != nil {
-		return "", err
-	}
-	if err := ensureSecureDirectory(w.root); err != nil {
-		return "", fmt.Errorf("validate prompt debug root %q: %w", w.root, err)
+		return "", nil, err
 	}
 	directory := filepath.Join(w.root, string(ref.ReportID), ref.ValidDate, ref.RunID)
 	if !isWithinDirectory(w.root, directory) {
-		return "", fmt.Errorf("prompt debug path escapes root")
+		return "", nil, fmt.Errorf("prompt debug path escapes root")
 	}
-	if err := ensureSecureDirectory(directory); err != nil {
-		return "", fmt.Errorf("create prompt debug directory %q: %w", directory, err)
+	secureDirectory, err := w.directory.openDirectory(string(ref.ReportID), ref.ValidDate, ref.RunID)
+	if err != nil {
+		return "", nil, fmt.Errorf("create prompt debug directory %q: %w", directory, err)
 	}
-	return directory, nil
+	return directory, secureDirectory, nil
 }
 
 func validatePromptDebugRef(ref PromptDebugRef) error {
@@ -254,100 +267,11 @@ func validatePromptDebugSegment(name string, value string) error {
 	return nil
 }
 
-func ensureSecureDirectory(path string) error {
-	if !filepath.IsAbs(path) {
-		return fmt.Errorf("directory must be absolute")
-	}
-	cleaned := filepath.Clean(path)
-	volume := filepath.VolumeName(cleaned)
-	current := volume + string(filepath.Separator)
-	for _, component := range strings.Split(strings.TrimPrefix(cleaned, current), string(filepath.Separator)) {
-		if component == "" {
-			continue
-		}
-		current = filepath.Join(current, component)
-		info, err := os.Lstat(current)
-		if os.IsNotExist(err) {
-			if err := os.Mkdir(current, debugDirectoryMode); err != nil {
-				if !os.IsExist(err) {
-					return err
-				}
-				info, err = os.Lstat(current)
-				if err != nil {
-					return err
-				}
-				if info.Mode()&os.ModeSymlink != 0 {
-					return fmt.Errorf("directory component %q must not be a symlink", current)
-				}
-				if !info.IsDir() {
-					return fmt.Errorf("directory component %q is not a directory", current)
-				}
-				if err := os.Chmod(current, debugDirectoryMode); err != nil {
-					return err
-				}
-				continue
-			}
-			if err := os.Chmod(current, debugDirectoryMode); err != nil {
-				return err
-			}
-			continue
-		}
-		if err != nil {
-			return err
-		}
-		if info.Mode()&os.ModeSymlink != 0 {
-			return fmt.Errorf("directory component %q must not be a symlink", current)
-		}
-		if !info.IsDir() {
-			return fmt.Errorf("directory component %q is not a directory", current)
-		}
-	}
-	if err := os.Chmod(cleaned, debugDirectoryMode); err != nil {
-		return err
-	}
-	return nil
-}
-
 func isWithinDirectory(root string, path string) bool {
 	relative, err := filepath.Rel(root, path)
 	return err == nil && relative != ".." && !strings.HasPrefix(relative, ".."+string(filepath.Separator)) && !filepath.IsAbs(relative)
 }
 
-func writeSecureJSON(path string, value any) error {
-	data, err := json.MarshalIndent(value, "", "  ")
-	if err != nil {
-		return fmt.Errorf("marshal %q: %w", path, err)
-	}
-	if info, err := os.Lstat(path); err == nil {
-		if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
-			return fmt.Errorf("prompt debug file %q is not a regular file", path)
-		}
-	} else if !os.IsNotExist(err) {
-		return err
-	}
-	temporary, err := os.CreateTemp(filepath.Dir(path), "."+filepath.Base(path)+".*.tmp")
-	if err != nil {
-		return err
-	}
-	temporaryPath := temporary.Name()
-	defer os.Remove(temporaryPath)
-	if err := temporary.Chmod(debugFileMode); err != nil {
-		temporary.Close()
-		return err
-	}
-	if _, err := temporary.Write(data); err != nil {
-		temporary.Close()
-		return err
-	}
-	if err := temporary.Close(); err != nil {
-		return err
-	}
-	if err := os.Rename(temporaryPath, path); err != nil {
-		return err
-	}
-	return nil
-}
-
 func promptDebugPreparation(value promptexec.Preparation) PromptDebugPreparation {
 	return PromptDebugPreparation{
 		PromptID: value.PromptID, PromptVersion: value.PromptVersion, PromptHash: value.PromptHash,