ezshare/store/memory.go

package store

import (
	"crypto/ecdsa"
	"crypto/sha256"
	"crypto/x509"
	"fmt"
	"strings"
	"sync"

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

type MemoryStore struct {
	// Filestore
	filesLock sync.RWMutex
	files     map[string]*pb.File
	// Certstore
	certLock     sync.RWMutex
	certs        map[string][]byte
	keyLock      sync.RWMutex
	keys         map[string][]byte
	revokedCerts map[string]struct{}
	revokedLock  sync.RWMutex
	// Userstore
	usersLock sync.RWMutex
	users     map[string]*pb.User
	// Binarystore
	binaryLock sync.RWMutex
	binaries   map[[32]byte]*pb.Binary
}

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),
		revokedCerts: make(map[string]struct{}),
		binaries:     make(map[[32]byte]*pb.Binary),
	}
}

///////////////
// 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, ErrNoSuchItem
}

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 ErrNoSuchItem
	}

	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(serial string) (*x509.Certificate, error) {
	s.certLock.Lock()
	defer s.certLock.Unlock()

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

	return x509.ParseCertificate(data)
}

func (s *MemoryStore) StoreCertificate(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[cert.SerialNumber.String()] = 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, ErrNoSuchItem
	}

	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
}

func (s *MemoryStore) Revoke(serial string) error {
	s.revokedLock.Lock()
	defer s.revokedLock.Unlock()

	s.revokedCerts[serial] = struct{}{}
	return nil
}

func (s *MemoryStore) IsRevoked(serial string) (bool, error) {
	s.revokedLock.RLock()
	defer s.revokedLock.RUnlock()

	_, revoked := s.revokedCerts[serial]
	return revoked, 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, ErrNoSuchItem
	}
	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, ErrNoSuchItem
}

// Binary store
func (s *MemoryStore) StoreBinary(release *pb.Binary) error {
	s.binaryLock.Lock()
	defer s.binaryLock.Unlock()

	hasher := sha256.New()
	_, err := hasher.Write(release.Data)
	if err != nil {
		return fmt.Errorf("unable to hash binary: %w", err)
	}
	sum := hasher.Sum(nil)
	var byteSum [32]byte
	copy(byteSum[:], sum[:32])

	s.binaries[byteSum] = release
	return nil
}

func (s *MemoryStore) GetBinary(version string, os string, arch string) (*pb.Binary, error) {
	s.binaryLock.RLock()
	defer s.binaryLock.RUnlock()

	for _, release := range s.binaries {
		if release.Arch == arch && release.Os == os {
			if release.Version == version || strings.EqualFold(os, "latest") {
				return release, nil
			}
		}
	}

	return nil, ErrNoSuchItem
}