motr-enclave/internal/keybase/conn.go

// Package keybase contains the SQLite database for cryptographic keys.
package keybase

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

	"enclave/internal/migrations"

	_ "github.com/ncruces/go-sqlite3/driver"
	_ "github.com/ncruces/go-sqlite3/embed"
)

// Keybase encapsulates the encrypted key storage database.
type Keybase struct {
	db      *sql.DB
	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
	}

	conn, err := sql.Open("sqlite3", ":memory:")
	if err != nil {
		return nil, fmt.Errorf("keybase: open database: %w", err)
	}

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

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

	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) < 10 {
		return "", fmt.Errorf("keybase: invalid database format")
	}

	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.mu.Lock()
	k.did = docs[0].Did
	k.didID = docs[0].ID
	k.mu.Unlock()

	return k.did, nil
}

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

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

	return k.exportDump()
}

func (k *Keybase) exportDump() ([]byte, error) {
	var dump strings.Builder
	dump.WriteString(migrations.SchemaSQL + "\n")

	tables := []string{
		"did_documents", "verification_methods", "credentials",
		"key_shares", "accounts", "ucan_tokens", "ucan_revocations",
		"sessions", "services", "grants", "delegations", "sync_checkpoints",
	}

	for _, table := range tables {
		if err := k.exportTable(&dump, table); err != nil {
			continue
		}
	}

	return []byte(dump.String()), nil
}

func (k *Keybase) exportTable(dump *strings.Builder, table string) error {
	rows, err := k.db.Query(fmt.Sprintf("SELECT * FROM %s", table))
	if err != nil {
		return err
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		return err
	}

	for rows.Next() {
		values := make([]any, len(cols))
		valuePtrs := make([]any, len(cols))
		for i := range values {
			valuePtrs[i] = &values[i]
		}

		if err := rows.Scan(valuePtrs...); err != nil {
			continue
		}

		dump.WriteString(fmt.Sprintf("INSERT INTO %s (", table))
		dump.WriteString(strings.Join(cols, ", "))
		dump.WriteString(") VALUES (")

		for i, val := range values {
			if i > 0 {
				dump.WriteString(", ")
			}
			dump.WriteString(formatSQLValue(val))
		}
		dump.WriteString(");\n")
	}

	return rows.Err()
}

func formatSQLValue(val any) string {
	if val == nil {
		return "NULL"
	}

	switch v := val.(type) {
	case int64:
		return fmt.Sprintf("%d", v)
	case float64:
		return fmt.Sprintf("%f", v)
	case bool:
		if v {
			return "1"
		}
		return "0"
	case []byte:
		return fmt.Sprintf("'%s'", escapeSQLString(string(v)))
	case string:
		return fmt.Sprintf("'%s'", escapeSQLString(v))
	default:
		return fmt.Sprintf("'%s'", escapeSQLString(fmt.Sprintf("%v", v)))
	}
}

func escapeSQLString(s string) string {
	return strings.ReplaceAll(s, "'", "''")
}

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

	statements := strings.Split(string(data), ";\n")
	for _, stmt := range statements {
		stmt = strings.TrimSpace(stmt)
		if stmt == "" || strings.HasPrefix(stmt, "--") {
			continue
		}
		if strings.HasPrefix(stmt, "INSERT INTO") {
			if _, err := k.db.Exec(stmt); err != nil {
				return fmt.Errorf("keybase: failed to execute statement: %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()
}