git » chasquid » main » tree

[main] / internal / safeio / safeio.go 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
// Package safeio implements convenient I/O routines that provide additional
// levels of safety in the presence of unexpected failures.
package safeio

import (
	"os"
	"path"
	"syscall"
)

// osFile is an interface to the methods of os.File that we need, so we can
// simulate failures in tests.
type osFile interface {
	Name() string
	Chmod(os.FileMode) error
	Chown(int, int) error
	Write([]byte) (int, error)
	Close() error
}

var createTemp func(dir, pattern string) (osFile, error) = func(
	dir, pattern string) (osFile, error) {
	return os.CreateTemp(dir, pattern)
}

// FileOp represents an operation on a file (passed by its name).
type FileOp func(fname string) error

// WriteFile writes data to a file named by filename, atomically.
//
// It's a wrapper to os.WriteFile, but provides atomicity (and increased
// safety) by writing to a temporary file and renaming it at the end.
//
// Before the final rename, the given ops (if any) are called. They can be
// used to manipulate the file before it is atomically renamed.
// If any operation fails, the file is removed and the error is returned.
//
// Note this relies on same-directory Rename being atomic, which holds in most
// reasonably modern filesystems.
func WriteFile(filename string, data []byte, perm os.FileMode, ops ...FileOp) error {
	// Note we create the temporary file in the same directory, otherwise we
	// would have no expectation of Rename being atomic.
	// We make the file names start with "." so there's no confusion with the
	// originals.
	tmpf, err := createTemp(path.Dir(filename), "."+path.Base(filename))
	if err != nil {
		return err
	}

	if err = tmpf.Chmod(perm); err != nil {
		tmpf.Close()
		os.Remove(tmpf.Name())
		return err
	}

	if uid, gid := getOwner(filename); uid >= 0 {
		if err = tmpf.Chown(uid, gid); err != nil {
			tmpf.Close()
			os.Remove(tmpf.Name())
			return err
		}
	}

	if _, err = tmpf.Write(data); err != nil {
		tmpf.Close()
		os.Remove(tmpf.Name())
		return err
	}

	if err = tmpf.Close(); err != nil {
		os.Remove(tmpf.Name())
		return err
	}

	for _, op := range ops {
		if err = op(tmpf.Name()); err != nil {
			os.Remove(tmpf.Name())
			return err
		}
	}

	return os.Rename(tmpf.Name(), filename)
}

func getOwner(fname string) (uid, gid int) {
	uid = -1
	gid = -1
	stat, err := os.Stat(fname)
	if err == nil {
		if sysstat, ok := stat.Sys().(*syscall.Stat_t); ok {
			uid = int(sysstat.Uid)
			gid = int(sysstat.Gid)
		}
	}

	return
}