1
// Copyright © 2014 Steve Francia <[email protected]>.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
// http://www.apache.org/licenses/LICENSE-2.0
7
//
8
// Unless required by applicable law or agreed to in writing, software
9
// distributed under the License is distributed on an "AS IS" BASIS,
10
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11
// See the License for the specific language governing permissions and
12
// limitations under the License.
13

14
package afero
15

16
import (
17
	"fmt"
18
	"io"
19
	"log"
20
	"os"
21
	"path/filepath"
22
	"sort"
23
	"strings"
24
	"sync"
25
	"time"
26

27
	"github.com/spf13/afero/mem"
28
)
29

30
const chmodBits = os.ModePerm | os.ModeSetuid | os.ModeSetgid | os.ModeSticky // Only a subset of bits are allowed to be changed. Documented under os.Chmod()
31

32
type MemMapFs struct {
33
	mu   sync.RWMutex
34
	data map[string]*mem.FileData
35
	init sync.Once
36
}
37

38
func NewMemMapFs() Fs {
39
	return &MemMapFs{}
40
}
41

42
func (m *MemMapFs) getData() map[string]*mem.FileData {
43
	m.init.Do(func() {
44
		m.data = make(map[string]*mem.FileData)
45
		// Root should always exist, right?
46
		// TODO: what about windows?
47
		root := mem.CreateDir(FilePathSeparator)
48
		mem.SetMode(root, os.ModeDir|0o755)
49
		m.data[FilePathSeparator] = root
50
	})
51
	return m.data
52
}
53

54
func (*MemMapFs) Name() string { return "MemMapFS" }
55

56
func (m *MemMapFs) Create(name string) (File, error) {
57
	name = normalizePath(name)
58
	m.mu.Lock()
59
	file := mem.CreateFile(name)
60
	m.getData()[name] = file
61
	m.registerWithParent(file, 0)
62
	m.mu.Unlock()
63
	return mem.NewFileHandle(file), nil
64
}
65

66
func (m *MemMapFs) unRegisterWithParent(fileName string) error {
67
	f, err := m.lockfreeOpen(fileName)
68
	if err != nil {
69
		return err
70
	}
71
	parent := m.findParent(f)
72
	if parent == nil {
73
		log.Panic("parent of ", f.Name(), " is nil")
74
	}
75

76
	parent.Lock()
77
	mem.RemoveFromMemDir(parent, f)
78
	parent.Unlock()
79
	return nil
80
}
81

82
func (m *MemMapFs) findParent(f *mem.FileData) *mem.FileData {
83
	pdir, _ := filepath.Split(f.Name())
84
	pdir = filepath.Clean(pdir)
85
	pfile, err := m.lockfreeOpen(pdir)
86
	if err != nil {
87
		return nil
88
	}
89
	return pfile
90
}
91

92
func (m *MemMapFs) findDescendants(name string) []*mem.FileData {
93
	fData := m.getData()
94
	descendants := make([]*mem.FileData, 0, len(fData))
95
	for p, dFile := range fData {
96
		if strings.HasPrefix(p, name+FilePathSeparator) {
97
			descendants = append(descendants, dFile)
98
		}
99
	}
100

101
	sort.Slice(descendants, func(i, j int) bool {
102
		cur := len(strings.Split(descendants[i].Name(), FilePathSeparator))
103
		next := len(strings.Split(descendants[j].Name(), FilePathSeparator))
104
		return cur < next
105
	})
106

107
	return descendants
108
}
109

110
func (m *MemMapFs) registerWithParent(f *mem.FileData, perm os.FileMode) {
111
	if f == nil {
112
		return
113
	}
114
	parent := m.findParent(f)
115
	if parent == nil {
116
		pdir := filepath.Dir(filepath.Clean(f.Name()))
117
		err := m.lockfreeMkdir(pdir, perm)
118
		if err != nil {
119
			// log.Println("Mkdir error:", err)
120
			return
121
		}
122
		parent, err = m.lockfreeOpen(pdir)
123
		if err != nil {
124
			// log.Println("Open after Mkdir error:", err)
125
			return
126
		}
127
	}
128

129
	parent.Lock()
130
	mem.InitializeDir(parent)
131
	mem.AddToMemDir(parent, f)
132
	parent.Unlock()
133
}
134

135
func (m *MemMapFs) lockfreeMkdir(name string, perm os.FileMode) error {
136
	name = normalizePath(name)
137
	x, ok := m.getData()[name]
138
	if ok {
139
		// Only return ErrFileExists if it's a file, not a directory.
140
		i := mem.FileInfo{FileData: x}
141
		if !i.IsDir() {
142
			return ErrFileExists
143
		}
144
	} else {
145
		item := mem.CreateDir(name)
146
		mem.SetMode(item, os.ModeDir|perm)
147
		m.getData()[name] = item
148
		m.registerWithParent(item, perm)
149
	}
150
	return nil
151
}
152

