ceremonyclient/alias/store.go

package aliases

import (
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"

	"github.com/pkg/errors"
	"gopkg.in/yaml.v3"
)

// AddressBytes encodes to YAML as a hex string.
type AddressBytes []byte

func (a AddressBytes) MarshalYAML() (any, error) {
	return fmt.Sprintf("%x", []byte(a)), nil
}

func (a *AddressBytes) UnmarshalYAML(node *yaml.Node) error {
	if node.Kind != yaml.ScalarNode {
		return errors.Wrap(
			fmt.Errorf("address must be a scalar"),
			"unmarshal yaml",
		)
	}
	b, err := parseAddressLiteral(node.Value)
	if err != nil {
		return err
	}
	*a = AddressBytes(b)
	return nil
}

type Alias struct {
	Address AddressBytes `yaml:"address"`
	Type    string       `yaml:"type,omitempty"`
}

func (a *Alias) UnmarshalYAML(node *yaml.Node) error {
	switch node.Kind {
	case yaml.ScalarNode:
		var addr AddressBytes
		if err := node.Decode(&addr); err != nil {
			return err
		}
		*a = Alias{Address: addr}
		return nil
	case yaml.MappingNode:
		type alias Alias
		var tmp alias
		if err := node.Decode(&tmp); err != nil {
			return err
		}
		*a = Alias(tmp)
		return nil
	default:
		return errors.Wrap(
			fmt.Errorf("alias must be a scalar or mapping"),
			"unmarshal yaml",
		)
	}
}

type File struct {
	Aliases map[string]Alias `yaml:"aliases"`
}

// Store keeps aliases in memory and persists to disk.
type Store struct {
	mu   sync.Mutex
	data File
	path string // file path for autosave
}

// NewInMemory creates an empty store without a path (no autosave).
func NewInMemory() *Store {
	return &Store{data: File{Aliases: map[string]Alias{}}}
}

// NewOnDisk creates (or loads) a file-backed store with no aliases yet.
func NewOnDisk(path string) (*Store, error) {
	// Try to load if present
	if st, err := Load(path); err == nil {
		return st, nil
	} else if !os.IsNotExist(err) {
		return nil, errors.Wrap(err, "new on disk")
	}

	// Not found
	err := os.MkdirAll(filepath.Dir(path), 0o755)
	if err != nil && !os.IsExist(err) {
		return nil, err
	}
	s := &Store{data: File{Aliases: map[string]Alias{}}, path: path}
	if err := s.saveLocked(); err != nil {
		return nil, err
	}
	return s, nil
}

// Load reads from path and returns a file-backed store (autosave enabled).
func Load(path string) (*Store, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return LoadFromReader(path, f)
}

// LoadFromReader reads a store from r; if path != "" autosave is enabled.
func LoadFromReader(path string, r io.Reader) (*Store, error) {
	var file File
	dec := yaml.NewDecoder(r)
	dec.KnownFields(true)
	if err := dec.Decode(&file); err != nil {
		return nil, err
	}
	if file.Aliases == nil {
		file.Aliases = make(map[string]Alias)
	}
	return &Store{data: file, path: path}, nil
}

// Save writes the store to its path (or to provided path if non-empty).
func (s *Store) Save(path string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if path != "" {
		s.path = path
	}
	return s.saveLocked()
}

// Put inserts or replaces an alias and autosaves.
func (s *Store) Put(name string, addr []byte, typeHint string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.data.Aliases == nil {
		s.data.Aliases = make(map[string]Alias)
	}
	s.data.Aliases[name] = Alias{
		Address: AddressBytes(append([]byte(nil), addr...)),
		Type:    typeHint,
	}
	return s.saveLocked()
}

// Remove deletes an alias (if present) and autosaves when a deletion happens.
// Returns (deleted, error).
func (s *Store) Delete(name string) (bool, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.data.Aliases[name]; !ok {
		return false, nil
	}
	delete(s.data.Aliases, name)
	return true, s.saveLocked()
}

// List returns sorted alias names.
func (s *Store) List() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]string, 0, len(s.data.Aliases))
	for k := range s.data.Aliases {
		out = append(out, k)
	}
	sort.Strings(out)
	return out
}

// Get returns (addr, type, ok).
func (s *Store) Get(name string) ([]byte, string, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	al, ok := s.data.Aliases[name]
	if !ok {
		return nil, "", false
	}
	return append([]byte(nil), al.Address...), al.Type, true
}

// FindByAddress returns (name, type, ok) for exact byte match.
func (s *Store) FindByAddress(addr []byte) (string, string, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for k, v := range s.data.Aliases {
		if bytesEqual(v.Address, addr) {
			return k, v.Type, true
		}
	}
	return "", "", false
}

// Resolve: alias name -> (addr, type), else parse literal hex -> (addr, "")
func (s *Store) Resolve(key string) ([]byte, string, bool) {
	if addr, typ, ok := s.Get(key); ok {
		return addr, typ, true
	}
	if b, err := parseAddressLiteral(key); err == nil {
		return b, "", true
	}
	return nil, "", false
}

func (s *Store) saveLocked() error {
	if s.path == "" {
		return errors.New(
			"no path set for autosave; call Save(path) once or use NewOnDisk/Load",
		)
	}
	tmp := s.path + ".tmp"
	f, err := os.Create(tmp)
	if err != nil {
		return err
	}
	enc := yaml.NewEncoder(f)
	enc.SetIndent(2)
	if err := enc.Encode(&s.data); err != nil {
		f.Close()
		_ = os.Remove(tmp)
		return err
	}
	if err := enc.Close(); err != nil {
		f.Close()
		_ = os.Remove(tmp)
		return err
	}
	if err := f.Close(); err != nil {
		_ = os.Remove(tmp)
		return err
	}
	return os.Rename(tmp, s.path)
}

func bytesEqual(a, b []byte) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

func parseAddressLiteral(s string) ([]byte, error) {
	t := strings.TrimSpace(s)
	if t == "" {
		return nil, errors.New("empty address")
	}

	return hex.DecodeString(t)
}