1package fsutil
2
3import (
4	"context"
5	"hash"
6	"io"
7	"os"
8	"path/filepath"
9	"strconv"
10	"sync"
11	"syscall"
12	"time"
13
14	"github.com/opencontainers/go-digest"
15	"github.com/pkg/errors"
16	"github.com/tonistiigi/fsutil/types"
17	"golang.org/x/sync/errgroup"
18)
19
20type WriteToFunc func(context.Context, string, io.WriteCloser) error
21
22type DiskWriterOpt struct {
23	AsyncDataCb   WriteToFunc
24	SyncDataCb    WriteToFunc
25	NotifyCb      func(ChangeKind, string, os.FileInfo, error) error
26	ContentHasher ContentHasher
27	Filter        FilterFunc
28}
29
30type FilterFunc func(string, *types.Stat) bool
31
32type DiskWriter struct {
33	opt  DiskWriterOpt
34	dest string
35
36	ctx    context.Context
37	cancel func()
38	eg     *errgroup.Group
39	filter FilterFunc
40}
41
42func NewDiskWriter(ctx context.Context, dest string, opt DiskWriterOpt) (*DiskWriter, error) {
43	if opt.SyncDataCb == nil && opt.AsyncDataCb == nil {
44		return nil, errors.New("no data callback specified")
45	}
46	if opt.SyncDataCb != nil && opt.AsyncDataCb != nil {
47		return nil, errors.New("can't specify both sync and async data callbacks")
48	}
49
50	ctx, cancel := context.WithCancel(ctx)
51	eg, ctx := errgroup.WithContext(ctx)
52
53	return &DiskWriter{
54		opt:    opt,
55		dest:   dest,
56		eg:     eg,
57		ctx:    ctx,
58		cancel: cancel,
59		filter: opt.Filter,
60	}, nil
61}
62
63func (dw *DiskWriter) Wait(ctx context.Context) error {
64	return dw.eg.Wait()
65}
66
67func (dw *DiskWriter) HandleChange(kind ChangeKind, p string, fi os.FileInfo, err error) (retErr error) {
68	if err != nil {
69		return err
70	}
71
72	select {
73	case <-dw.ctx.Done():
74		return dw.ctx.Err()
75	default:
76	}
77
78	defer func() {
79		if retErr != nil {
80			dw.cancel()
81		}
82	}()
83
84	destPath := filepath.Join(dw.dest, filepath.FromSlash(p))
85
86	if kind == ChangeKindDelete {
87		if dw.filter != nil {
88			var empty types.Stat
89			if ok := dw.filter(p, &empty); !ok {
90				return nil
91			}
92		}
93		// todo: no need to validate if diff is trusted but is it always?
94		if err := os.RemoveAll(destPath); err != nil {
95			return errors.Wrapf(err, "failed to remove: %s", destPath)
96		}
97		if dw.opt.NotifyCb != nil {
98			if err := dw.opt.NotifyCb(kind, p, nil, nil); err != nil {
99				return err
100			}
101		}
102		return nil
103	}
104
105	stat, ok := fi.Sys().(*types.Stat)
106	if !ok {
107		return errors.WithStack(&os.PathError{Path: p, Err: syscall.EBADMSG, Op: "change without stat info"})
108	}
109
110	statCopy := *stat
111
112	if dw.filter != nil {
113		if ok := dw.filter(p, &statCopy); !ok {
114			return nil
115		}
116	}
117
118	rename := true
119	oldFi, err := os.Lstat(destPath)
120	if err != nil {
121		if errors.Is(err, os.ErrNotExist) {
122			if kind != ChangeKindAdd {
123				return errors.Wrap(err, "modify/rm")
124			}
125			rename = false
126		} else {
127			return errors.WithStack(err)
128		}
129	}
130
131	if oldFi != nil && fi.IsDir() && oldFi.IsDir() {
132		if err := rewriteMetadata(destPath, &statCopy); err != nil {
133			return errors.Wrapf(err, "error setting dir metadata for %s", destPath)
134		}
135		return nil
136	}
137
138	newPath := destPath
139	if rename {
140		newPath = filepath.Join(filepath.Dir(destPath), ".tmp."+nextSuffix())
141	}
142
143	isRegularFile := false
144
145	switch {
146	case fi.IsDir():
147		if err := os.Mkdir(newPath, fi.Mode()); err != nil {
148			return errors.Wrapf(err, "failed to create dir %s", newPath)
149		}
150	case fi.Mode()&os.ModeDevice != 0 || fi.Mode()&os.ModeNamedPipe != 0:
151		if err := handleTarTypeBlockCharFifo(newPath, &statCopy); err != nil {
152			return errors.Wrapf(err, "failed to create device %s", newPath)
153		}
154	case fi.Mode()&os.ModeSymlink != 0:
155		if err := os.Symlink(statCopy.Linkname, newPath); err != nil {
156			return errors.Wrapf(err, "failed to symlink %s", newPath)
157		}
158	case statCopy.Linkname != "":
159		if err := os.Link(filepath.Join(dw.dest, statCopy.Linkname), newPath); err != nil {
160			return errors.Wrapf(err, "failed to link %s to %s", newPath, statCopy.Linkname)
161		}
162	default:
163		isRegularFile = true
164		file, err := os.OpenFile(newPath, os.O_CREATE|os.O_WRONLY, fi.Mode()) //todo: windows
165		if err != nil {
166			return errors.Wrapf(err, "failed to create %s", newPath)
167		}
168		if dw.opt.SyncDataCb != nil {
169			if err := dw.processChange(ChangeKindAdd, p, fi, file); err != nil {
170				file.Close()
171				return err
172			}
173			break
174		}
175		if err := file.Close(); err != nil {
176			return errors.Wrapf(err, "failed to close %s", newPath)
177		}
178	}
179
180	if err := rewriteMetadata(newPath, &statCopy); err != nil {
181		return errors.Wrapf(err, "error setting metadata for %s", newPath)
182	}
183
184	if rename {
185		if oldFi.IsDir() != fi.IsDir() {
186			if err := os.RemoveAll(destPath); err != nil {
187				return errors.Wrapf(err, "failed to remove %s", destPath)
188			}
189		}
190		if err := os.Rename(newPath, destPath); err != nil {
191			return errors.Wrapf(err, "failed to rename %s to %s", newPath, destPath)
192		}
193	}
194
195	if isRegularFile {
196		if dw.opt.AsyncDataCb != nil {
197			dw.requestAsyncFileData(p, destPath, fi, &statCopy)
198		}
199	} else {
200		return dw.processChange(kind, p, fi, nil)
201	}
202
203	return nil
204}
205
206func (dw *DiskWriter) requestAsyncFileData(p, dest string, fi os.FileInfo, st *types.Stat) {
207	// todo: limit worker threads
208	dw.eg.Go(func() error {
209		if err := dw.processChange(ChangeKindAdd, p, fi, &lazyFileWriter{
210			dest: dest,
211		}); err != nil {
212			return err
213		}
214		return chtimes(dest, st.ModTime) // TODO: parent dirs
215	})
216}
217
218func (dw *DiskWriter) processChange(kind ChangeKind, p string, fi os.FileInfo, w io.WriteCloser) error {
219	origw := w
220	var hw *hashedWriter
221	if dw.opt.NotifyCb != nil {
222		var err error
223		if hw, err = newHashWriter(dw.opt.ContentHasher, fi, w); err != nil {
224			return err
225		}
226		w = hw
227	}
228	if origw != nil {
229		fn := dw.opt.SyncDataCb
230		if fn == nil && dw.opt.AsyncDataCb != nil {
231			fn = dw.opt.AsyncDataCb
232		}
233		if err := fn(dw.ctx, p, w); err != nil {
234			return err
235		}
236	} else {
237		if hw != nil {
238			hw.Close()
239		}
240	}
241	if hw != nil {
242		return dw.opt.NotifyCb(kind, p, hw, nil)
243	}
244	return nil
245}
246
247type hashedWriter struct {
248	os.FileInfo
249	io.Writer
250	h    hash.Hash
251	w    io.WriteCloser
252	dgst digest.Digest
253}
254
255func newHashWriter(ch ContentHasher, fi os.FileInfo, w io.WriteCloser) (*hashedWriter, error) {
256	stat, ok := fi.Sys().(*types.Stat)
257	if !ok {
258		return nil, errors.Errorf("invalid change without stat information")
259	}
260
261	h, err := ch(stat)
262	if err != nil {
263		return nil, err
264	}
265	hw := &hashedWriter{
266		FileInfo: fi,
267		Writer:   io.MultiWriter(w, h),
268		h:        h,
269		w:        w,
270	}
271	return hw, nil
272}
273
274func (hw *hashedWriter) Close() error {
275	hw.dgst = digest.NewDigest(digest.SHA256, hw.h)
276	if hw.w != nil {
277		return hw.w.Close()
278	}
279	return nil
280}
281
282func (hw *hashedWriter) Digest() digest.Digest {
283	return hw.dgst
284}
285
286type lazyFileWriter struct {
287	dest     string
288	f        *os.File
289	fileMode *os.FileMode
290}
291
292func (lfw *lazyFileWriter) Write(dt []byte) (int, error) {
293	if lfw.f == nil {
294		file, err := os.OpenFile(lfw.dest, os.O_WRONLY, 0) //todo: windows
295		if os.IsPermission(err) {
296			// retry after chmod
297			fi, er := os.Stat(lfw.dest)
298			if er == nil {
299				mode := fi.Mode()
300				lfw.fileMode = &mode
301				er = os.Chmod(lfw.dest, mode|0222)
302				if er == nil {
303					file, err = os.OpenFile(lfw.dest, os.O_WRONLY, 0)
304				}
305			}
306		}
307		if err != nil {
308			return 0, errors.Wrapf(err, "failed to open %s", lfw.dest)
309		}
310		lfw.f = file
311	}
312	return lfw.f.Write(dt)
313}
314
315func (lfw *lazyFileWriter) Close() error {
316	var err error
317	if lfw.f != nil {
318		err = lfw.f.Close()
319	}
320	if err == nil && lfw.fileMode != nil {
321		err = os.Chmod(lfw.dest, *lfw.fileMode)
322	}
323	return err
324}
325
326func mkdev(major int64, minor int64) uint32 {
327	return uint32(((minor & 0xfff00) << 12) | ((major & 0xfff) << 8) | (minor & 0xff))
328}
329
330// Random number state.
331// We generate random temporary file names so that there's a good
332// chance the file doesn't exist yet - keeps the number of tries in
333// TempFile to a minimum.
334var rand uint32
335var randmu sync.Mutex
336
337func reseed() uint32 {
338	return uint32(time.Now().UnixNano() + int64(os.Getpid()))
339}
340
341func nextSuffix() string {
342	randmu.Lock()
343	r := rand
344	if r == 0 {
345		r = reseed()
346	}
347	r = r*1664525 + 1013904223 // constants from Numerical Recipes
348	rand = r
349	randmu.Unlock()
350	return strconv.Itoa(int(1e9 + r%1e9))[1:]
351}
352