--- /dev/null
+package main
+
+import (
+ "archive/tar"
+ "compress/gzip"
+ "context"
+ "errors"
+ "fmt"
+ "io"
+ "io/fs"
+ "os"
+ "path/filepath"
+ "strings"
+)
+
+// backupArchive writes the files in root to path through a private temporary
+// file and a rename. Only directories and regular files enter the archive.
+func backupArchive(ctx context.Context, root, path string) (size int64, err error) {
+ file, err := os.CreateTemp(filepath.Dir(path), ".onechat-backup-*")
+ if err != nil {
+ return 0, fmt.Errorf("create archive: %w", err)
+ }
+
+ defer func() {
+ err = errors.Join(err, file.Close())
+ removeErr := os.Remove(file.Name())
+ if !errors.Is(removeErr, os.ErrNotExist) {
+ err = errors.Join(err, removeErr)
+ }
+ }()
+
+ compressed := gzip.NewWriter(file)
+ archive := tar.NewWriter(compressed)
+ err = filepath.WalkDir(root, func(path string, entry fs.DirEntry, walkErr error) error {
+ if walkErr != nil {
+ return fmt.Errorf("walk backup: %w", walkErr)
+ }
+
+ if ctx.Err() != nil {
+ return ctx.Err()
+ }
+
+ if path == root {
+ return nil
+ }
+
+ name, nameErr := filepath.Rel(root, path)
+ if nameErr != nil {
+ return fmt.Errorf("name archive entry: %w", nameErr)
+ }
+
+ info, infoErr := entry.Info()
+ if infoErr != nil {
+ return fmt.Errorf("read backup entry: %w", infoErr)
+ }
+
+ if !info.IsDir() && !info.Mode().IsRegular() {
+ return fmt.Errorf("unsupported backup entry %q", name)
+ }
+
+ header, headerErr := tar.FileInfoHeader(info, "")
+ if headerErr != nil {
+ return fmt.Errorf("make archive header: %w", headerErr)
+ }
+
+ header.Name = filepath.ToSlash(name)
+ header.Uid, header.Gid = 0, 0
+ header.Uname, header.Gname = "", ""
+ headerErr = archive.WriteHeader(header)
+ if headerErr != nil {
+ return fmt.Errorf("write archive header: %w", headerErr)
+ }
+
+ if info.IsDir() {
+ return nil
+ }
+
+ return backupCopyFile(ctx, archive, path, info.Size())
+ })
+ err = errors.Join(err, archive.Close(), compressed.Close())
+ if err != nil {
+ return 0, fmt.Errorf("write archive: %w", err)
+ }
+
+ err = file.Sync()
+ if err != nil {
+ return 0, fmt.Errorf("sync archive: %w", err)
+ }
+
+ info, err := file.Stat()
+ if err != nil {
+ return 0, fmt.Errorf("read archive size: %w", err)
+ }
+
+ err = os.Rename(file.Name(), path)
+ if err != nil {
+ return 0, fmt.Errorf("install archive: %w", err)
+ }
+
+ return info.Size(), nil
+}
+
+// backupExtract restores regular files and directories in an empty staging
+// directory. Links, duplicate entries, unsafe names, and special modes fail.
+// Directory modes apply last so a read-only directory can hold files.
+func backupExtract(ctx context.Context, path, root string) (err error) {
+ file, err := os.Open(path)
+ if err != nil {
+ return fmt.Errorf("open archive: %w", err)
+ }
+
+ defer func() {
+ err = errors.Join(err, file.Close())
+ }()
+
+ compressed, err := gzip.NewReader(file)
+ if err != nil {
+ return fmt.Errorf("open compressed archive: %w", err)
+ }
+
+ defer func() {
+ err = errors.Join(err, compressed.Close())
+ }()
+
+ archive := tar.NewReader(BackupReader{ctx: ctx, reader: compressed})
+ seen := make(map[string]bool)
+
+ var directories []*tar.Header
+
+ for {
+ header, readErr := archive.Next()
+ if errors.Is(readErr, io.EOF) {
+ break
+ }
+
+ if readErr != nil {
+ return fmt.Errorf("read archive entry: %w", readErr)
+ }
+
+ name := strings.TrimSuffix(header.Name, "/")
+ if !filepath.IsLocal(name) || strings.Contains(name, "\\") || filepath.Clean(name) != name || backupExcluded(name) {
+ return fmt.Errorf("unsafe archive path %q", header.Name)
+ }
+
+ if seen[name] {
+ return fmt.Errorf("duplicate archive path %q", name)
+ }
+
+ seen[name] = true
+ if header.Mode & ^int64(0o777) != 0 {
+ return fmt.Errorf("unsafe archive mode for %q", name)
+ }
+
+ target := filepath.Join(root, name)
+ switch header.Typeflag {
+ case tar.TypeDir:
+ err = os.MkdirAll(target, 0o700)
+ directories = append(directories, header)
+ case tar.TypeReg:
+ err = os.MkdirAll(filepath.Dir(target), 0o700)
+ if err == nil {
+ err = backupExtractFile(archive, target, os.FileMode(header.Mode))
+ }
+ default:
+ return fmt.Errorf("unsupported archive entry %q", name)
+ }
+
+ if err != nil {
+ return fmt.Errorf("restore archive entry %q: %w", name, err)
+ }
+ }
+
+ // Read through the gzip trailer to check its checksum.
+ _, err = io.Copy(io.Discard, compressed)
+ if err != nil {
+ return fmt.Errorf("check archive checksum: %w", err)
+ }
+
+ for i := len(directories) - 1; i >= 0; i-- {
+ header := directories[i]
+ err = os.Chmod(filepath.Join(root, header.Name), os.FileMode(header.Mode))
+ if err != nil {
+ return fmt.Errorf("restore directory mode: %w", err)
+ }
+ }
+
+ return nil
+}
+
+// BackupReader checks cancellation between reads of an archive or file.
+type BackupReader struct {
+ ctx context.Context
+ reader io.Reader
+}
+
+func (r BackupReader) Read(data []byte) (int, error) {
+ if r.ctx.Err() != nil {
+ return 0, r.ctx.Err()
+ }
+
+ count, err := r.reader.Read(data)
+ if err != nil {
+ return count, fmt.Errorf("read backup data: %w", err)
+ }
+
+ return count, nil
+}
+
+// backupCopyFile copies the size fixed at the snapshot of path. Growth after
+// that point is not part of the copy, and a shorter file is an error.
+func backupCopyFile(ctx context.Context, out io.Writer, path string, size int64) (err error) {
+ file, err := os.Open(path)
+ if err != nil {
+ return fmt.Errorf("open backup file: %w", err)
+ }
+
+ defer func() {
+ err = errors.Join(err, file.Close())
+ }()
+
+ _, err = io.CopyN(out, BackupReader{ctx: ctx, reader: file}, size)
+ if err != nil {
+ return fmt.Errorf("copy backup file: %w", err)
+ }
+
+ return nil
+}
+
+func backupExtractFile(in io.Reader, path string, mode os.FileMode) (err error) {
+ file, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
+ if err != nil {
+ return fmt.Errorf("create restored file: %w", err)
+ }
+
+ defer func() {
+ err = errors.Join(err, file.Close())
+ }()
+
+ _, err = io.Copy(file, in)
+ if err != nil {
+ return fmt.Errorf("write restored file: %w", err)
+ }
+
+ err = file.Chmod(mode)
+ if err != nil {
+ return fmt.Errorf("restore file mode: %w", err)
+ }
+
+ return nil
+}
+
+// backupExcluded keeps runtime locks, sockets, probe files, and import copies
+// out of the snapshot.
+func backupExcluded(name string) bool {
+ first, _, _ := strings.Cut(filepath.ToSlash(name), "/")
+ if first == socketFileName || first == tokenLockFileName || first == "probe" || strings.HasPrefix(first, "agent-backup-") {
+ return true
+ }
+
+ if filepath.Base(name) == "lock" {
+ return true
+ }
+
+ return false
+}
--- /dev/null
+package main
+
+import (
+ "archive/tar"
+ "compress/gzip"
+ "os"
+ "path/filepath"
+ "testing"
+
+ "github.com/stretchr/testify/require"
+)
+
+func TestBackupArchiveRestoresBytesAndModes(t *testing.T) {
+ root := t.TempDir()
+ require.NoError(t, os.Mkdir(filepath.Join(root, "agent"), 0o750))
+ require.NoError(t, os.WriteFile(filepath.Join(root, "agent", "instructions.md"), []byte("Keep the whole chat.\n"), 0o640))
+
+ archive := filepath.Join(t.TempDir(), "backup.tar.gz")
+ size, err := backupArchive(t.Context(), root, archive)
+ require.NoError(t, err)
+
+ info, err := os.Stat(archive)
+ require.NoError(t, err)
+ require.Equal(t, size, info.Size())
+ require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
+
+ restored := t.TempDir()
+ require.NoError(t, backupExtract(t.Context(), archive, restored))
+
+ data, err := os.ReadFile(filepath.Join(restored, "agent", "instructions.md"))
+ require.NoError(t, err)
+ require.Equal(t, "Keep the whole chat.\n", string(data))
+
+ info, err = os.Stat(filepath.Join(restored, "agent"))
+ require.NoError(t, err)
+ require.Equal(t, os.FileMode(0o750), info.Mode().Perm())
+
+ info, err = os.Stat(filepath.Join(restored, "agent", "instructions.md"))
+ require.NoError(t, err)
+ require.Equal(t, os.FileMode(0o640), info.Mode().Perm())
+}
+
+func TestRestoreRejectsUnsafeArchiveEntries(t *testing.T) {
+ testCases := []struct {
+ name string
+ header tar.Header
+ }{
+ {
+ name: "absolute paths fail",
+ header: tar.Header{
+ Name: "/escape",
+ Typeflag: tar.TypeReg,
+ Mode: 0o600,
+ },
+ },
+ {
+ name: "parent paths fail",
+ header: tar.Header{
+ Name: "../escape",
+ Typeflag: tar.TypeReg,
+ Mode: 0o600,
+ },
+ },
+ {
+ name: "embedded parent paths fail",
+ header: tar.Header{
+ Name: "agent/../../escape",
+ Typeflag: tar.TypeReg,
+ Mode: 0o600,
+ },
+ },
+ {
+ name: "symbolic links fail",
+ header: tar.Header{
+ Name: "agent",
+ Typeflag: tar.TypeSymlink,
+ Linkname: "../escape",
+ Mode: 0o777,
+ },
+ },
+ {
+ name: "hard links fail",
+ header: tar.Header{
+ Name: "agent",
+ Typeflag: tar.TypeLink,
+ Linkname: "/escape",
+ Mode: 0o600,
+ },
+ },
+ {
+ name: "device files fail",
+ header: tar.Header{
+ Name: "device",
+ Typeflag: tar.TypeChar,
+ Mode: 0o600,
+ },
+ },
+ {
+ name: "setuid modes fail",
+ header: tar.Header{
+ Name: "agent",
+ Typeflag: tar.TypeReg,
+ Mode: 0o4600,
+ },
+ },
+ {
+ name: "lock replacement fails",
+ header: tar.Header{
+ Name: "auth.lock",
+ Typeflag: tar.TypeReg,
+ Mode: 0o600,
+ },
+ },
+ }
+ for _, tc := range testCases {
+ t.Run(tc.name, func(t *testing.T) {
+ archive := backupTestArchive(t, tc.header)
+ require.Error(t, backupExtract(t.Context(), archive, t.TempDir()))
+ })
+ }
+}
+
+func TestBackupFailureKeepsTheDestination(t *testing.T) {
+ root := t.TempDir()
+ require.NoError(t, os.Symlink("missing", filepath.Join(root, "link")))
+
+ archive := filepath.Join(t.TempDir(), "backup.tar.gz")
+ require.NoError(t, os.WriteFile(archive, []byte("keep"), 0o600))
+
+ _, err := backupArchive(t.Context(), root, archive)
+ require.ErrorContains(t, err, "unsupported backup entry")
+
+ data, err := os.ReadFile(archive)
+ require.NoError(t, err)
+ require.Equal(t, "keep", string(data))
+
+ entries, err := os.ReadDir(filepath.Dir(archive))
+ require.NoError(t, err)
+ require.Len(t, entries, 1)
+}
+
+func backupTestArchive(t *testing.T, headers ...tar.Header) string {
+ t.Helper()
+
+ path := filepath.Join(t.TempDir(), "input.tar.gz")
+ file, err := os.Create(path)
+ require.NoError(t, err)
+
+ compressed := gzip.NewWriter(file)
+
+ archive := tar.NewWriter(compressed)
+ for _, header := range headers {
+ require.NoError(t, archive.WriteHeader(&header))
+ }
+
+ require.NoError(t, archive.Close())
+ require.NoError(t, compressed.Close())
+ require.NoError(t, file.Close())
+
+ return path
+}
projects / onechat / commitdiff