ezshare/store/bolt.go

package store

import (
	"crypto/ecdsa"
	"crypto/x509"
	"fmt"

	"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")

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
		}
		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(id string) (*x509.Certificate, error) {
	var raw []byte
	err := s.db.View(func(t *bolt.Tx) error {
		bkt := t.Bucket(bktKeyCerts)

		raw = bkt.Get([]byte(id))
		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(id string, 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(id), 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
}

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
}