ezshare/store/memory.go

package store

import (
	"crypto/ecdsa"
	"crypto/x509"
	"sync"

	"gitea.benny.dog/torjus/ezshare/pb"
	"github.com/google/uuid"
)

type MemoryStore struct {
	filesLock sync.RWMutex
	files     map[string]*pb.File
	certLock  sync.RWMutex
	certs     map[string][]byte
	keyLock   sync.RWMutex
	keys      map[string][]byte
	usersLock sync.RWMutex
	users     map[string]*pb.User
}

func NewMemoryStore() *MemoryStore {
	return &MemoryStore{
		files: make(map[string]*pb.File),
		certs: make(map[string][]byte),
		keys:  make(map[string][]byte),
		users: make(map[string]*pb.User),
	}
}

///////////////
// FileStore //
///////////////

var _ FileStore = &MemoryStore{}

func (s *MemoryStore) GetFile(id string) (*pb.File, error) {
	s.filesLock.RLock()
	defer s.filesLock.RUnlock()

	if file, ok := s.files[id]; ok {
		return file, nil
	}

	return nil, ErrNoSuchFile
}

func (s *MemoryStore) StoreFile(file *pb.File) (string, error) {
	s.filesLock.Lock()
	defer s.filesLock.Unlock()

	id := uuid.Must(uuid.NewRandom()).String()
	file.FileId = id

	s.files[id] = file

	return id, nil
}

func (s *MemoryStore) DeleteFile(id string) error {
	s.filesLock.Lock()
	defer s.filesLock.Unlock()
	if _, ok := s.files[id]; !ok {
		return ErrNoSuchFile
	}

	delete(s.files, id)
	return nil
}

func (s *MemoryStore) ListFiles() ([]*pb.ListFilesResponse_ListFileInfo, error) {
	s.filesLock.RLock()
	defer s.filesLock.RUnlock()

	var response []*pb.ListFilesResponse_ListFileInfo

	for _, f := range s.files {
		response = append(response, &pb.ListFilesResponse_ListFileInfo{FileId: f.FileId, Metadata: f.Metadata})
	}

	return response, nil
}

//////////////////////
// CertificateStore //
//////////////////////

var _ CertificateStore = &MemoryStore{}

func (s *MemoryStore) GetCertificate(id string) (*x509.Certificate, error) {
	s.certLock.Lock()
	defer s.certLock.Unlock()

	data, ok := s.certs[id]
	if !ok {
		// TODO: Make separate error, or rename error
		return nil, ErrNoSuchFile
	}

	return x509.ParseCertificate(data)
}

func (s *MemoryStore) StoreCertificate(id string, cert *x509.Certificate) error {
	s.certLock.Lock()
	defer s.certLock.Unlock()

	// Copy cert data
	data := make([]byte, len(cert.Raw))
	copy(data, cert.Raw)

	s.certs[id] = data
	return nil
}

func (s *MemoryStore) GetKey(id string) (*ecdsa.PrivateKey, error) {
	s.keyLock.RLock()
	defer s.keyLock.RUnlock()
	data, ok := s.keys[id]
	if !ok {
		return nil, ErrNoSuchFile
	}

	return x509.ParseECPrivateKey(data)
}

func (s *MemoryStore) StoreKey(id string, key *ecdsa.PrivateKey) error {
	s.keyLock.Lock()
	defer s.keyLock.Unlock()

	data, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return err
	}

	s.keys[id] = data
	return nil
}

func (s *MemoryStore) ListCertificates() ([]string, error) {
	s.certLock.RLock()
	defer s.certLock.RUnlock()
	var certIDs []string
	for key := range s.certs {
		certIDs = append(certIDs, key)
	}
	return certIDs, nil
}

///////////////
// UserStore //
///////////////

var _ UserStore = &MemoryStore{}

func (s *MemoryStore) StoreUser(user *pb.User) error {
	s.usersLock.Lock()
	defer s.usersLock.Unlock()

	s.users[user.Id] = user
	return nil
}

func (s *MemoryStore) GetUser(id string) (*pb.User, error) {
	s.usersLock.RLock()
	defer s.usersLock.RUnlock()
	user, ok := s.users[id]
	if !ok {
		// TODO: Update error
		return nil, ErrNoSuchFile
	}
	return user, nil
}

func (s *MemoryStore) ListUsers() ([]string, error) {
	s.usersLock.RLock()
	defer s.usersLock.RUnlock()

	var ids []string
	for id := range s.users {
		ids = append(ids, id)
	}
	return ids, nil
}

func (s *MemoryStore) GetUserByUsername(username string) (*pb.User, error) {
	s.usersLock.RLock()
	defer s.usersLock.RUnlock()
	for _, user := range s.users {
		if user.Username == username {
			return user, nil
		}
	}
	return nil, ErrNoSuchFile
}