meowlib/client/serverstorage.go

package client

//
// Storage
//
import (
	"crypto/sha256"
	"encoding/json"
	"path/filepath"
	"sync"

	"forge.redroom.link/yves/meowlib"
	"github.com/dgraph-io/badger"
)

type ServerStorage struct {
	DbFile string `json:"db_file,omitempty"`
	mu     sync.Mutex
	db     *badger.DB
}

// open opens the Badger database. Caller must hold mu.
func (ss *ServerStorage) open() error {
	opts := badger.DefaultOptions(filepath.Join(GetConfig().StoragePath, GetConfig().GetIdentity().Uuid, ss.DbFile))
	opts.Logger = nil
	var err error
	ss.db, err = badger.Open(opts)
	return err
}

// close closes the Badger database. Caller must hold mu.
func (ss *ServerStorage) close() {
	ss.db.Close()
}

// StoreServer stores a server in the Badger database with Server.GetUid() as key.
func (ss *ServerStorage) StoreServer(sc *Server) error {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	return ss.storeServerLocked(sc)
}

// storeServerLocked is StoreServer without acquiring the lock. Caller must hold mu.
func (ss *ServerStorage) storeServerLocked(sc *Server) error {
	if err := ss.open(); err != nil {
		return err
	}
	defer ss.close()
	jsonsrv, err := json.Marshal(sc)
	if err != nil {
		return err
	}
	password, err := GetConfig().GetMemPass()
	if err != nil {
		return err
	}
	data, err := meowlib.SymEncrypt(password, jsonsrv)
	if err != nil {
		return err
	}
	shakey := sha256.Sum256([]byte(sc.GetServerCard().GetUid()))
	key := shakey[:]
	return ss.db.Update(func(txn *badger.Txn) error {
		return txn.Set(key, data)
	})
}

// ServerExists checks if a server exists in the Badger database.
func (ss *ServerStorage) ServerExists(sc *Server) (bool, error) {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	return ss.serverExistsLocked(sc)
}

// serverExistsLocked is ServerExists without acquiring the lock. Caller must hold mu.
func (ss *ServerStorage) serverExistsLocked(sc *Server) (bool, error) {
	if err := ss.open(); err != nil {
		return false, err
	}
	defer ss.close()
	shakey := sha256.Sum256([]byte(sc.GetServerCard().GetUid()))
	key := shakey[:]
	err := ss.db.View(func(txn *badger.Txn) error {
		_, err := txn.Get(key)
		return err
	})
	if err != nil {
		return false, nil
	}
	return true, nil
}

// StoreServerIfNotExists stores a server only if it is not already present.
func (ss *ServerStorage) StoreServerIfNotExists(sc *Server) error {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	exists, err := ss.serverExistsLocked(sc)
	if err != nil {
		return err
	}
	if !exists {
		return ss.storeServerLocked(sc)
	}
	return nil
}

// LoadServer loads a Server from the Badger database by uid.
func (ss *ServerStorage) LoadServer(uid string) (*Server, error) {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	var sc Server
	if err := ss.open(); err != nil {
		return nil, err
	}
	defer ss.close()
	password, err := GetConfig().GetMemPass()
	if err != nil {
		return nil, err
	}
	shakey := sha256.Sum256([]byte(uid))
	key := shakey[:]
	err = ss.db.View(func(txn *badger.Txn) error {
		item, err := txn.Get(key)
		if err != nil {
			return err
		}
		return item.Value(func(val []byte) error {
			jsonsrv, err := meowlib.SymDecrypt(password, val)
			if err != nil {
				return err
			}
			return json.Unmarshal(jsonsrv, &sc)
		})
	})
	return &sc, err
}

// DeleteServer deletes a Server from the Badger database by uid.
func (ss *ServerStorage) DeleteServer(uid string) error {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	if err := ss.open(); err != nil {
		return err
	}
	defer ss.close()
	shakey := sha256.Sum256([]byte(uid))
	key := shakey[:]
	return ss.db.Update(func(txn *badger.Txn) error {
		return txn.Delete(key)
	})
}

