steamcache2/vfs/memory/memory.go

// vfs/memory/memory.go
package memory

import (
	"bytes"
	"container/list"
	"io"
	"s1d3sw1ped/SteamCache2/vfs"
	"s1d3sw1ped/SteamCache2/vfs/vfserror"
	"strings"
	"sync"
)

// Ensure MemoryFS implements VFS.
var _ vfs.VFS = (*MemoryFS)(nil)

// MemoryFS is an in-memory virtual file system
type MemoryFS struct {
	data     map[string]*bytes.Buffer
	info     map[string]*vfs.FileInfo
	capacity int64
	size     int64
	mu       sync.RWMutex
	keyLocks sync.Map // map[string]*sync.RWMutex
	LRU      *lruList
}

// lruList for time-decayed LRU eviction
type lruList struct {
	list *list.List
	elem map[string]*list.Element
}

func newLruList() *lruList {
	return &lruList{
		list: list.New(),
		elem: make(map[string]*list.Element),
	}
}

func (l *lruList) Add(key string, fi *vfs.FileInfo) {
	elem := l.list.PushFront(fi)
	l.elem[key] = elem
}

func (l *lruList) MoveToFront(key string) {
	if elem, exists := l.elem[key]; exists {
		l.list.MoveToFront(elem)
		// Update the FileInfo in the element with new access time
		if fi := elem.Value.(*vfs.FileInfo); fi != nil {
			fi.UpdateAccess()
		}
	}
}

func (l *lruList) Remove(key string) *vfs.FileInfo {
	if elem, exists := l.elem[key]; exists {
		delete(l.elem, key)
		if fi := l.list.Remove(elem).(*vfs.FileInfo); fi != nil {
			return fi
		}
	}
	return nil
}

func (l *lruList) Len() int {
	return l.list.Len()
}

// New creates a new MemoryFS
func New(capacity int64) *MemoryFS {
	if capacity <= 0 {
		panic("memory capacity must be greater than 0")
	}

	return &MemoryFS{
		data:     make(map[string]*bytes.Buffer),
		info:     make(map[string]*vfs.FileInfo),
		capacity: capacity,
		size:     0,
		LRU:      newLruList(),
	}
}

// Name returns the name of this VFS
func (m *MemoryFS) Name() string {
	return "MemoryFS"
}

// Size returns the current size
func (m *MemoryFS) Size() int64 {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.size
}

// Capacity returns the maximum capacity
func (m *MemoryFS) Capacity() int64 {
	return m.capacity
}

// getKeyLock returns a lock for the given key
func (m *MemoryFS) getKeyLock(key string) *sync.RWMutex {
	keyLock, _ := m.keyLocks.LoadOrStore(key, &sync.RWMutex{})
	return keyLock.(*sync.RWMutex)
}

// Create creates a new file
func (m *MemoryFS) Create(key string, size int64) (io.WriteCloser, error) {
	if key == "" {
		return nil, vfserror.ErrInvalidKey
	}
	if key[0] == '/' {
		return nil, vfserror.ErrInvalidKey
	}

	// Sanitize key to prevent path traversal
	if strings.Contains(key, "..") {
		return nil, vfserror.ErrInvalidKey
	}

	keyMu := m.getKeyLock(key)
	keyMu.Lock()
	defer keyMu.Unlock()

	m.mu.Lock()
	// Check if file already exists and handle overwrite
	if fi, exists := m.info[key]; exists {
		m.size -= fi.Size
		m.LRU.Remove(key)
		delete(m.info, key)
		delete(m.data, key)
	}

	buffer := &bytes.Buffer{}
	m.data[key] = buffer
	fi := vfs.NewFileInfo(key, size)
	m.info[key] = fi
	m.LRU.Add(key, fi)
	m.size += size
	m.mu.Unlock()

	return &memoryWriteCloser{
		buffer: buffer,
		memory: m,
		key:    key,
	}, nil
}

// memoryWriteCloser implements io.WriteCloser for memory files
type memoryWriteCloser struct {
	buffer *bytes.Buffer
	memory *MemoryFS
	key    string
}

func (mwc *memoryWriteCloser) Write(p []byte) (n int, err error) {
	return mwc.buffer.Write(p)
}

func (mwc *memoryWriteCloser) Close() error {
	// Update the actual size in FileInfo
	mwc.memory.mu.Lock()
	if fi, exists := mwc.memory.info[mwc.key]; exists {
		actualSize := int64(mwc.buffer.Len())
		sizeDiff := actualSize - fi.Size
		fi.Size = actualSize
		mwc.memory.size += sizeDiff
	}
	mwc.memory.mu.Unlock()
	return nil
}

// Open opens a file for reading
func (m *MemoryFS) Open(key string) (io.ReadCloser, error) {
	if key == "" {
		return nil, vfserror.ErrInvalidKey
	}
	if key[0] == '/' {
		return nil, vfserror.ErrInvalidKey
	}

	if strings.Contains(key, "..") {
		return nil, vfserror.ErrInvalidKey
	}

	keyMu := m.getKeyLock(key)
	keyMu.RLock()
	defer keyMu.RUnlock()

	m.mu.Lock()
	fi, exists := m.info[key]
	if !exists {
		m.mu.Unlock()
		return nil, vfserror.ErrNotFound
	}
	fi.UpdateAccess()
	m.LRU.MoveToFront(key)

	buffer, exists := m.data[key]
	if !exists {
		m.mu.Unlock()
		return nil, vfserror.ErrNotFound
	}

	// Create a copy of the buffer for reading
	data := make([]byte, buffer.Len())
	copy(data, buffer.Bytes())
	m.mu.Unlock()

	return &memoryReadCloser{
		reader: bytes.NewReader(data),
	}, nil
}

// memoryReadCloser implements io.ReadCloser for memory files
type memoryReadCloser struct {
	reader *bytes.Reader
}

func (mrc *memoryReadCloser) Read(p []byte) (n int, err error) {
	return mrc.reader.Read(p)
}

func (mrc *memoryReadCloser) Close() error {
	return nil
}

// Delete removes a file
func (m *MemoryFS) Delete(key string) error {
	if key == "" {
		return vfserror.ErrInvalidKey
	}
	if key[0] == '/' {
		return vfserror.ErrInvalidKey
	}

	if strings.Contains(key, "..") {
		return vfserror.ErrInvalidKey
	}

	keyMu := m.getKeyLock(key)
	keyMu.Lock()
	defer keyMu.Unlock()

	m.mu.Lock()
	fi, exists := m.info[key]
	if !exists {
		m.mu.Unlock()
		return vfserror.ErrNotFound
	}
	m.size -= fi.Size
	m.LRU.Remove(key)
	delete(m.info, key)
	delete(m.data, key)
	m.mu.Unlock()

	return nil
}

// Stat returns file information
func (m *MemoryFS) Stat(key string) (*vfs.FileInfo, error) {
	if key == "" {
		return nil, vfserror.ErrInvalidKey
	}
	if key[0] == '/' {
		return nil, vfserror.ErrInvalidKey
	}

	if strings.Contains(key, "..") {
		return nil, vfserror.ErrInvalidKey
	}

	keyMu := m.getKeyLock(key)
	keyMu.RLock()
	defer keyMu.RUnlock()

	m.mu.RLock()
	defer m.mu.RUnlock()

	if fi, ok := m.info[key]; ok {
		return fi, nil
	}

	return nil, vfserror.ErrNotFound
}