144 lines
3.7 KiB
Go
144 lines
3.7 KiB
Go
//go:build !unix
|
|
|
|
package promptdebug
|
|
|
|
import (
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
)
|
|
|
|
type secureDirectory struct {
|
|
root *os.Root
|
|
}
|
|
|
|
func openSecureDirectory(path string) (*secureDirectory, error) {
|
|
if !filepath.IsAbs(path) {
|
|
return nil, fmt.Errorf("directory must be absolute")
|
|
}
|
|
if err := os.MkdirAll(path, debugDirectoryMode); err != nil {
|
|
return nil, err
|
|
}
|
|
info, err := os.Lstat(path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
|
|
return nil, fmt.Errorf("prompt debug root is not a directory")
|
|
}
|
|
if err := os.Chmod(path, debugDirectoryMode); err != nil {
|
|
return nil, err
|
|
}
|
|
root, err := os.OpenRoot(path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &secureDirectory{root: root}, nil
|
|
}
|
|
|
|
func (directory *secureDirectory) Close() error {
|
|
if directory == nil || directory.root == nil {
|
|
return nil
|
|
}
|
|
root := directory.root
|
|
directory.root = nil
|
|
return root.Close()
|
|
}
|
|
|
|
func (directory *secureDirectory) openDirectory(components ...string) (*secureDirectory, error) {
|
|
if directory == nil || directory.root == nil {
|
|
return nil, fmt.Errorf("prompt debug directory is closed")
|
|
}
|
|
for _, component := range components {
|
|
if err := validateSecureDirectoryName(component); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
path := filepath.Join(components...)
|
|
if err := directory.root.MkdirAll(path, debugDirectoryMode); err != nil {
|
|
return nil, err
|
|
}
|
|
child, err := directory.root.OpenRoot(path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &secureDirectory{root: child}, nil
|
|
}
|
|
|
|
func validateSecureDirectoryName(name string) error {
|
|
if name == "" || name == "." || name == ".." || strings.ContainsRune(name, filepath.Separator) {
|
|
return fmt.Errorf("prompt debug directory component %q is invalid", name)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (directory *secureDirectory) writeJSON(name string, value any) error {
|
|
if directory == nil || directory.root == nil {
|
|
return fmt.Errorf("prompt debug directory is closed")
|
|
}
|
|
if err := validateSecureDirectoryName(name); err != nil {
|
|
return err
|
|
}
|
|
data, err := json.MarshalIndent(value, "", " ")
|
|
if err != nil {
|
|
return fmt.Errorf("marshal prompt debug artifact: %w", err)
|
|
}
|
|
if info, err := directory.root.Lstat(name); err == nil {
|
|
if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
|
|
return fmt.Errorf("prompt debug file %q is not a regular file", name)
|
|
}
|
|
} else if !os.IsNotExist(err) {
|
|
return err
|
|
}
|
|
temporaryName, temporary, err := directory.createTemporaryFile(name)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer func() {
|
|
if temporary != nil {
|
|
_ = temporary.Close()
|
|
}
|
|
_ = directory.root.Remove(temporaryName)
|
|
}()
|
|
if err := temporary.Chmod(debugFileMode); err != nil {
|
|
return err
|
|
}
|
|
if written, err := temporary.Write(data); err != nil {
|
|
return err
|
|
} else if written != len(data) {
|
|
return io.ErrShortWrite
|
|
}
|
|
if err := temporary.Close(); err != nil {
|
|
return err
|
|
}
|
|
temporary = nil
|
|
if err := directory.root.Rename(temporaryName, name); err != nil {
|
|
return fmt.Errorf("replace prompt debug file %q: %w", name, err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (directory *secureDirectory) createTemporaryFile(name string) (string, *os.File, error) {
|
|
for attempt := 0; attempt < 16; attempt++ {
|
|
random := make([]byte, 12)
|
|
if _, err := rand.Read(random); err != nil {
|
|
return "", nil, err
|
|
}
|
|
temporaryName := "." + name + "." + hex.EncodeToString(random) + ".tmp"
|
|
temporary, err := directory.root.OpenFile(temporaryName, os.O_WRONLY|os.O_CREATE|os.O_EXCL, debugFileMode)
|
|
if os.IsExist(err) {
|
|
continue
|
|
}
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
return temporaryName, temporary, nil
|
|
}
|
|
return "", nil, fmt.Errorf("create temporary prompt debug file: too many name collisions")
|
|
}
|