// LoadAllServers loads all Servers from the Badger database.
func (ss *ServerStorage) LoadAllServers() ([]*Server, error) {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	var scs []*Server
	if err := ss.open(); err != nil {
		return nil, err
	}
	defer ss.close()
	password, err := GetConfig().GetMemPass()
	if err != nil {
		return nil, err
	}
	err = ss.db.View(func(txn *badger.Txn) error {
		opts := badger.DefaultIteratorOptions
		opts.PrefetchSize = 10
		it := txn.NewIterator(opts)
		defer it.Close()
		for it.Rewind(); it.Valid(); it.Next() {
			item := it.Item()
			var sc Server
			err := item.Value(func(val []byte) error {
				jsonsrv, err := meowlib.SymDecrypt(password, val)
				if err != nil {
					return err
				}
				return json.Unmarshal(jsonsrv, &sc)
			})
			if err != nil {
				return err
			}
			scs = append(scs, &sc)
		}
		return nil
	})
	return scs, err
}

// LoadAllServerCards loads all ServerCards from the Badger database.
func (ss *ServerStorage) LoadAllServerCards() ([]*meowlib.ServerCard, error) {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	var scs []*meowlib.ServerCard
	if err := ss.open(); err != nil {
		return nil, err
	}
	defer ss.close()
	password, err := GetConfig().GetMemPass()
	if err != nil {
		return nil, err
	}
	err = ss.db.View(func(txn *badger.Txn) error {
		opts := badger.DefaultIteratorOptions
		opts.PrefetchSize = 10
		it := txn.NewIterator(opts)
		defer it.Close()
		for it.Rewind(); it.Valid(); it.Next() {
			item := it.Item()
			var sc Server
			err := item.Value(func(val []byte) error {
				jsonsrv, err := meowlib.SymDecrypt(password, val)
				if err != nil {
					return err
				}
				return json.Unmarshal(jsonsrv, &sc)
			})
			if err != nil {
				return err
			}
			scs = append(scs, sc.GetServerCard())
		}
		return nil
	})
	return scs, err
}

// LoadServersFromUids loads Servers whose UIDs are in the provided slice.
func (ss *ServerStorage) LoadServersFromUids(uids []string) ([]*Server, error) {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	var scs []*Server
	if err := ss.open(); err != nil {
		return nil, err
	}
	defer ss.close()
	password, err := GetConfig().GetMemPass()
	if err != nil {
		return nil, err
	}
	err = ss.db.View(func(txn *badger.Txn) error {
		for _, uid := range uids {
			shakey := sha256.Sum256([]byte(uid))
			key := shakey[:]
			item, err := txn.Get(key)
			if err != nil {
				return err
			}
			var sc Server
			err = item.Value(func(val []byte) error {
				jsonsrv, err := meowlib.SymDecrypt(password, val)
				if err != nil {
					return err
				}
				return json.Unmarshal(jsonsrv, &sc)
			})
			if err != nil {
				return err
			}
			scs = append(scs, &sc)
		}
		return nil
	})
	return scs, err
}

// LoadServerCardsFromUids loads ServerCards whose UIDs are in the provided slice.
func (ss *ServerStorage) LoadServerCardsFromUids(uids []string) ([]*meowlib.ServerCard, error) {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	var scs []*meowlib.ServerCard
	if err := ss.open(); err != nil {
		return nil, err
	}
	defer ss.close()
	password, err := GetConfig().GetMemPass()
	if err != nil {
		return nil, err
	}
	err = ss.db.View(func(txn *badger.Txn) error {
		for _, uid := range uids {
			shakey := sha256.Sum256([]byte(uid))
			key := shakey[:]
			item, err := txn.Get(key)
			if err != nil {
				return err
			}
			var sc Server
			err = item.Value(func(val []byte) error {
				jsonsrv, err := meowlib.SymDecrypt(password, val)
				if err != nil {
					return err
				}
				return json.Unmarshal(jsonsrv, &sc)
			})
			if err != nil {
				return err
			}
			scs = append(scs, sc.GetServerCard())
		}
		return nil
	})
	return scs, err
}