motr-enclave/internal/keybase/conn.go

package keybase

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"sync"

	"enclave/internal/migrations"

	"github.com/ncruces/go-sqlite3"
	"github.com/ncruces/go-sqlite3/driver"
	_ "github.com/ncruces/go-sqlite3/embed"
	"github.com/ncruces/go-sqlite3/ext/hash"
	"github.com/ncruces/go-sqlite3/ext/serdes"
	"github.com/ncruces/go-sqlite3/ext/uuid"
)

// Keybase encapsulates the encrypted key storage database.
type Keybase struct {
	db      *sql.DB
	conn    *sqlite3.Conn // raw connection for serdes
	queries *Queries
	did     string
	didID   int64
	mu      sync.RWMutex
}

var (
	instance *Keybase
	initMu   sync.Mutex
)

// Open creates or returns the singleton Keybase instance with an in-memory database.
func Open() (*Keybase, error) {
	initMu.Lock()
	defer initMu.Unlock()

	if instance != nil {
		return instance, nil
	}

	var rawConn *sqlite3.Conn
	initCallback := func(conn *sqlite3.Conn) error {
		rawConn = conn
		if err := hash.Register(conn); err != nil {
			return fmt.Errorf("register hash extension: %w", err)
		}
		if err := uuid.Register(conn); err != nil {
			return fmt.Errorf("register uuid extension: %w", err)
		}
		return nil
	}

	db, err := driver.Open(":memory:", initCallback)
	if err != nil {
		return nil, fmt.Errorf("keybase: open database: %w", err)
	}

	if _, err := db.Exec(migrations.SchemaSQL); err != nil {
		db.Close()
		return nil, fmt.Errorf("keybase: init schema: %w", err)
	}

	instance = &Keybase{
		db:      db,
		conn:    rawConn,
		queries: New(db),
	}

	return instance, nil
}

// Get returns the existing Keybase instance or nil if not initialized.
func Get() *Keybase {
	initMu.Lock()
	defer initMu.Unlock()
	return instance
}

// MustGet returns the existing Keybase instance or panics if not initialized.
func MustGet() *Keybase {
	kb := Get()
	if kb == nil {
		panic("keybase: not initialized")
	}
	return kb
}

// Close closes the database connection and clears the singleton.
func Close() error {
	initMu.Lock()
	defer initMu.Unlock()

	if instance == nil {
		return nil
	}

	err := instance.db.Close()
	instance = nil
	return err
}

// Reset clears the singleton instance (useful for testing).
func Reset() {
	initMu.Lock()
	defer initMu.Unlock()

	if instance != nil {
		instance.db.Close()
		instance = nil
	}
}

// DB returns the underlying sql.DB connection.
func (k *Keybase) DB() *sql.DB {
	k.mu.RLock()
	defer k.mu.RUnlock()
	return k.db
}

// Queries returns the SQLC-generated query interface.
func (k *Keybase) Queries() *Queries {
	k.mu.RLock()
	defer k.mu.RUnlock()
	return k.queries
}

// DID returns the current DID identifier.
func (k *Keybase) DID() string {
	k.mu.RLock()
	defer k.mu.RUnlock()
	return k.did
}

// DIDID returns the database ID of the current DID.
func (k *Keybase) DIDID() int64 {
	k.mu.RLock()
	defer k.mu.RUnlock()
	return k.didID
}

// IsInitialized returns true if a DID has been set.
func (k *Keybase) IsInitialized() bool {
	k.mu.RLock()
	defer k.mu.RUnlock()
	return k.did != ""
}

// SetDID sets the current DID context.
func (k *Keybase) SetDID(did string, didID int64) {
	k.mu.Lock()
	defer k.mu.Unlock()
	k.did = did
	k.didID = didID
}

// Initialize creates a new DID document from a WebAuthn credential.
func (k *Keybase) Initialize(ctx context.Context, credentialBytes []byte) (string, error) {
	k.mu.Lock()
	defer k.mu.Unlock()

	did := fmt.Sprintf("did:sonr:%x", credentialBytes[:16])
	docJSON, _ := json.Marshal(map[string]any{
		"@context": []string{"https://www.w3.org/ns/did/v1"},
		"id":       did,
	})

	doc, err := k.queries.CreateDID(ctx, CreateDIDParams{
		Did:        did,
		Controller: did,
		Document:   docJSON,
		Sequence:   0,
	})
	if err != nil {
		return "", fmt.Errorf("keybase: create DID: %w", err)
	}

	k.did = did
	k.didID = doc.ID

	return did, nil
}

// Load restores the database state from serialized bytes and sets the current DID.
func (k *Keybase) Load(ctx context.Context, data []byte) (string, error) {
	if len(data) < 100 {
		return "", fmt.Errorf("keybase: invalid database format")
	}

	k.mu.Lock()
	defer k.mu.Unlock()

	if k.conn == nil {
		return "", fmt.Errorf("keybase: database not initialized")
	}

	if err := serdes.Deserialize(k.conn, "main", data); err != nil {
		return "", fmt.Errorf("keybase: deserialize database: %w", err)
	}

	docs, err := k.queries.ListAllDIDs(ctx)
	if err != nil {
		return "", fmt.Errorf("keybase: list DIDs: %w", err)
	}

	if len(docs) == 0 {
		return "", fmt.Errorf("keybase: no DID found in database")
	}

	k.did = docs[0].Did
	k.didID = docs[0].ID

	return k.did, nil
}

// Serialize exports the database state as bytes using native SQLite serialization.
func (k *Keybase) Serialize() ([]byte, error) {
	k.mu.RLock()
	defer k.mu.RUnlock()

	if k.conn == nil {
		return nil, fmt.Errorf("keybase: database not initialized")
	}

	return serdes.Serialize(k.conn, "main")
}

func (k *Keybase) RestoreFromDump(data []byte) error {
	k.mu.Lock()
	defer k.mu.Unlock()

	if k.conn == nil {
		return fmt.Errorf("keybase: database not initialized")
	}

	if err := serdes.Deserialize(k.conn, "main", data); err != nil {
		return fmt.Errorf("keybase: deserialize database: %w", err)
	}

	docs, err := k.queries.ListAllDIDs(context.Background())
	if err != nil {
		return fmt.Errorf("keybase: failed to list DIDs: %w", err)
	}

	if len(docs) > 0 {
		k.did = docs[0].Did
		k.didID = docs[0].ID
	}

	return nil
}

// WithTx executes a function within a database transaction.
func (k *Keybase) WithTx(ctx context.Context, fn func(*Queries) error) error {
	tx, err := k.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("keybase: begin tx: %w", err)
	}

	if err := fn(k.queries.WithTx(tx)); err != nil {
		tx.Rollback()
		return err
	}

	return tx.Commit()
}