motr-enclave/main.go

package main

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"enclave/internal/keybase"
	"enclave/internal/state"
	"enclave/internal/types"

	"github.com/extism/go-pdk"
)

func main() { state.Default() }

//go:wasmexport ping
func ping() int32 {
	pdk.Log(pdk.LogInfo, "ping: received request")

	var input types.PingInput
	if err := pdk.InputJSON(&input); err != nil {
		output := types.PingOutput{
			Success: false,
			Message: fmt.Sprintf("failed to parse input: %s", err),
		}
		pdk.OutputJSON(output)
		return 0
	}

	output := types.PingOutput{
		Success: true,
		Message: "pong",
		Echo:    input.Message,
	}

	if err := pdk.OutputJSON(output); err != nil {
		pdk.Log(pdk.LogError, fmt.Sprintf("ping: failed to output: %s", err))
		return 1
	}

	pdk.Log(pdk.LogInfo, fmt.Sprintf("ping: responded with echo=%s", input.Message))
	return 0
}

//go:wasmexport generate
func generate() int32 {
	pdk.Log(pdk.LogInfo, "generate: starting database initialization")

	var input types.GenerateInput
	if err := pdk.InputJSON(&input); err != nil {
		pdk.SetError(fmt.Errorf("generate: failed to parse input: %w", err))
		return 1
	}

	if input.Credential == "" {
		pdk.SetError(errors.New("generate: credential is required"))
		return 1
	}

	credentialBytes, err := base64.StdEncoding.DecodeString(input.Credential)
	if err != nil {
		pdk.SetError(fmt.Errorf("generate: invalid base64 credential: %w", err))
		return 1
	}

	did, err := initializeDatabase(credentialBytes)
	if err != nil {
		pdk.SetError(fmt.Errorf("generate: failed to initialize database: %w", err))
		return 1
	}

	state.SetInitialized(true)
	state.SetDID(did)

	dbBytes, err := serializeDatabase()
	if err != nil {
		pdk.SetError(fmt.Errorf("generate: failed to serialize database: %w", err))
		return 1
	}

	output := types.GenerateOutput{
		DID:      did,
		Database: dbBytes,
	}

	if err := pdk.OutputJSON(output); err != nil {
		pdk.SetError(fmt.Errorf("generate: failed to output result: %w", err))
		return 1
	}

	pdk.Log(pdk.LogInfo, fmt.Sprintf("generate: created DID %s", did))
	return 0
}

//go:wasmexport load
func load() int32 {
	pdk.Log(pdk.LogInfo, "load: loading database from buffer")

	var input types.LoadInput
	if err := pdk.InputJSON(&input); err != nil {
		pdk.SetError(fmt.Errorf("load: failed to parse input: %w", err))
		return 1
	}

	if len(input.Database) == 0 {
		pdk.SetError(errors.New("load: database buffer is required"))
		return 1
	}

	did, err := loadDatabase(input.Database)
	if err != nil {
		output := types.LoadOutput{
			Success: false,
			Error:   err.Error(),
		}
		pdk.OutputJSON(output)
		return 1
	}

	state.SetInitialized(true)
	state.SetDID(did)

	output := types.LoadOutput{
		Success: true,
		DID:     did,
	}

	if err := pdk.OutputJSON(output); err != nil {
		pdk.SetError(fmt.Errorf("load: failed to output result: %w", err))
		return 1
	}

	pdk.Log(pdk.LogInfo, fmt.Sprintf("load: loaded database for DID %s", did))
	return 0
}

//go:wasmexport exec
func exec() int32 {
	pdk.Log(pdk.LogInfo, "exec: executing action")

	if !state.IsInitialized() {
		output := types.ExecOutput{Success: false, Error: "database not initialized, call generate or load first"}
		pdk.OutputJSON(output)
		return 0
	}

	var input types.ExecInput
	if err := pdk.InputJSON(&input); err != nil {
		output := types.ExecOutput{Success: false, Error: fmt.Sprintf("failed to parse input: %s", err)}
		pdk.OutputJSON(output)
		return 0
	}

	if input.Filter == "" {
		output := types.ExecOutput{Success: false, Error: "filter is required"}
		pdk.OutputJSON(output)
		return 0
	}

	params, err := parseFilter(input.Filter)
	if err != nil {
		output := types.ExecOutput{Success: false, Error: fmt.Sprintf("invalid filter: %s", err)}
		pdk.OutputJSON(output)
		return 0
	}

	if input.Token != "" {
		if err := validateUCAN(input.Token, params); err != nil {
			output := types.ExecOutput{
				Success: false,
				Error:   fmt.Sprintf("authorization failed: %s", err.Error()),
			}
			pdk.OutputJSON(output)
			return 1
		}
	}

	result, err := executeAction(params)
	if err != nil {
		output := types.ExecOutput{
			Success: false,
			Error:   err.Error(),
		}
		pdk.OutputJSON(output)
		return 1
	}

	output := types.ExecOutput{
		Success: true,
		Result:  result,
	}

	pdk.OutputJSON(output)
	pdk.Log(pdk.LogInfo, fmt.Sprintf("exec: completed %s on %s", params.Action, params.Resource))
	return 0
}

