motr-enclave/internal/keybase/store.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"
)

type Keybase struct {
	db      *sql.DB
	conn    *sqlite3.Conn
	queries *Queries
	did     string
	didID   int64
	mu      sync.RWMutex
}

var (
	instance *Keybase
	initMu   sync.Mutex
)

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("open database: %w", err)
	}

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

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

	if rawConn != nil {
		if err := RegisterMPCFunctions(rawConn); err != nil {
			db.Close()
			instance = nil
			return nil, fmt.Errorf("register mpc functions: %w", err)
		}
	}

	return instance, nil
}

func Get() *Keybase {
	initMu.Lock()
	defer initMu.Unlock()
	return instance
}

func Close() error {
	initMu.Lock()
	defer initMu.Unlock()

	if instance == nil {
		return nil
	}

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

func (k *Keybase) DB() *sql.DB {
	k.mu.RLock()
	defer k.mu.RUnlock()
	return k.db
}

func (k *Keybase) Queries() *Queries {
	k.mu.RLock()
	defer k.mu.RUnlock()
	return k.queries
}

func (k *Keybase) DID() string {
	k.mu.RLock()
	defer k.mu.RUnlock()
	return k.did
}

func (k *Keybase) DIDID() int64 {
	k.mu.RLock()
	defer k.mu.RUnlock()
	return k.didID
}

func (k *Keybase) SetDID(did string, didID int64) {
	k.mu.Lock()
	defer k.mu.Unlock()
	k.did = did
	k.didID = didID
}

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("create DID: %w", err)
	}

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

	return did, nil
}

func (k *Keybase) Load(ctx context.Context, data []byte) (string, error) {
	if len(data) < 100 {
		return "", fmt.Errorf("invalid database format")
	}

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

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

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

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

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

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

	return k.did, nil
}

func (k *Keybase) Serialize() ([]byte, error) {
	k.mu.RLock()
	defer k.mu.RUnlock()

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

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

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("begin tx: %w", err)
	}

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

	return tx.Commit()
}