153
func (m *MemMapFs) Mkdir(name string, perm os.FileMode) error {
154
	perm &= chmodBits
155
	name = normalizePath(name)
156

157
	m.mu.RLock()
158
	_, ok := m.getData()[name]
159
	m.mu.RUnlock()
160
	if ok {
161
		return &os.PathError{Op: "mkdir", Path: name, Err: ErrFileExists}
162
	}
163

164
	m.mu.Lock()
165
	// Dobule check that it doesn't exist.
166
	if _, ok := m.getData()[name]; ok {
167
		m.mu.Unlock()
168
		return &os.PathError{Op: "mkdir", Path: name, Err: ErrFileExists}
169
	}
170
	item := mem.CreateDir(name)
171
	mem.SetMode(item, os.ModeDir|perm)
172
	m.getData()[name] = item
173
	m.registerWithParent(item, perm)
174
	m.mu.Unlock()
175

176
	return m.setFileMode(name, perm|os.ModeDir)
177
}
178

179
func (m *MemMapFs) MkdirAll(path string, perm os.FileMode) error {
180
	err := m.Mkdir(path, perm)
181
	if err != nil {
182
		if err.(*os.PathError).Err == ErrFileExists {
183
			return nil
184
		}
185
		return err
186
	}
187
	return nil
188
}
189

190
// Handle some relative paths
191
func normalizePath(path string) string {
192
	path = filepath.Clean(path)
193

194
	switch path {
195
	case ".":
196
		return FilePathSeparator
197
	case "..":
198
		return FilePathSeparator
199
	default:
200
		return path
201
	}
202
}
203

204
func (m *MemMapFs) Open(name string) (File, error) {
205
	f, err := m.open(name)
206
	if f != nil {
207
		return mem.NewReadOnlyFileHandle(f), err
208
	}
209
	return nil, err
210
}
211

212
func (m *MemMapFs) openWrite(name string) (File, error) {
213
	f, err := m.open(name)
214
	if f != nil {
215
		return mem.NewFileHandle(f), err
216
	}
217
	return nil, err
218
}
219

220
func (m *MemMapFs) open(name string) (*mem.FileData, error) {
221
	name = normalizePath(name)
222

223
	m.mu.RLock()
224
	f, ok := m.getData()[name]
225
	m.mu.RUnlock()
226
	if !ok {
227
		return nil, &os.PathError{Op: "open", Path: name, Err: ErrFileNotFound}
228
	}
229
	return f, nil
230
}
231

232
func (m *MemMapFs) lockfreeOpen(name string) (*mem.FileData, error) {
233
	name = normalizePath(name)
234
	f, ok := m.getData()[name]
235
	if ok {
236
		return f, nil
237
	} else {
238
		return nil, ErrFileNotFound
239
	}
240
}
241

242
func (m *MemMapFs) OpenFile(name string, flag int, perm os.FileMode) (File, error) {
243
	perm &= chmodBits
244
	chmod := false
245
	file, err := m.openWrite(name)
246
	if err == nil && (flag&os.O_EXCL > 0) {
247
		return nil, &os.PathError{Op: "open", Path: name, Err: ErrFileExists}
248
	}
249
	if os.IsNotExist(err) && (flag&os.O_CREATE > 0) {
250
		file, err = m.Create(name)
251
		chmod = true
252
	}
253
	if err != nil {
254
		return nil, err
255
	}
256
	if flag == os.O_RDONLY {
257
		file = mem.NewReadOnlyFileHandle(file.(*mem.File).Data())
258
	}
259
	if flag&os.O_APPEND > 0 {
260
		_, err = file.Seek(0, io.SeekEnd)
261
		if err != nil {
262
			file.Close()
263
			return nil, err
264
		}
265
	}
266
	if flag&os.O_TRUNC > 0 && flag&(os.O_RDWR|os.O_WRONLY) > 0 {
267
		err = file.Truncate(0)
268
		if err != nil {
269
			file.Close()
270
			return nil, err
271
		}
272
	}
273
	if chmod {
274
		return file, m.setFileMode(name, perm)
275
	}
276
	return file, nil
277
}
278

279
func (m *MemMapFs) Remove(name string) error {
280
	name = normalizePath(name)
281

282
	m.mu.Lock()
283
	defer m.mu.Unlock()
284

285
	if _, ok := m.getData()[name]; ok {
286
		err := m.unRegisterWithParent(name)
287
		if err != nil {
288
			return &os.PathError{Op: "remove", Path: name, Err: err}
289
		}
290
		delete(m.getData(), name)
291
	} else {
292
		return &os.PathError{Op: "remove", Path: name, Err: os.ErrNotExist}
293
	}
294
	return nil
295
}
296

297
func (m *MemMapFs) RemoveAll(path string) error {
298
	path = normalizePath(path)
299
	m.mu.Lock()
300
	m.unRegisterWithParent(path)
301
	m.mu.Unlock()
302

303
	m.mu.RLock()
304
	defer m.mu.RUnlock()
305

306
	for p := range m.getData() {
307
		if p == path || strings.HasPrefix(p, path+FilePathSeparator) {
308
			m.mu.RUnlock()
309
			m.mu.Lock()
310
			delete(m.getData(), p)
311
			m.mu.Unlock()
312
			m.mu.RLock()
313
		}
314
	}
315
	return nil
316
}
317

