memmap.go raw

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