//go:wasmexport query
func query() int32 {
	pdk.Log(pdk.LogInfo, "query: resolving DID document")

	if !state.IsInitialized() {
		pdk.SetError(errors.New("database not initialized, call generate or load first"))
		return 1
	}

	var input types.QueryInput
	if err := pdk.InputJSON(&input); err != nil {
		pdk.SetError(fmt.Errorf("query: failed to parse input: %w", err))
		return 1
	}

	if input.DID == "" {
		input.DID = state.GetDID()
	}

	if !strings.HasPrefix(input.DID, "did:") {
		pdk.SetError(errors.New("query: invalid DID format"))
		return 1
	}

	output, err := resolveDID(input.DID)
	if err != nil {
		pdk.SetError(fmt.Errorf("query: failed to resolve DID: %w", err))
		return 1
	}

	if err := pdk.OutputJSON(output); err != nil {
		pdk.SetError(fmt.Errorf("query: failed to output result: %w", err))
		return 1
	}

	pdk.Log(pdk.LogInfo, fmt.Sprintf("query: resolved DID %s", input.DID))
	return 0
}

func initializeDatabase(credentialBytes []byte) (string, error) {
	kb, err := keybase.Open()
	if err != nil {
		return "", fmt.Errorf("open database: %w", err)
	}

	ctx := context.Background()
	did, err := kb.Initialize(ctx, credentialBytes)
	if err != nil {
		return "", fmt.Errorf("initialize: %w", err)
	}

	pdk.Log(pdk.LogDebug, "initializeDatabase: created schema and initial records")
	return did, nil
}

func serializeDatabase() ([]byte, error) {
	kb := keybase.Get()
	if kb == nil {
		return nil, errors.New("database not initialized")
	}
	return kb.Serialize()
}

func loadDatabase(data []byte) (string, error) {
	if len(data) < 10 {
		return "", errors.New("invalid database format")
	}

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

	ctx := context.Background()
	did, err := kb.Load(ctx, data)
	if err != nil {
		return "", fmt.Errorf("load DID: %w", err)
	}

	pdk.Log(pdk.LogDebug, "loadDatabase: database loaded successfully")
	return did, nil
}

func parseFilter(filter string) (*types.FilterParams, error) {
	params := &types.FilterParams{}
	parts := strings.FieldsSeq(filter)

	for part := range parts {
		kv := strings.SplitN(part, ":", 2)
		if len(kv) != 2 {
			continue
		}

		key, value := kv[0], kv[1]
		switch key {
		case "resource":
			params.Resource = value
		case "action":
			params.Action = value
		case "subject":
			params.Subject = value
		}
	}

	if params.Resource == "" {
		return nil, errors.New("resource is required")
	}
	if params.Action == "" {
		return nil, errors.New("action is required")
	}

	return params, nil
}

func validateUCAN(token string, params *types.FilterParams) error {
	if token == "" {
		return errors.New("token is required")
	}

	pdk.Log(pdk.LogDebug, fmt.Sprintf("validateUCAN: validated token for %s:%s", params.Resource, params.Action))
	return nil
}

func executeAction(params *types.FilterParams) (json.RawMessage, error) {
	switch params.Resource {
	case "accounts":
		return executeAccountAction(params)
	case "credentials":
		return executeCredentialAction(params)
	case "sessions":
		return executeSessionAction(params)
	case "grants":
		return executeGrantAction(params)
	default:
		return nil, fmt.Errorf("unknown resource: %s", params.Resource)
	}
}

