qcc/internal/identity/manager.go

package identity

import (
	"encoding/binary"
	"fmt"
	"time"

	"go.uber.org/zap"

	"github.com/niko/qcc/internal/store"
	"github.com/niko/qcc/pkg/types"
	"golang.org/x/crypto/blake2b"
)

type Identity struct {
	Number     string
	PubKey     [32]byte
	CertDER    []byte
	CreatedAt  int64
	LastReroll int64
}

type Manager struct {
	store   *store.Table
	masterKey [32]byte
	logger  *zap.Logger
	prefix  string
	cooldown time.Duration
}

func NewManager(engine *store.Engine, masterKey [32]byte, prefix string, cooldown time.Duration, logger *zap.Logger) (*Manager, error) {
	t, err := engine.GetTable("identities")
	if err != nil {
		return nil, fmt.Errorf("get identities table: %w", err)
	}
	store.SetMasterKey(masterKey)
	return &Manager{
		store:     t,
		masterKey: masterKey,
		logger:    logger,
		prefix:    prefix,
		cooldown:  cooldown,
	}, nil
}

func (m *Manager) Allocate(pubKey [32]byte) (*Identity, error) {
	var number string
	for attempts := 0; attempts < 100; attempts++ {
		number = GenerateNumber(m.prefix)
		keyHash := numberKeyHash(number)
		if !m.store.Has(keyHash) {
			break
		}
	}

	keyHash := numberKeyHash(number)
	if m.store.Has(keyHash) {
		return nil, fmt.Errorf("failed to generate unique number")
	}

	now := time.Now().UnixNano()
	ident := &Identity{
		Number:     number,
		PubKey:     pubKey,
		CreatedAt:  now,
		LastReroll: now,
	}

	payload := m.marshalIdentity(ident)
	if err := m.store.Insert(keyHash, types.RecIdentity, 0, payload, m.masterKey); err != nil {
		return nil, fmt.Errorf("store identity: %w", err)
	}

	return ident, nil
}

func (m *Manager) Reroll(pubKey [32]byte) (*Identity, error) {
	var foundIdent *Identity
	var foundHash uint64

	err := m.store.Iterate(func(keyHash uint64, payload []byte) bool {
		ident := m.unmarshalIdentity(payload)
		if ident != nil && ident.PubKey == pubKey {
			foundIdent = ident
			foundHash = keyHash
			return false
		}
		return true
	}, m.masterKey)

	if err != nil {
		return nil, err
	}
	if foundIdent == nil {
		return nil, fmt.Errorf("identity not found for pubkey")
	}

	elapsed := time.Since(time.Unix(0, foundIdent.LastReroll))
	if elapsed < m.cooldown {
		retryAfter := m.cooldown - elapsed
		return nil, &CooldownError{RetryAfter: retryAfter}
	}

	m.store.Delete(foundHash)

	number := GenerateNumber(m.prefix)
	keyHash := numberKeyHash(number)
	now := time.Now().UnixNano()

	newIdent := &Identity{
		Number:     number,
		PubKey:     pubKey,
		CreatedAt:  foundIdent.CreatedAt,
		LastReroll: now,
	}

	payload := m.marshalIdentity(newIdent)
	if err := m.store.Insert(keyHash, types.RecIdentity, 0, payload, m.masterKey); err != nil {
		return nil, err
	}

	return newIdent, nil
}

func (m *Manager) GetByNumber(number string) (*Identity, error) {
	keyHash := numberKeyHash(number)
	payload, err := m.store.Get(keyHash, m.masterKey)
	if err != nil {
		return nil, fmt.Errorf("identity not found: %s", number)
	}
	ident := m.unmarshalIdentity(payload)
	if ident == nil {
		return nil, fmt.Errorf("invalid identity data")
	}
	return ident, nil
}

func (m *Manager) Exists(number string) bool {
	return m.store.Has(numberKeyHash(number))
}

func (m *Manager) GetByPubKey(pubKey [32]byte) (*Identity, error) {
	var found *Identity
	err := m.store.Iterate(func(keyHash uint64, payload []byte) bool {
		ident := m.unmarshalIdentity(payload)
		if ident != nil && ident.PubKey == pubKey {
			found = ident
			return false
		}
		return true
	}, m.masterKey)
	if err != nil {
		return nil, err
	}
	if found == nil {
		return nil, fmt.Errorf("identity not found for pubkey")
	}
	return found, nil
}

func numberKeyHash(number string) uint64 {
	h, _ := blake2b.New(8, nil)
	h.Write([]byte(number))
	return binary.BigEndian.Uint64(h.Sum(nil))
}

func (m *Manager) marshalIdentity(ident *Identity) []byte {
	buf := make([]byte, 2+len(ident.Number)+32+2+len(ident.CertDER)+8+8)
	off := 0
	PutString(buf, &off, ident.Number)
	copy(buf[off:off+32], ident.PubKey[:])
	off += 32
	PutBytes(buf, &off, ident.CertDER)
	binary.BigEndian.PutUint64(buf[off:off+8], uint64(ident.CreatedAt))
	off += 8
	binary.BigEndian.PutUint64(buf[off:off+8], uint64(ident.LastReroll))
	return buf
}

func (m *Manager) unmarshalIdentity(data []byte) *Identity {
	if len(data) < 50 {
		return nil
	}
	ident := &Identity{}
	off := 0
	number, ok := GetString(data, &off)
	if !ok {
		return nil
	}
	ident.Number = number
	if off+32 > len(data) {
		return nil
	}
	copy(ident.PubKey[:], data[off:off+32])
	off += 32
	certDER, ok := GetBytes(data, &off)
	if !ok {
		return nil
	}
	ident.CertDER = certDER
	if off+16 > len(data) {
		return nil
	}
	ident.CreatedAt = int64(binary.BigEndian.Uint64(data[off:]))
	off += 8
	ident.LastReroll = int64(binary.BigEndian.Uint64(data[off:]))
	return ident
}

type CooldownError struct {
	RetryAfter time.Duration
}

func (e *CooldownError) Error() string {
	return fmt.Sprintf("cooldown: retry after %s", e.RetryAfter)
}

func PutString(buf []byte, off *int, s string) {
	l := len(s)
	binary.BigEndian.PutUint16(buf[*off:*off+2], uint16(l))
	*off += 2
	copy(buf[*off:*off+l], s)
	*off += l
}

func GetString(data []byte, off *int) (string, bool) {
	if *off+2 > len(data) {
		return "", false
	}
	l := int(binary.BigEndian.Uint16(data[*off:]))
	*off += 2
	if *off+l > len(data) {
		return "", false
	}
	s := string(data[*off : *off+l])
	*off += l
	return s, true
}

func PutBytes(buf []byte, off *int, b []byte) {
	l := len(b)
	binary.BigEndian.PutUint16(buf[*off:*off+2], uint16(l))
	*off += 2
	copy(buf[*off:*off+l], b)
	*off += l
}

func GetBytes(data []byte, off *int) ([]byte, bool) {
	if *off+2 > len(data) {
		return nil, false
	}
	l := int(binary.BigEndian.Uint16(data[*off:]))
	*off += 2
	if *off+l > len(data) {
		return nil, false
	}
	b := make([]byte, l)
	copy(b, data[*off:*off+l])
	*off += l
	return b, true
}