talos

1
// This Source Code Form is subject to the terms of the Mozilla Public
2
// License, v. 2.0. If a copy of the MPL was not distributed with this
3
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4

5
package archiver
6

7
import (
8
	"context"
9
	"os"
10
	"path/filepath"
11
	"strings"
12
)
13

14
// FileItem is unit of work for archive.
15
type FileItem struct {
16
	FullPath string
17
	RelPath  string
18
	FileInfo os.FileInfo
19
	Link     string
20
	Error    error
21
}
22

23
// FileType is a file type.
24
type FileType int
25

26
// File types.
27
const (
28
	RegularFileType FileType = iota
29
	DirectoryFileType
30
	SymlinkFileType
31
)
32

33
type walkerOptions struct {
34
	skipRoot        bool
35
	maxRecurseDepth int
36
	fnmatchPatterns []string
37
	types           map[FileType]struct{}
38
}
39

40
// WalkerOption configures Walker.
41
type WalkerOption func(*walkerOptions)
42

43
// WithSkipRoot skips root path if it's a directory.
44
func WithSkipRoot() WalkerOption {
45
	return func(o *walkerOptions) {
46
		o.skipRoot = true
47
	}
48
}
49

50
// WithMaxRecurseDepth controls maximum recursion depth while walking file tree.
51
//
52
// Value of -1 disables depth control.
53
func WithMaxRecurseDepth(maxDepth int) WalkerOption {
54
	return func(o *walkerOptions) {
55
		o.maxRecurseDepth = maxDepth
56
	}
57
}
58

59
// WithFnmatchPatterns filters results to match the patterns.
60
//
61
// Default is not to do any filtering.
62
func WithFnmatchPatterns(patterns ...string) WalkerOption {
63
	return func(o *walkerOptions) {
64
		o.fnmatchPatterns = append(o.fnmatchPatterns, patterns...)
65
	}
66
}
67

68
// WithFileTypes filters results by file types.
69
//
70
// Default is not to do any filtering.
71
func WithFileTypes(fileType ...FileType) WalkerOption {
72
	return func(o *walkerOptions) {
73
		o.types = make(map[FileType]struct{}, len(fileType))
74
		for _, t := range fileType {
75
			o.types[t] = struct{}{}
76
		}
77
	}
78
}
79

80
// Walker provides a channel of file info/paths for archival.
81
//
82
//nolint:gocyclo,cyclop
83
func Walker(ctx context.Context, rootPath string, options ...WalkerOption) (<-chan FileItem, error) {
84
	var opts walkerOptions
85
	opts.maxRecurseDepth = -1
86

87
	for _, o := range options {
88
		o(&opts)
89
	}
90

91
	info, err := os.Lstat(rootPath)
92
	if err != nil {
93
		return nil, err
94
	}
95

96
	if info.Mode()&os.ModeSymlink == os.ModeSymlink {
97
		rootPath, err = filepath.EvalSymlinks(rootPath)
98
		if err != nil {
99
			return nil, err
100
		}
101
	}
102

103
	ch := make(chan FileItem)
104

105
	go func() {
106
		defer close(ch)
107

108
		err := filepath.Walk(rootPath, func(path string, fileInfo os.FileInfo, walkErr error) error {
109
			item := FileItem{
110
				FullPath: path,
111
				FileInfo: fileInfo,
112
				Error:    walkErr,
113
			}
114

115
			if path == rootPath && !fileInfo.IsDir() {
116
				// only one file
117
				item.RelPath = filepath.Base(path)
118
			} else if item.Error == nil {
119
				item.RelPath, item.Error = filepath.Rel(rootPath, path)
120
			}
121

122
			// TODO: refactor all those `if item.Error == nil &&` conditions
123

124
			if item.Error == nil && len(opts.types) > 0 {
125
				var matches bool
126

127
				for t := range opts.types {
128
					switch t {
129
					case RegularFileType:
130
						matches = fileInfo.Mode()&os.ModeType == 0
131
					case DirectoryFileType:
132
						matches = fileInfo.Mode()&os.ModeDir != 0
133
					case SymlinkFileType:
134
						matches = fileInfo.Mode()&os.ModeSymlink != 0
135
					}
136

137
					if matches {
138
						break
139
					}
140
				}
141

142
				if !matches {
143
					return nil
144
				}
145
			}
146

147
			if item.Error == nil && path == rootPath && opts.skipRoot && fileInfo.IsDir() {
148
				// skip containing directory
149
				return nil
150
			}
151

152
			if item.Error == nil && fileInfo.Mode()&os.ModeSymlink == os.ModeSymlink {
153
				item.Link, item.Error = os.Readlink(path)
154
			}
155

156
			if item.Error == nil && len(opts.fnmatchPatterns) > 0 {
157
				var matches bool
158

159
				for _, pattern := range opts.fnmatchPatterns {
160
					if matches, _ = filepath.Match(pattern, item.RelPath); matches { //nolint:errcheck
161
						break
162
					}
163
				}
164

165
				if !matches {
166
					return nil
167
				}
168
			}
169

170
			select {
171
			case <-ctx.Done():
172
				return ctx.Err()
173
			case ch <- item:
174
			}
175

176
			if item.Error == nil && fileInfo.IsDir() && atMaxDepth(opts.maxRecurseDepth, rootPath, path) {
177
				return filepath.SkipDir
178
			}
179

180
			return nil
181
		})
182
		if err != nil {
183
			select {
184
			case <-ctx.Done():
185
			case ch <- FileItem{Error: err}:
186
			}
187
		}
188
	}()
189

190
	return ch, nil
191
}
192

193
// OSPathSeparator is the string version of the os.PathSeparator.
194
const OSPathSeparator = string(os.PathSeparator)
195

196
func atMaxDepth(max int, root, cur string) bool {
197
	if max < 0 {
198
		return false
199
	}
200

201
	if root == cur {
202
		// always recurse the root directory
203
		return false
204
	}
205

206
	return (strings.Count(cur, OSPathSeparator) - strings.Count(root, OSPathSeparator)) >= max
207
}
208