gpaste/files/filestore_memory.go

package files

import (
	"bytes"
	"fmt"
	"io"
	"sync"
	"time"
)

type fileData struct {
	ID   string
	Body bytes.Buffer

	MaxViews  uint
	ExpiresOn time.Time
	FileSize  int64
}

type MemoryFileStore struct {
	lock sync.RWMutex
	data map[string]*fileData
}

func NewMemoryFileStore() *MemoryFileStore {
	return &MemoryFileStore{
		data: make(map[string]*fileData),
	}
}

func (s *MemoryFileStore) Store(f *File) error {
	data := &fileData{
		ID:        f.ID,
		MaxViews:  f.MaxViews,
		ExpiresOn: f.ExpiresOn,
	}

	n, err := io.Copy(&data.Body, f.Body)
	_ = f.Body.Close()

	data.FileSize = n

	s.lock.Lock()
	defer s.lock.Unlock()

	s.data[f.ID] = data

	return err
}

func (s *MemoryFileStore) Get(id string) (*File, error) {
	s.lock.RLock()
	defer s.lock.RUnlock()

	fd, ok := s.data[id]
	if !ok {
		return nil, fmt.Errorf("no such item")
	}

	body := new(bytes.Buffer)
	if _, err := body.Write(fd.Body.Bytes()); err != nil {
		return nil, err
	}

	f := &File{
		ID:        fd.ID,
		MaxViews:  fd.MaxViews,
		ExpiresOn: fd.ExpiresOn,
		Body:      io.NopCloser(body),
		FileSize:  fd.FileSize,
	}

	return f, nil
}

func (s *MemoryFileStore) Delete(id string) error {
	s.lock.Lock()
	defer s.lock.Unlock()

	delete(s.data, id)

	return nil
}

func (s *MemoryFileStore) List() ([]string, error) {
	ids := make([]string, 0, len(s.data))

	s.lock.RLock()
	defer s.lock.RUnlock()

	for id := range s.data {
		ids = append(ids, id)
	}

	return ids, nil
}