ezshare/store/bolt.go
2021-12-08 06:17:11 +01:00

360 lines
7.8 KiB
Go

package store
import (
"crypto/ecdsa"
"crypto/x509"
"fmt"
"strings"
"gitea.benny.dog/torjus/ezshare/pb"
"github.com/google/uuid"
bolt "go.etcd.io/bbolt"
"google.golang.org/protobuf/proto"
)
var _ FileStore = &BoltStore{}
type BoltStore struct {
db *bolt.DB
}
var bktKey = []byte("files")
var bktKeyCerts = []byte("certs")
var bktKeyKeys = []byte("keys")
var bktKeyUsers = []byte("users")
var bktKeyRevoked = []byte("revoked")
var bktKeyBinary = []byte("binaries")
func NewBoltStore(path string) (*BoltStore, error) {
s := &BoltStore{}
db, err := bolt.Open(path, 0660, &bolt.Options{})
if err != nil {
return nil, fmt.Errorf("unable to open store: %w", err)
}
s.db = db
err = db.Update(func(t *bolt.Tx) error {
if _, err := t.CreateBucketIfNotExists(bktKey); err != nil {
return err
}
if _, err := t.CreateBucketIfNotExists(bktKeyCerts); err != nil {
return err
}
if _, err := t.CreateBucketIfNotExists(bktKeyKeys); err != nil {
return err
}
if _, err := t.CreateBucketIfNotExists(bktKeyUsers); err != nil {
return err
}
if _, err := t.CreateBucketIfNotExists(bktKeyRevoked); err != nil {
return err
}
if _, err := t.CreateBucketIfNotExists(bktKeyBinary); err != nil {
return err
}
return nil
})
if err != nil {
return nil, fmt.Errorf("unable to create bucket: %w", err)
}
return s, nil
}
func (s *BoltStore) Close() error {
return s.db.Close()
}
func (s *BoltStore) GetFile(id string) (*pb.File, error) {
var file pb.File
var data []byte
if err := s.db.View(func(t *bolt.Tx) error {
bkt := t.Bucket(bktKey)
data = bkt.Get([]byte(id))
return nil
}); err != nil {
return nil, fmt.Errorf("error getting file: %w", err)
}
if data == nil {
return nil, ErrNoSuchItem
}
err := proto.Unmarshal(data, &file)
if err != nil {
return nil, fmt.Errorf("unable to unmarshal file: %w", err)
}
return &file, nil
}
func (s *BoltStore) StoreFile(file *pb.File) (string, error) {
file.FileId = uuid.Must(uuid.NewRandom()).String()
data, err := proto.Marshal(file)
if err != nil {
return "", fmt.Errorf("unable to marshal file: %w", err)
}
err = s.db.Update(func(t *bolt.Tx) error {
bkt := t.Bucket(bktKey)
return bkt.Put([]byte(file.FileId), data)
})
if err != nil {
return "", fmt.Errorf("unable to store file: %w", err)
}
return file.FileId, nil
}
func (s *BoltStore) DeleteFile(id string) error {
return s.db.Update(func(t *bolt.Tx) error {
bkt := t.Bucket(bktKey)
data := bkt.Get([]byte(id))
if data == nil {
return ErrNoSuchItem
}
return bkt.Delete([]byte(id))
})
}
func (s *BoltStore) ListFiles() ([]*pb.ListFilesResponse_ListFileInfo, error) {
var response []*pb.ListFilesResponse_ListFileInfo
err := s.db.View(func(t *bolt.Tx) error {
bkt := t.Bucket(bktKey)
return bkt.ForEach(func(k, v []byte) error {
var f pb.File
err := proto.Unmarshal(v, &f)
if err != nil {
return err
}
response = append(response, &pb.ListFilesResponse_ListFileInfo{FileId: f.FileId, Metadata: f.Metadata})
return nil
})
})
if err != nil {
return nil, err
}
return response, nil
}
// Certificate store
var _ CertificateStore = &BoltStore{}
func (s *BoltStore) GetCertificate(serial string) (*x509.Certificate, error) {
var raw []byte
err := s.db.View(func(t *bolt.Tx) error {
bkt := t.Bucket(bktKeyCerts)
raw = bkt.Get([]byte(serial))
return nil
})
if err != nil {
return nil, err
}
if raw == nil {
return nil, ErrNoSuchItem
}
cert, err := x509.ParseCertificate(raw)
if err != nil {
return nil, fmt.Errorf("unable to parse certificate: %w", err)
}
return cert, nil
}
func (s *BoltStore) StoreCertificate(cert *x509.Certificate) error {
data := make([]byte, len(cert.Raw))
copy(data, cert.Raw)
return s.db.Update(func(t *bolt.Tx) error {
bkt := t.Bucket(bktKeyCerts)
return bkt.Put([]byte(cert.SerialNumber.String()), cert.Raw)
})
}
func (s *BoltStore) GetKey(id string) (*ecdsa.PrivateKey, error) {
var data []byte
err := s.db.View(func(t *bolt.Tx) error {
bkt := t.Bucket(bktKeyKeys)
data = bkt.Get([]byte(id))
return nil
})
if err != nil {
return nil, err
}
return x509.ParseECPrivateKey(data)
}
func (s *BoltStore) StoreKey(id string, key *ecdsa.PrivateKey) error {
data, err := x509.MarshalECPrivateKey(key)
if err != nil {
return fmt.Errorf("unable to marshal key: %w", err)
}
return s.db.Update(func(t *bolt.Tx) error {
bkt := t.Bucket(bktKeyKeys)
return bkt.Put([]byte(id), data)
})
}
func (s *BoltStore) ListCertificates() ([]string, error) {
var ids []string
err := s.db.View(func(tx *bolt.Tx) error {
bkt := tx.Bucket(bktKeyCerts)
return bkt.ForEach(func(k, v []byte) error {
ids = append(ids, string(k))
return nil
})
})
if err != nil {
return nil, err
}
return ids, nil
}
func (s *BoltStore) Revoke(serial string) error {
return s.db.Update(func(tx *bolt.Tx) error {
bkt := tx.Bucket(bktKeyRevoked)
return bkt.Put([]byte(serial), []byte{'r'})
})
}
func (s *BoltStore) IsRevoked(serial string) (bool, error) {
var revoked bool
err := s.db.View(func(tx *bolt.Tx) error {
bkt := tx.Bucket(bktKeyRevoked)
status := bkt.Get([]byte(serial))
if status != nil {
revoked = true
}
return nil
})
if err != nil {
return false, err
}
return revoked, nil
}
var _ UserStore = &BoltStore{}
func (s *BoltStore) StoreUser(user *pb.User) error {
return s.db.Update(func(tx *bolt.Tx) error {
bkt := tx.Bucket(bktKeyUsers)
data, err := proto.Marshal(user)
if err != nil {
return err
}
return bkt.Put([]byte(user.Id), data)
})
}
func (s *BoltStore) GetUser(id string) (*pb.User, error) {
var data []byte
err := s.db.View(func(tx *bolt.Tx) error {
bkt := tx.Bucket(bktKeyUsers)
data = bkt.Get([]byte(id))
return nil
})
if err != nil {
return nil, err
}
var user pb.User
err = proto.Unmarshal(data, &user)
return &user, err
}
func (s *BoltStore) ListUsers() ([]string, error) {
var ids []string
err := s.db.View(func(tx *bolt.Tx) error {
bkt := tx.Bucket(bktKeyUsers)
return bkt.ForEach(func(k, _ []byte) error {
ids = append(ids, string(k))
return nil
})
})
return ids, err
}
func (s *BoltStore) GetUserByUsername(username string) (*pb.User, error) {
var user pb.User
err := s.db.View(func(tx *bolt.Tx) error {
bkt := tx.Bucket(bktKeyUsers)
c := bkt.Cursor()
for k, v := c.First(); k != nil; c.Next() {
err := proto.Unmarshal(v, &user)
if err != nil {
// TODO: Log that db has invalid user
continue
}
if user.Username == username {
return nil
}
}
return nil
})
if err != nil {
return nil, err
}
if user.Username == username {
return &user, nil
}
return nil, ErrNoSuchItem
}
/////////////////
// BinaryStore //
/////////////////
var _ BinaryStore = &BoltStore{}
func binaryKey(version string, os string, arch string) []byte {
return []byte(fmt.Sprintf("%s,%s,%s", version, os, arch))
}
func (s *BoltStore) StoreBinary(release *pb.Binary) error {
return s.db.Update(func(tx *bolt.Tx) error {
bkt := tx.Bucket(bktKeyBinary)
key := binaryKey(release.Version, release.Os, release.Arch)
return bkt.Put(key, release.Data)
})
}
func (s *BoltStore) GetBinary(version string, os string, arch string) (*pb.Binary, error) {
binary := &pb.Binary{}
err := s.db.View(func(tx *bolt.Tx) error {
bkt := tx.Bucket(bktKeyBinary)
key := binaryKey(version, os, arch)
binary.Data = bkt.Get(key)
return nil
})
if err != nil {
return nil, err
}
if binary.Data == nil {
return nil, ErrNoSuchItem
}
return binary, nil
}
func (s *BoltStore) List() ([]string, error) {
var releases []string
err := s.db.View(func(tx *bolt.Tx) error {
bkt := tx.Bucket(bktKeyBinary)
return bkt.ForEach(func(k, _ []byte) error {
key := strings.Split(string(k), ",")
releases = append(releases, fmt.Sprintf("ezshare-%s-%s-%s", key[0][1:], key[1], key[2]))
return nil
})
})
if err != nil {
return nil, err
}
return releases, nil
}