318
func (m *MemMapFs) Rename(oldname, newname string) error {
319
	oldname = normalizePath(oldname)
320
	newname = normalizePath(newname)
321

322
	if oldname == newname {
323
		return nil
324
	}
325

326
	m.mu.RLock()
327
	defer m.mu.RUnlock()
328
	if _, ok := m.getData()[oldname]; ok {
329
		m.mu.RUnlock()
330
		m.mu.Lock()
331
		err := m.unRegisterWithParent(oldname)
332
		if err != nil {
333
			return err
334
		}
335

336
		fileData := m.getData()[oldname]
337
		mem.ChangeFileName(fileData, newname)
338
		m.getData()[newname] = fileData
339

340
		err = m.renameDescendants(oldname, newname)
341
		if err != nil {
342
			return err
343
		}
344

345
		delete(m.getData(), oldname)
346

347
		m.registerWithParent(fileData, 0)
348
		m.mu.Unlock()
349
		m.mu.RLock()
350
	} else {
351
		return &os.PathError{Op: "rename", Path: oldname, Err: ErrFileNotFound}
352
	}
353
	return nil
354
}
355

356
func (m *MemMapFs) renameDescendants(oldname, newname string) error {
357
	descendants := m.findDescendants(oldname)
358
	removes := make([]string, 0, len(descendants))
359
	for _, desc := range descendants {
360
		descNewName := strings.Replace(desc.Name(), oldname, newname, 1)
361
		err := m.unRegisterWithParent(desc.Name())
362
		if err != nil {
363
			return err
364
		}
365

366
		removes = append(removes, desc.Name())
367
		mem.ChangeFileName(desc, descNewName)
368
		m.getData()[descNewName] = desc
369

370
		m.registerWithParent(desc, 0)
371
	}
372
	for _, r := range removes {
373
		delete(m.getData(), r)
374
	}
375

376
	return nil
377
}
378

379
func (m *MemMapFs) LstatIfPossible(name string) (os.FileInfo, bool, error) {
380
	fileInfo, err := m.Stat(name)
381
	return fileInfo, false, err
382
}
383

384
func (m *MemMapFs) Stat(name string) (os.FileInfo, error) {
385
	f, err := m.Open(name)
386
	if err != nil {
387
		return nil, err
388
	}
389
	fi := mem.GetFileInfo(f.(*mem.File).Data())
390
	return fi, nil
391
}
392

393
func (m *MemMapFs) Chmod(name string, mode os.FileMode) error {
394
	mode &= chmodBits
395

396
	m.mu.RLock()
397
	f, ok := m.getData()[name]
398
	m.mu.RUnlock()
399
	if !ok {
400
		return &os.PathError{Op: "chmod", Path: name, Err: ErrFileNotFound}
401
	}
402
	prevOtherBits := mem.GetFileInfo(f).Mode() & ^chmodBits
403

404
	mode = prevOtherBits | mode
405
	return m.setFileMode(name, mode)
406
}
407

408
func (m *MemMapFs) setFileMode(name string, mode os.FileMode) error {
409
	name = normalizePath(name)
410

411
	m.mu.RLock()
412
	f, ok := m.getData()[name]
413
	m.mu.RUnlock()
414
	if !ok {
415
		return &os.PathError{Op: "chmod", Path: name, Err: ErrFileNotFound}
416
	}
417

418
	m.mu.Lock()
419
	mem.SetMode(f, mode)
420
	m.mu.Unlock()
421

422
	return nil
423
}
424

425
func (m *MemMapFs) Chown(name string, uid, gid int) error {
426
	name = normalizePath(name)
427

428
	m.mu.RLock()
429
	f, ok := m.getData()[name]
430
	m.mu.RUnlock()
431
	if !ok {
432
		return &os.PathError{Op: "chown", Path: name, Err: ErrFileNotFound}
433
	}
434

435
	mem.SetUID(f, uid)
436
	mem.SetGID(f, gid)
437

438
	return nil
439
}
440

441
func (m *MemMapFs) Chtimes(name string, atime time.Time, mtime time.Time) error {
442
	name = normalizePath(name)
443

444
	m.mu.RLock()
445
	f, ok := m.getData()[name]
446
	m.mu.RUnlock()
447
	if !ok {
448
		return &os.PathError{Op: "chtimes", Path: name, Err: ErrFileNotFound}
449
	}
450

451
	m.mu.Lock()
452
	mem.SetModTime(f, mtime)
453
	m.mu.Unlock()
454

455
	return nil
456
}
457

458
func (m *MemMapFs) List() {
459
	for _, x := range m.data {
460
		y := mem.FileInfo{FileData: x}
461
		fmt.Println(x.Name(), y.Size())
462
	}
463
}
464

465