func executeAccountAction(params *types.FilterParams) (json.RawMessage, error) {
	switch params.Action {
	case "list":
		accounts := []types.Account{
			{
				Address:      "sonr1abc123...",
				ChainID:      "sonr-mainnet-1",
				CoinType:     118,
				AccountIndex: 0,
				AddressIndex: 0,
				Label:        "Primary",
				IsDefault:    true,
			},
		}
		return json.Marshal(accounts)
	case "balances":
		return fetchAccountBalances(params.Subject)
	case "sign":
		return json.Marshal(map[string]string{"signature": "placeholder"})
	default:
		return nil, fmt.Errorf("unknown action for accounts: %s", params.Action)
	}
}

func fetchAccountBalances(address string) (json.RawMessage, error) {
	if address == "" {
		address = state.GetDID()
	}

	apiBase, ok := state.GetConfig("api_endpoint")
	if !ok {
		apiBase = "https://api.sonr.io"
	}

	url := fmt.Sprintf("%s/cosmos/bank/v1beta1/balances/%s", apiBase, address)
	pdk.Log(pdk.LogInfo, fmt.Sprintf("fetchAccountBalances: GET %s", url))

	req := pdk.NewHTTPRequest(pdk.MethodGet, url)
	req.SetHeader("Accept", "application/json")

	res := req.Send()
	status := res.Status()

	if status < 200 || status >= 300 {
		pdk.Log(pdk.LogError, fmt.Sprintf("fetchAccountBalances: HTTP %d", status))
		return json.Marshal(map[string]any{
			"error":   "failed to fetch balances",
			"status":  status,
			"address": address,
		})
	}

	body := res.Body()
	pdk.Log(pdk.LogDebug, fmt.Sprintf("fetchAccountBalances: received %d bytes", len(body)))

	return body, nil
}

func executeCredentialAction(params *types.FilterParams) (json.RawMessage, error) {
	switch params.Action {
	case "list":
		credentials := []types.Credential{
			{
				CredentialID:  "cred_abc123",
				DeviceName:    "MacBook Pro",
				DeviceType:    "platform",
				Authenticator: "Touch ID",
				Transports:    []string{"internal"},
				CreatedAt:     "2025-01-01T00:00:00Z",
				LastUsed:      "2025-01-01T00:00:00Z",
			},
		}
		return json.Marshal(credentials)
	default:
		return nil, fmt.Errorf("unknown action for credentials: %s", params.Action)
	}
}

func executeSessionAction(params *types.FilterParams) (json.RawMessage, error) {
	switch params.Action {
	case "list":
		return json.Marshal([]map[string]any{})
	case "create":
		return json.Marshal(map[string]string{"session_id": "sess_placeholder"})
	case "revoke":
		return json.Marshal(map[string]bool{"revoked": true})
	default:
		return nil, fmt.Errorf("unknown action for sessions: %s", params.Action)
	}
}

func executeGrantAction(params *types.FilterParams) (json.RawMessage, error) {
	switch params.Action {
	case "list":
		return json.Marshal([]map[string]any{})
	case "create":
		return json.Marshal(map[string]string{"grant_id": "grant_placeholder"})
	case "revoke":
		return json.Marshal(map[string]bool{"revoked": true})
	default:
		return nil, fmt.Errorf("unknown action for grants: %s", params.Action)
	}
}

func resolveDID(did string) (*types.QueryOutput, error) {
	output := &types.QueryOutput{
		DID:        did,
		Controller: did,
		VerificationMethods: []types.VerificationMethod{
			{
				ID:         did + "#key-1",
				Type:       "Ed25519VerificationKey2020",
				Controller: did,
				PublicKey:  "placeholder_public_key",
				Purpose:    "authentication",
			},
		},
		Accounts: []types.Account{
			{
				Address:      "sonr1abc123...",
				ChainID:      "sonr-mainnet-1",
				CoinType:     118,
				AccountIndex: 0,
				AddressIndex: 0,
				Label:        "Primary",
				IsDefault:    true,
			},
		},
		Credentials: []types.Credential{
			{
				CredentialID:  "cred_abc123",
				DeviceName:    "MacBook Pro",
				DeviceType:    "platform",
				Authenticator: "Touch ID",
				Transports:    []string{"internal"},
				CreatedAt:     "2025-01-01T00:00:00Z",
				LastUsed:      "2025-01-01T00:00:00Z",
			},
		},
	}

	return output, nil
}