1// Copyright © 2014 Steve Francia <spf@spf13.com>.
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
14package afero
15
16import (
17	"fmt"
18	"log"
19	"os"
20	"path/filepath"
21	"strings"
22	"sync"
23	"time"
24
25	"github.com/spf13/afero/mem"
26)
27
28const chmodBits = os.ModePerm | os.ModeSetuid | os.ModeSetgid | os.ModeSticky // Only a subset of bits are allowed to be changed. Documented under os.Chmod()
29
30type MemMapFs struct {
31	mu   sync.RWMutex
32	data map[string]*mem.FileData
33	init sync.Once
34}
35
36func NewMemMapFs() Fs {
37	return &MemMapFs{}
38}
39
40func (m *MemMapFs) getData() map[string]*mem.FileData {
41	m.init.Do(func() {
42		m.data = make(map[string]*mem.FileData)
43		// Root should always exist, right?
44		// TODO: what about windows?
45		root := mem.CreateDir(FilePathSeparator)
46		mem.SetMode(root, os.ModeDir|0755)
47		m.data[FilePathSeparator] = root
48	})
49	return m.data
50}
51
52func (*MemMapFs) Name() string { return "MemMapFS" }
53
54func (m *MemMapFs) Create(name string) (File, error) {
55	name = normalizePath(name)
56	m.mu.Lock()
57	file := mem.CreateFile(name)
58	m.getData()[name] = file
59	m.registerWithParent(file, 0)
60	m.mu.Unlock()
61	return mem.NewFileHandle(file), nil
62}
63
64func (m *MemMapFs) unRegisterWithParent(fileName string) error {
65	f, err := m.lockfreeOpen(fileName)
66	if err != nil {
67		return err
68	}
69	parent := m.findParent(f)
70	if parent == nil {
71		log.Panic("parent of ", f.Name(), " is nil")
72	}
73
74	parent.Lock()
75	mem.RemoveFromMemDir(parent, f)
76	parent.Unlock()
77	return nil
78}
79
80func (m *MemMapFs) findParent(f *mem.FileData) *mem.FileData {
81	pdir, _ := filepath.Split(f.Name())
82	pdir = filepath.Clean(pdir)
83	pfile, err := m.lockfreeOpen(pdir)
84	if err != nil {
85		return nil
86	}
87	return pfile
88}
89
90func (m *MemMapFs) registerWithParent(f *mem.FileData, perm os.FileMode) {
91	if f == nil {
92		return
93	}
94	parent := m.findParent(f)
95	if parent == nil {
96		pdir := filepath.Dir(filepath.Clean(f.Name()))
97		err := m.lockfreeMkdir(pdir, perm)
98		if err != nil {
99			//log.Println("Mkdir error:", err)
100			return
101		}
102		parent, err = m.lockfreeOpen(pdir)
103		if err != nil {
104			//log.Println("Open after Mkdir error:", err)
105			return
106		}
107	}
108
109	parent.Lock()
110	mem.InitializeDir(parent)
111	mem.AddToMemDir(parent, f)
112	parent.Unlock()
113}
114
115func (m *MemMapFs) lockfreeMkdir(name string, perm os.FileMode) error {
116	name = normalizePath(name)
117	x, ok := m.getData()[name]
118	if ok {
119		// Only return ErrFileExists if it's a file, not a directory.
120		i := mem.FileInfo{FileData: x}
121		if !i.IsDir() {
122			return ErrFileExists
123		}
124	} else {
125		item := mem.CreateDir(name)
126		mem.SetMode(item, os.ModeDir|perm)
127		m.getData()[name] = item
128		m.registerWithParent(item, perm)
129	}
130	return nil
131}
132
133func (m *MemMapFs) Mkdir(name string, perm os.FileMode) error {
134	perm &= chmodBits
135	name = normalizePath(name)
136
137	m.mu.RLock()
138	_, ok := m.getData()[name]
139	m.mu.RUnlock()
140	if ok {
141		return &os.PathError{Op: "mkdir", Path: name, Err: ErrFileExists}
142	}
143
144	m.mu.Lock()
145	item := mem.CreateDir(name)
146	mem.SetMode(item, os.ModeDir|perm)
147	m.getData()[name] = item
148	m.registerWithParent(item, perm)
149	m.mu.Unlock()
150
151	return m.setFileMode(name, perm|os.ModeDir)
152}
153
154func (m *MemMapFs) MkdirAll(path string, perm os.FileMode) error {
155	err := m.Mkdir(path, perm)
156	if err != nil {
157		if err.(*os.PathError).Err == ErrFileExists {
158			return nil
159		}
160		return err
161	}
162	return nil
163}
164
165// Handle some relative paths
166func normalizePath(path string) string {
167	path = filepath.Clean(path)
168
169	switch path {
170	case ".":
171		return FilePathSeparator
172	case "..":
173		return FilePathSeparator
174	default:
175		return path
176	}
177}
178
179func (m *MemMapFs) Open(name string) (File, error) {
180	f, err := m.open(name)
181	if f != nil {
182		return mem.NewReadOnlyFileHandle(f), err
183	}
184	return nil, err
185}
186
187func (m *MemMapFs) openWrite(name string) (File, error) {
188	f, err := m.open(name)
189	if f != nil {
190		return mem.NewFileHandle(f), err
191	}
192	return nil, err
193}
194
195func (m *MemMapFs) open(name string) (*mem.FileData, error) {
196	name = normalizePath(name)
197
198	m.mu.RLock()
199	f, ok := m.getData()[name]
200	m.mu.RUnlock()
201	if !ok {
202		return nil, &os.PathError{Op: "open", Path: name, Err: ErrFileNotFound}
203	}
204	return f, nil
205}
206
207func (m *MemMapFs) lockfreeOpen(name string) (*mem.FileData, error) {
208	name = normalizePath(name)
209	f, ok := m.getData()[name]
210	if ok {
211		return f, nil
212	} else {
213		return nil, ErrFileNotFound
214	}
215}
216
217func (m *MemMapFs) OpenFile(name string, flag int, perm os.FileMode) (File, error) {
218	perm &= chmodBits
219	chmod := false
220	file, err := m.openWrite(name)
221	if err == nil && (flag&os.O_EXCL > 0) {
222		return nil, &os.PathError{Op: "open", Path: name, Err: ErrFileExists}
223	}
224	if os.IsNotExist(err) && (flag&os.O_CREATE > 0) {
225		file, err = m.Create(name)
226		chmod = true
227	}
228	if err != nil {
229		return nil, err
230	}
231	if flag == os.O_RDONLY {
232		file = mem.NewReadOnlyFileHandle(file.(*mem.File).Data())
233	}
234	if flag&os.O_APPEND > 0 {
235		_, err = file.Seek(0, os.SEEK_END)
236		if err != nil {
237			file.Close()
238			return nil, err
239		}
240	}
241	if flag&os.O_TRUNC > 0 && flag&(os.O_RDWR|os.O_WRONLY) > 0 {
242		err = file.Truncate(0)
243		if err != nil {
244			file.Close()
245			return nil, err
246		}
247	}
248	if chmod {
249		return file, m.setFileMode(name, perm)
250	}
251	return file, nil
252}
253
254func (m *MemMapFs) Remove(name string) error {
255	name = normalizePath(name)
256
257	m.mu.Lock()
258	defer m.mu.Unlock()
259
260	if _, ok := m.getData()[name]; ok {
261		err := m.unRegisterWithParent(name)
262		if err != nil {
263			return &os.PathError{Op: "remove", Path: name, Err: err}
264		}
265		delete(m.getData(), name)
266	} else {
267		return &os.PathError{Op: "remove", Path: name, Err: os.ErrNotExist}
268	}
269	return nil
270}
271
272func (m *MemMapFs) RemoveAll(path string) error {
273	path = normalizePath(path)
274	m.mu.Lock()
275	m.unRegisterWithParent(path)
276	m.mu.Unlock()
277
278	m.mu.RLock()
279	defer m.mu.RUnlock()
280
281	for p := range m.getData() {
282		if strings.HasPrefix(p, path) {
283			m.mu.RUnlock()
284			m.mu.Lock()
285			delete(m.getData(), p)
286			m.mu.Unlock()
287			m.mu.RLock()
288		}
289	}
290	return nil
291}
292
293func (m *MemMapFs) Rename(oldname, newname string) error {
294	oldname = normalizePath(oldname)
295	newname = normalizePath(newname)
296
297	if oldname == newname {
298		return nil
299	}
300
301	m.mu.RLock()
302	defer m.mu.RUnlock()
303	if _, ok := m.getData()[oldname]; ok {
304		m.mu.RUnlock()
305		m.mu.Lock()
306		m.unRegisterWithParent(oldname)
307		fileData := m.getData()[oldname]
308		delete(m.getData(), oldname)
309		mem.ChangeFileName(fileData, newname)
310		m.getData()[newname] = fileData
311		m.registerWithParent(fileData, 0)
312		m.mu.Unlock()
313		m.mu.RLock()
314	} else {
315		return &os.PathError{Op: "rename", Path: oldname, Err: ErrFileNotFound}
316	}
317	return nil
318}
319
320func (m *MemMapFs) LstatIfPossible(name string) (os.FileInfo, bool, error) {
321	fileInfo, err := m.Stat(name)
322	return fileInfo, false, err
323}
324
325func (m *MemMapFs) Stat(name string) (os.FileInfo, error) {
326	f, err := m.Open(name)
327	if err != nil {
328		return nil, err
329	}
330	fi := mem.GetFileInfo(f.(*mem.File).Data())
331	return fi, nil
332}
333
334func (m *MemMapFs) Chmod(name string, mode os.FileMode) error {
335	mode &= chmodBits
336
337	m.mu.RLock()
338	f, ok := m.getData()[name]
339	m.mu.RUnlock()
340	if !ok {
341		return &os.PathError{Op: "chmod", Path: name, Err: ErrFileNotFound}
342	}
343	prevOtherBits := mem.GetFileInfo(f).Mode() & ^chmodBits
344
345	mode = prevOtherBits | mode
346	return m.setFileMode(name, mode)
347}
348
349func (m *MemMapFs) setFileMode(name string, mode os.FileMode) error {
350	name = normalizePath(name)
351
352	m.mu.RLock()
353	f, ok := m.getData()[name]
354	m.mu.RUnlock()
355	if !ok {
356		return &os.PathError{Op: "chmod", Path: name, Err: ErrFileNotFound}
357	}
358
359	m.mu.Lock()
360	mem.SetMode(f, mode)
361	m.mu.Unlock()
362
363	return nil
364}
365
366func (m *MemMapFs) Chown(name string, uid, gid int) error {
367	name = normalizePath(name)
368
369	m.mu.RLock()
370	f, ok := m.getData()[name]
371	m.mu.RUnlock()
372	if !ok {
373		return &os.PathError{Op: "chown", Path: name, Err: ErrFileNotFound}
374	}
375
376	mem.SetUID(f, uid)
377	mem.SetGID(f, gid)
378
379	return nil
380}
381
382func (m *MemMapFs) Chtimes(name string, atime time.Time, mtime time.Time) error {
383	name = normalizePath(name)
384
385	m.mu.RLock()
386	f, ok := m.getData()[name]
387	m.mu.RUnlock()
388	if !ok {
389		return &os.PathError{Op: "chtimes", Path: name, Err: ErrFileNotFound}
390	}
391
392	m.mu.Lock()
393	mem.SetModTime(f, mtime)
394	m.mu.Unlock()
395
396	return nil
397}
398
399func (m *MemMapFs) List() {
400	for _, x := range m.data {
401		y := mem.FileInfo{FileData: x}
402		fmt.Println(x.Name(), y.Size())
403	}
404}
405