diff --git a/cmd/fees/worker/main.go b/cmd/fees/worker/main.go index 7da9b84..3ac5ed1 100644 --- a/cmd/fees/worker/main.go +++ b/cmd/fees/worker/main.go @@ -2,11 +2,13 @@ package main import ( "context" + "encoding/json" "fmt" "time" "github.com/DataDog/datadog-go/statsd" "github.com/hibiken/asynq" + "github.com/robfig/cron/v3" "github.com/sirupsen/logrus" "github.com/vultisig/plugin/internal/keysign" "github.com/vultisig/verifier/tx_indexer" @@ -20,6 +22,45 @@ import ( "github.com/vultisig/plugin/storage/postgres" ) +func startLoadingFees(asynqClient *asynq.Client, logger *logrus.Logger) { + logger.Info("Loading fees") + payload, err := json.Marshal(fees.FeeCollectionFormat{ + FeeCollectionType: fees.FeeCollectionTypeAll, + }) + if err != nil { + logger.WithError(err).Error("Failed to marshal fee loading config in demo run") + return + } + asynqClient.Enqueue( + asynq.NewTask(fees.TypeFeeLoad, payload), + asynq.MaxRetry(0), + asynq.Timeout(2*time.Minute), + asynq.Retention(5*time.Minute), + asynq.Queue(tasks.QUEUE_NAME)) +} + +func startTransactingFees(asynqClient *asynq.Client, logger *logrus.Logger) { + logger.Info("Transacting fees") + payload := make([]byte, 0) + asynqClient.Enqueue( + asynq.NewTask(fees.TypeFeeTransact, payload), + asynq.MaxRetry(0), + asynq.Timeout(2*time.Minute), + asynq.Retention(5*time.Minute), + asynq.Queue(tasks.QUEUE_NAME)) +} + +func startPostTx(asynqClient *asynq.Client, logger *logrus.Logger) { + logger.Info("Checking status") + payload := make([]byte, 0) + asynqClient.Enqueue( + asynq.NewTask(fees.TypeFeePostTx, payload), + asynq.MaxRetry(0), + asynq.Timeout(2*time.Minute), + asynq.Retention(5*time.Minute), + asynq.Queue(tasks.QUEUE_NAME)) +} + func main() { ctx := context.Background() @@ -132,9 +173,33 @@ func main() { mux.HandleFunc(tasks.TypeReshareDKLS, vaultService.HandleReshareDKLS) //Plugin specific functions. - mux.HandleFunc(fees.TypeFeeCollection, feePlugin.HandleCollections) + mux.HandleFunc(fees.TypeFeeLoad, feePlugin.LoadFees) + mux.HandleFunc(fees.TypeFeeTransact, feePlugin.HandleTransactions) + mux.HandleFunc(fees.TypeFeePostTx, feePlugin.HandlePostTx) + + // Load fees every 10 minutes + loadFees := cron.New() + loadFees.AddFunc(feePluginConfig.Jobs.Load.Cronexpr, func() { + startLoadingFees(asynqClient, logger) + }) + loadFees.Start() + + // Transact fees every Friday at 12:00 PM + transactFees := cron.New() + transactFees.AddFunc(feePluginConfig.Jobs.Transact.Cronexpr, func() { + startTransactingFees(asynqClient, logger) + }) + transactFees.Start() + + // // Update verifier every 10 minutes + updateVerifier := cron.New() + updateVerifier.AddFunc(feePluginConfig.Jobs.Post.Cronexpr, func() { + startPostTx(asynqClient, logger) + }) + updateVerifier.Start() logger.Info("Starting asynq listener") + if err := srv.Run(mux); err != nil { panic(fmt.Errorf("could not run server: %w", err)) } diff --git a/etc/vultisig/fee.yml b/etc/vultisig/fee.yml index 48d88d6..1486f67 100644 --- a/etc/vultisig/fee.yml +++ b/etc/vultisig/fee.yml @@ -1,9 +1,16 @@ type: fee version: 1.0.0 -rpc_url: https://ethereum.publicnode.com/ -collector_whitelist_addresses: - - 0x7d760c17d798a7A9a4c4AcAf311A02dC95972503 - - 0x5BB06B9C5e4f7624Df7100Badb2F5AA1C86d8498 -collector_address: 0x7d760c17d798a7A9a4c4AcAf311A02dC95972503 usdc_address: 0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48 verifier_token: localhost-fee-apikey +eth_provider: https://eth-mainnet.g.alchemy.com/v2/HAtIwB5y82TNVwaHkwcMj +chain_id: 1 +jobs: + load: + cronexpr: "@every 2m" + max_concurrent_jobs: 10 + transact: + cronexpr: "0 12 * * 5" + max_concurrent_jobs: 10 + post: + cronexpr: "@every 5m" + max_concurrent_jobs: 10 \ No newline at end of file diff --git a/internal/types/fees.go b/internal/types/fees.go index 6dea895..78c1310 100644 --- a/internal/types/fees.go +++ b/internal/types/fees.go @@ -11,8 +11,10 @@ import ( type FeeRunState string const ( - FeeRunStateDraft FeeRunState = "draft" - FeeRunStateSent FeeRunState = "sent" + FeeRunStateDraft FeeRunState = "draft" + FeeRunStateSent FeeRunState = "sent" + FeeRunStateSuccess FeeRunState = "completed" + FeeRunStateFailed FeeRunState = "failed" ) // individual fee record in the db @@ -29,8 +31,9 @@ type FeeRun struct { Status FeeRunState `db:"status"` CreatedAt time.Time `db:"created_at"` UpdatedAt time.Time `db:"updated_at"` - TxID *uuid.UUID `db:"tx_id"` + TxHash *string `db:"tx_hash"` PolicyID uuid.UUID `db:"policy_id"` TotalAmount int `db:"total_amount"` FeeCount int `db:"fee_count"` + Fees []Fee `db:"fees"` } diff --git a/internal/verifierapi/fees.go b/internal/verifierapi/fees.go index b864f15..e1dbfdc 100644 --- a/internal/verifierapi/fees.go +++ b/internal/verifierapi/fees.go @@ -4,6 +4,7 @@ import ( "encoding/json" "fmt" "net/http" + "time" "github.com/google/uuid" ) @@ -55,3 +56,29 @@ func (v *VerifierApi) GetPublicKeysFees(ecdsaPublicKey string) (*FeeHistoryDto, return &feeHistory.Data, nil } + +func (v *VerifierApi) MarkFeeAsCollected(txHash string, collectedAt time.Time, feeIds ...uuid.UUID) error { + + var body = struct { + IDs []uuid.UUID `json:"ids"` + TxHash string `json:"tx_hash"` + CollectedAt time.Time `json:"collected_at"` + }{ + IDs: feeIds, + TxHash: txHash, + CollectedAt: collectedAt, + } + + url := "/fees/collected" + response, err := v.postAuth(url, body) + if err != nil { + return fmt.Errorf("failed to mark fee as collected: %w", err) + } + defer response.Body.Close() + + if response.StatusCode != http.StatusOK { + return fmt.Errorf("failed to mark fee as collected, status code: %d", response.StatusCode) + } + + return nil +} diff --git a/internal/verifierapi/verifierapi.go b/internal/verifierapi/verifierapi.go index ddb3902..c0d540d 100644 --- a/internal/verifierapi/verifierapi.go +++ b/internal/verifierapi/verifierapi.go @@ -1,6 +1,8 @@ package verifierapi import ( + "bytes" + "encoding/json" "net/http" "github.com/sirupsen/logrus" @@ -49,3 +51,20 @@ func (v *VerifierApi) getAuth(endpoint string) (*http.Response, error) { r.Header.Set("Authorization", "Bearer "+v.token) return v.client.Do(r) } + +func (v *VerifierApi) postAuth(endpoint string, body any) (*http.Response, error) { + jsonBody, err := json.Marshal(body) + if err != nil { + return nil, err + } + + request, err := http.NewRequest(http.MethodPost, v.url+endpoint, bytes.NewBuffer(jsonBody)) + if err != nil { + return nil, err + } + + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Authorization", "Bearer "+v.token) + + return v.client.Do(request) +} diff --git a/plugin/fees/config.go b/plugin/fees/config.go index 36fdd74..d4a6218 100644 --- a/plugin/fees/config.go +++ b/plugin/fees/config.go @@ -4,40 +4,36 @@ import ( "errors" "fmt" "math/big" - "slices" "strings" "github.com/spf13/viper" ) -/* - { - "type": "fees", - "version": "1.0.0", - "rpc_url": "https://eth-mainnet.g.alchemy.com/v2/YOUR_API_KEY", - "usdc_address": "0xA0b86a33E6441b8C4C8C8C8C8C8C8C8C8C8C8C8C", - "gas": { - "limit_multiplier": 1, - "price_multiplier": 1 - }, - "monitoring": { - "timeout_minutes": 30, - "check_interval_seconds": 60 - } - } -*/ - // These are properties and parameters specific to the fee plugin config. They should be distinct from system/core config type FeeConfig struct { - Type string `mapstructure:"type"` - Version string `mapstructure:"version"` - RpcURL string `mapstructure:"rpc_url"` // URL for the RPC endpoint to interact with the ethereum/evm blockchain - CollectorWhitelistAddresses []string `mapstructure:"collector_whitelist_addresses"` // A list of whitelisted addresses for which fee transactions are collected against. These can include previous addresses. Fee plugins with a recipient address that is not in this list will not be processed. - CollectorAddress string `mapstructure:"collector_address"` // This address is what is used for new policies. Fee policies created with a different address will be rejected. - MaxFeeAmount uint64 `mapstructure:"max_fee_amount"` // Policies that are created/submitted which do not have this amount will be rejected. - UsdcAddress string `mapstructure:"usdc_address"` // The address of the USDC token on the Ethereum blockchain. - VerifierToken string `mapstructure:"verifier_token"` // The token to use for the verifier API. - ChainId *big.Int `mapstructure:"chain_id"` // The chain ID of the Ethereum blockchain. + Type string `mapstructure:"type"` + Version string `mapstructure:"version"` + MaxFeeAmount uint64 `mapstructure:"max_fee_amount"` // Policies that are created/submitted which do not have this amount will be rejected. + UsdcAddress string `mapstructure:"usdc_address"` // The address of the USDC token on the Ethereum blockchain. + VerifierToken string `mapstructure:"verifier_token"` // The token to use for the verifier API. + chainId uint64 `mapstructure:"chain_id"` // The chain ID of the Ethereum blockchain. + ChainId *big.Int + EthProvider string `mapstructure:"eth_provider"` // The Ethereum provider to use for the fee plugin. + Jobs struct { + Load struct { + MaxConcurrentJobs uint64 `mapstructure:"max_concurrent_jobs"` //How many consecutive tasks can take place + Cronexpr string `mapstructure:"cronexpr"` // Cron link expression on how often these tasks should run + } `mapstructure:"load"` + Transact struct { + MaxConcurrentJobs uint64 `mapstructure:"max_concurrent_jobs"` //How many consecutive tasks can take place + Cronexpr string `mapstructure:"cronexpr"` // Cron link expression on how often these tasks should run + } `mapstructure:"transact"` + Post struct { + SuccessConfirmations uint64 `mapstructure:"success_confirmations"` //How many consecutive tasks can take place + Cronexpr string `mapstructure:"cronexpr"` // Cron link expression on how often these tasks should run + MaxConcurrentJobs uint64 `mapstructure:"max_concurrent_jobs"` + } `mapstructure:"post"` + } } type ConfigOption func(*FeeConfig) error @@ -46,43 +42,60 @@ func withDefaults(c *FeeConfig) { c.ChainId = big.NewInt(1) c.Type = PLUGIN_TYPE c.Version = "1.0.0" - c.RpcURL = "https://ethereum.publicnode.com/" - c.CollectorWhitelistAddresses = []string{} - c.CollectorAddress = "" c.MaxFeeAmount = 500e6 // 500 USDC + + c.Jobs.Load.MaxConcurrentJobs = 10 + c.Jobs.Transact.MaxConcurrentJobs = 10 + c.Jobs.Post.MaxConcurrentJobs = 10 + c.Jobs.Post.SuccessConfirmations = 20 + + c.Jobs.Load.Cronexpr = "@every 2m" + c.Jobs.Transact.Cronexpr = "0 12 * * 5" + c.Jobs.Post.Cronexpr = "@every 5m" } -func WithEthConfig(rpcUrl string) ConfigOption { +func WithMaxFeeAmount(maxFeeAmount uint64) ConfigOption { return func(c *FeeConfig) error { - c.RpcURL = rpcUrl + c.MaxFeeAmount = maxFeeAmount return nil } } -func WithCollectorAddress(collectorAddress string) ConfigOption { +func WithChainId(chainId *big.Int) ConfigOption { return func(c *FeeConfig) error { - c.CollectorAddress = collectorAddress + c.ChainId = chainId return nil } } -func WithCollectorWhitelistAddresses(collectorWhitelistAddresses []string) ConfigOption { +func WithEthClient(url string) ConfigOption { return func(c *FeeConfig) error { - c.CollectorWhitelistAddresses = collectorWhitelistAddresses + c.EthProvider = url return nil } } -func WithMaxFeeAmount(maxFeeAmount uint64) ConfigOption { +func WithSuccessConfirmations(successConfirmations uint64) ConfigOption { return func(c *FeeConfig) error { - c.MaxFeeAmount = maxFeeAmount + c.Jobs.Post.SuccessConfirmations = successConfirmations return nil } } -func WithChainId(chainId *big.Int) ConfigOption { +func WithJobConcurrency(load, transact, post uint64) ConfigOption { return func(c *FeeConfig) error { - c.ChainId = chainId + c.Jobs.Load.MaxConcurrentJobs = load + c.Jobs.Transact.MaxConcurrentJobs = transact + c.Jobs.Post.MaxConcurrentJobs = post + return nil + } +} + +func WithCronexpr(load, transact, post string) ConfigOption { + return func(c *FeeConfig) error { + c.Jobs.Load.Cronexpr = load + c.Jobs.Transact.Cronexpr = transact + c.Jobs.Post.Cronexpr = post return nil } } @@ -112,6 +125,8 @@ func WithFileConfig(basePath string) ConfigOption { if err := v.Unmarshal(c); err != nil { return fmt.Errorf("failed to unmarshal config: %w", err) } + + c.ChainId = big.NewInt(0).SetUint64(c.chainId) return nil } } @@ -129,23 +144,6 @@ func NewFeeConfig(fns ...ConfigOption) (*FeeConfig, error) { if c.Type != PLUGIN_TYPE { return c, fmt.Errorf("invalid plugin type: %s", c.Type) } - if c.RpcURL == "" { - return c, errors.New("rpc_url is required") - } - // Collector address cannot be empty - if c.CollectorAddress == "" { - return c, errors.New("collector_address is required") - } - - // There must be at least one collector whitelist address - if len(c.CollectorWhitelistAddresses) == 0 { - return c, errors.New("collector_whitelist_addresses is required") - } - - // Collector address must be in the whitelist - if !slices.Contains(c.CollectorWhitelistAddresses, c.CollectorAddress) { - return c, fmt.Errorf("collector_address must be in the whitelist: %s, whitelist: %v", c.CollectorAddress, c.CollectorWhitelistAddresses) - } if c.VerifierToken == "" { return c, errors.New("verifier_token is required") @@ -155,6 +153,19 @@ func NewFeeConfig(fns ...ConfigOption) (*FeeConfig, error) { return c, errors.New("chain_id is required") } + if c.EthProvider == "" { + return c, errors.New("eth_provider is required") + } + + if c.Jobs.Load.MaxConcurrentJobs < 1 || + c.Jobs.Load.MaxConcurrentJobs > 100 || + c.Jobs.Transact.MaxConcurrentJobs < 1 || + c.Jobs.Transact.MaxConcurrentJobs > 100 || + c.Jobs.Post.MaxConcurrentJobs < 1 || + c.Jobs.Post.MaxConcurrentJobs > 100 { + return c, errors.New("max_concurrent_jobs must be greater than 0 and less than 100") + } + return c, nil } diff --git a/plugin/fees/constraints.go b/plugin/fees/constraints.go index 388e269..7748fbe 100644 --- a/plugin/fees/constraints.go +++ b/plugin/fees/constraints.go @@ -3,7 +3,9 @@ package fees const PLUGIN_TYPE = "fee" // Task Definitions -const TypeFeeCollection = "fees:collection" +const TypeFeeLoad = "fees:load" // Load list of pending fees into the db from the verifier +const TypeFeeTransact = "fees:transaction" // Collect a list of loaded fees from the users wallet +const TypeFeePostTx = "fees:post_tx" // Check the status of the fee runs var ERC20_TRANSFER_GAS int = 65000 //typically the upper bound from an ERC20 transfer diff --git a/plugin/fees/fees.go b/plugin/fees/fees.go index 7e90bb9..abd0c24 100644 --- a/plugin/fees/fees.go +++ b/plugin/fees/fees.go @@ -2,8 +2,8 @@ package fees import ( "context" - "encoding/json" "fmt" + "sync" "github.com/ethereum/go-ethereum/ethclient" "github.com/google/uuid" @@ -15,6 +15,7 @@ import ( "github.com/vultisig/verifier/tx_indexer" vtypes "github.com/vultisig/verifier/types" "golang.org/x/sync/errgroup" + "golang.org/x/sync/semaphore" "github.com/vultisig/verifier/vault" @@ -45,6 +46,8 @@ type FeePlugin struct { asynqInspector *asynq.Inspector asynqClient *asynq.Client encryptionSecret string + transactingMutex sync.Mutex // when actual transactions are happening we cannot load fees + ethClient *ethclient.Client } func NewFeePlugin(db storage.DatabaseStorage, @@ -63,7 +66,7 @@ func NewFeePlugin(db storage.DatabaseStorage, return nil, fmt.Errorf("database storage cannot be nil") } - rpcClient, err := ethclient.Dial(feeConfig.RpcURL) + rpcClient, err := ethclient.Dial(feeConfig.EthProvider) if err != nil { return nil, err } @@ -90,6 +93,11 @@ func NewFeePlugin(db storage.DatabaseStorage, return nil, fmt.Errorf("failed to create verifier api") } + ethClient, err := ethclient.Dial(feeConfig.EthProvider) + if err != nil { + return nil, fmt.Errorf("failed to dial eth client: %w", err) + } + return &FeePlugin{ db: db, eth: evm.NewSDK(feeConfig.ChainId, rpcClient, rpcClient.Client()), @@ -102,98 +110,52 @@ func NewFeePlugin(db storage.DatabaseStorage, asynqInspector: inspector, asynqClient: client, encryptionSecret: encryptionSecret, + ethClient: ethClient, }, nil } -/* -The handler of the asynq job. Fees can initialized and collected in 3 ways: -- By public key (queries the db for a single fee_policy and then kicks off the fee collection) -- By policy id (queries the db for a fee_policy matching that id and then kicks off the fee collection) -- All (queries the db for all fee_policies and then kicks off the fee collection for each policy) -*/ -func (fp *FeePlugin) HandleCollections(ctx context.Context, task *asynq.Task) error { - fp.logger.Info("Starting Fee Collection Job") +/* ------------------------------------------------------------------------------------------------ +LOADING FEES +here we pull in a list of fees (amounts and ids) that are pending collection and add them to a fee run +------------------------------------------------------------------------------------------------ */ - // Figure out if we're collecting fees by public key, policy, or plugin id - feeCollectionFormat := FeeCollectionFormat{ - FeeCollectionType: FeeCollectionTypeAll, - } - if len(task.Payload()) != 0 { - if err := json.Unmarshal(task.Payload(), &feeCollectionFormat); err != nil { - return fmt.Errorf("fp.HandleCollections, failed to unmarshall asynq task payload, %w", err) - } - } - switch feeCollectionFormat.FeeCollectionType { - case FeeCollectionTypeByPublicKey: - fp.logger.Info("Collecting fees by public key") - return fp.collectFeesByPublicKey(ctx, feeCollectionFormat.Value) - case FeeCollectionTypeByPolicy: - fp.logger.Info("Collecting fees by policy") - return fp.collectFeesByPolicyId(ctx, feeCollectionFormat.Value) - case FeeCollectionTypeAll: - fp.logger.Info("Collecting all fees") - return fp.collectAllFees(ctx) - default: - return fmt.Errorf("invalid fee collection type") - } -} +func (fp *FeePlugin) LoadFees(ctx context.Context, task *asynq.Task) error { + fp.transactingMutex.Lock() + defer fp.transactingMutex.Unlock() -func (fp *FeePlugin) collectFeesByPublicKey(ctx context.Context, publicKey string) error { - // Get the fee policy from the database - feePolicies, err := fp.db.GetAllPluginPolicies(ctx, publicKey, vtypes.PluginVultisigFees_feee, true) - if err != nil { - return fmt.Errorf("failed to get plugin policy: %w", err) - } - if len(feePolicies) == 0 { - return fmt.Errorf("no fee policy found for public key: %s", publicKey) - } - if len(feePolicies) > 1 { - return fmt.Errorf("multiple fee policies found for public key: %s", publicKey) - } - return fp.executeFeeCollection(ctx, feePolicies[0]) -} + fp.logger.Info("Starting Fee Loading Job") -func (fp *FeePlugin) collectFeesByPolicyId(ctx context.Context, policyId string) error { - policyIdUuid, err := uuid.Parse(policyId) - if err != nil { - return fmt.Errorf("failed to parse policy id: %w", err) - } - policy, err := fp.db.GetPluginPolicy(ctx, policyIdUuid) + feePolicies, err := fp.db.GetAllFeePolicies(ctx) if err != nil { return fmt.Errorf("failed to get plugin policy: %w", err) } - return fp.executeFeeCollection(ctx, *policy) -} - -func (fp *FeePlugin) collectAllFees(ctx context.Context) error { - fp.logger.Info("Collecting all fees") - feePolicies, err := fp.db.GetAllPluginPolicies(ctx, "", vtypes.PluginVultisigFees_feee, true) - if err != nil { - return fmt.Errorf("failed to get fee policies: %w", err) - } + // We limit the number of concurrent fee loading operations to 10 + sem := semaphore.NewWeighted(int64(fp.config.Jobs.Load.MaxConcurrentJobs)) + var wg sync.WaitGroup var eg errgroup.Group + for _, feePolicy := range feePolicies { - feePolicy := feePolicy // Capture by value + wg.Add(1) + feePolicy = feePolicy eg.Go(func() error { - return fp.executeFeeCollection(ctx, feePolicy) + defer wg.Done() + if err := sem.Acquire(context.Background(), 1); err != nil { + return fmt.Errorf("failed to acquire semaphore: %w", err) + } + defer sem.Release(1) + return fp.executeFeeLoading(ctx, feePolicy) }) } - return eg.Wait() + + wg.Wait() + if err := eg.Wait(); err != nil { + return fmt.Errorf("failed to execute fee loading: %w", err) + } + return nil } -/* -This function is the main function that collects fees. It is called by -- collectFeesByPublicKey, -- collectFeesByPolicyId -- collectAllFees -It does the following: -- Gets the list of fees from the verifier -- If there are fees to collect, it creates a fee run, errors if already being collected. -- It gets a vault and proposes the transactions -- It initializes the signing -*/ -func (fp *FeePlugin) executeFeeCollection(ctx context.Context, feePolicy vtypes.PluginPolicy) error { +func (fp *FeePlugin) executeFeeLoading(ctx context.Context, feePolicy vtypes.PluginPolicy) error { // Get list of fees from the verifier connected to the fee policy feesResponse, err := fp.verifierApi.GetPublicKeysFees(feePolicy.PublicKey) @@ -212,34 +174,151 @@ func (fp *FeePlugin) executeFeeCollection(ctx context.Context, feePolicy vtypes. "publicKey": feePolicy.PublicKey, }).Info("Fees pending collection: ", feesResponse.FeesPendingCollection) - // Get list of fee ids to be collected in this batch - // Verify that the sum of the fees is equal to the fees pending collection - feesToCollect := make([]uuid.UUID, 0, len(feesResponse.Fees)) checkAmount := 0 for _, fee := range feesResponse.Fees { if !fee.Collected { - feesToCollect = append(feesToCollect, fee.ID) checkAmount += fee.Amount } } if checkAmount != feesResponse.FeesPendingCollection { return fmt.Errorf("fees pending collection amount does not match the sum of the fees") } - fp.logger.WithFields(logrus.Fields{ - "publicKey": feePolicy.PublicKey, - "amount": checkAmount, - }).Info("Collecting fee ids: ", feesToCollect) - // Here we check if the fee collection is already in progress for any of the specific fee ids - feeRun, err := fp.db.CreateFeeRun(ctx, feePolicy.ID, types.FeeRunStateDraft, feesResponse.Fees) + for _, fee := range feesResponse.Fees { + if !fee.Collected { + + // Check if the fee has already been loaded and added to a fee run, if so, skip it + existingFee, err := fp.db.GetFees(ctx, fee.ID) + if err != nil { + return fmt.Errorf("failed to get fee: %w", err) + } + if len(existingFee) > 0 { + fp.logger.WithFields(logrus.Fields{ + "publicKey": feePolicy.PublicKey, + "feeId": fee.ID, + "runId": existingFee[0].FeeRunID, + }).Info("Fee already added to a fee run") + continue + } + + // If the fee hasn't been loaded, look for a draft run and add it to it + run, err := fp.db.GetPendingFeeRun(ctx, feePolicy.ID) + if err != nil { + return fmt.Errorf("failed to get pending fee run: %w", err) + } + + // If no draft run is found, create a new one and add the fee to it + if run == nil { + run, err = fp.db.CreateFeeRun(ctx, feePolicy.ID, types.FeeRunStateDraft, fee) + if err != nil { + return fmt.Errorf("failed to create fee run: %w", err) + } + fp.logger.WithFields(logrus.Fields{ + "publicKey": feePolicy.PublicKey, + "feeIds": []uuid.UUID{fee.ID}, + "runId": run.ID, + }).Info("Fee run created") + + // If a draft run is found, add the fee to it + } else { + fp.db.CreateFee(ctx, run.ID, fee) + fp.logger.WithFields(logrus.Fields{ + "publicKey": feePolicy.PublicKey, + "feeIds": []uuid.UUID{fee.ID}, + "runId": run.ID, + }).Info("Fee added to fee run") + } + } + } + + return nil +} + +/* ------------------------------------------------------------------------------------------------ +HANDLING TRANSACTIONS +here we handle the transactions for a fee run +------------------------------------------------------------------------------------------------ */ + +func (fp *FeePlugin) HandleTransactions(ctx context.Context, task *asynq.Task) error { + fp.logger.Info("Starting Fee Transaction Job. Acquiring mutex") + fp.transactingMutex.Lock() + fp.logger.Info("Mutex acquired") + + defer func() { + fp.transactingMutex.Unlock() + fp.logger.Info("Mutex released") + }() + + fp.logger.Info("Getting all fee runs") + runs, err := fp.db.GetAllFeeRuns(ctx) if err != nil { - return fmt.Errorf("failed to create fee run: %w", err) + fp.logger.WithError(err).Error("Failed to get fee runs") + return fmt.Errorf("failed to get fee runs: %w", err) + } + + sem := semaphore.NewWeighted(int64(fp.config.Jobs.Transact.MaxConcurrentJobs)) + var wg sync.WaitGroup + var eg errgroup.Group + for _, run := range runs { + run = run + eg.Go(func() error { + //TODO also check failed runs + if run.Status != types.FeeRunStateDraft { + return nil + } + + if run.TxHash != nil { + return nil + } + + if run.FeeCount == 0 || run.TotalAmount == 0 { + return nil + } + + fp.logger.WithFields(logrus.Fields{"runId": run.ID}).Info("Processing fee run") + feePolicy, err := fp.db.GetPluginPolicy(ctx, run.PolicyID) + if err != nil { + return fmt.Errorf("failed to get fee policy: %w", err) + } + wg.Add(1) + fp.logger.WithFields(logrus.Fields{"runId": run.ID, "policyId": run.PolicyID}).Info("Retrieved fee policy") + + defer wg.Done() + if err := sem.Acquire(context.Background(), 1); err != nil { + return fmt.Errorf("failed to acquire semaphore: %w", err) + } + defer sem.Release(1) + if err := fp.executeFeeTransaction(ctx, run, *feePolicy); err != nil { + fp.logger.WithFields(logrus.Fields{ + "runId": run.ID, + }).Error("Failed to execute fee transaction") + } + return nil + }) } + + wg.Wait() + if err := eg.Wait(); err != nil { + return fmt.Errorf("failed to execute fee transaction: %w", err) + } + + return nil +} + +func (fp *FeePlugin) executeFeeTransaction(ctx context.Context, run types.FeeRun, feePolicy vtypes.PluginPolicy) error { + fp.logger.WithFields(logrus.Fields{ - "publicKey": feePolicy.PublicKey, - }).Info("Fee run created with id: ", feeRun.ID) + "runId": run.ID, + "policyId": feePolicy.ID, + }).Info("Checking if fee run policy id matches fee policy id") + if run.PolicyID != feePolicy.ID { + return fmt.Errorf("fee run policy id does not match fee policy id") + } // Get a vault and sign the transactions + fp.logger.WithFields(logrus.Fields{ + "publicKey": feePolicy.PublicKey, + }).Info("Getting vault") vaultFileName := vcommon.GetVaultBackupFilename(feePolicy.PublicKey, vtypes.PluginVultisigFees_feee.String()) vaultContent, err := fp.vaultStorage.GetVault(vaultFileName) if err != nil { @@ -250,14 +329,19 @@ func (fp *FeePlugin) executeFeeCollection(ctx context.Context, feePolicy vtypes. } // Propose the transactions - keySignRequests, err := fp.ProposeTransactions(feePolicy) + fp.logger.WithFields(logrus.Fields{ + "publicKey": feePolicy.PublicKey, + }).Info("Proposing transactions") + keySignRequests, err := fp.proposeTransactions(ctx, feePolicy, run) if err != nil { return fmt.Errorf("failed to propose transactions: %w", err) } - + fp.logger.WithFields(logrus.Fields{ + "publicKey": feePolicy.PublicKey, + }).Info("Key sign requests proposed") for _, keySignRequest := range keySignRequests { req := keySignRequest - if err := fp.initSign(ctx, req, feePolicy, feeRun.ID); err != nil { + if err := fp.initSign(ctx, req, feePolicy, run.ID); err != nil { return fmt.Errorf("failed to init sign: %w", err) } } diff --git a/plugin/fees/helper.go b/plugin/fees/helper.go new file mode 100644 index 0000000..6587379 --- /dev/null +++ b/plugin/fees/helper.go @@ -0,0 +1,115 @@ +package fees + +import ( + "encoding/hex" + "fmt" + "math/big" + "strings" + + "github.com/ethereum/go-ethereum/accounts/abi" + ecommon "github.com/ethereum/go-ethereum/common" + "github.com/ethereum/go-ethereum/common/hexutil" + etypes "github.com/ethereum/go-ethereum/core/types" + "github.com/ethereum/go-ethereum/rlp" + reth "github.com/vultisig/recipes/ethereum" + "github.com/vultisig/recipes/sdk/evm" +) + +func getHash(inTx evm.UnsignedTx, r, s, v []byte, chainID *big.Int) (*etypes.Transaction, error) { + var sig []byte + sig = append(sig, r...) + sig = append(sig, s...) + sig = append(sig, v...) + + inTxDecoded, err := reth.DecodeUnsignedPayload(inTx) + if err != nil { + return nil, fmt.Errorf("reth.DecodeUnsignedPayload: %w", err) + } + + outTx, err := etypes.NewTx(inTxDecoded).WithSignature(etypes.LatestSignerForChainID(chainID), sig) + if err != nil { + return nil, fmt.Errorf("types.NewTx.WithSignature: %w", err) + } + + return outTx, nil +} + +type erc20tx struct { + to ecommon.Address `json:"to"` + amount *big.Int `json:"amount"` + token ecommon.Address `json:"token"` +} + +func decodeTx(rawHex string) (*erc20tx, error) { + type unsignedDynamicFeeTx struct { + ChainID *big.Int + Nonce uint64 + GasTipCap *big.Int + GasFeeCap *big.Int + Gas uint64 + To *ecommon.Address + Value *big.Int + Data []byte + AccessList etypes.AccessList + } + + rawHex = strings.TrimPrefix(rawHex, "0x") + + rawBytes, err := hex.DecodeString(rawHex) + if err != nil { + return nil, fmt.Errorf("hex decode failed: %w", err) + } + + // Check transaction type (EIP-1559 is 0x02) + if len(rawBytes) == 0 || rawBytes[0] != 0x02 { + return nil, fmt.Errorf("unsupported transaction type: 0x%02x", rawBytes[0]) + } + + tx := new(unsignedDynamicFeeTx) + err = rlp.DecodeBytes(rawBytes[1:], tx) + if err != nil { + return nil, fmt.Errorf("rlp decode failed: %w", err) + } + + // Parse the ERC20 transfer ABI + const transferABI = `[{"name":"transfer","type":"function","inputs":[{"name":"to","type":"address"},{"name":"value","type":"uint256"}],"outputs":[{"name":"","type":"bool"}]}]` + parsedABI, err := abi.JSON(strings.NewReader(transferABI)) + if err != nil { + return nil, fmt.Errorf("failed to parse ABI") + } + + // Get the method by selector + method, err := parsedABI.MethodById(tx.Data[:4]) + if err != nil { + return nil, fmt.Errorf("unknown method ID") + } + + // Decode the arguments + args := make(map[string]interface{}) + if err := method.Inputs.UnpackIntoMap(args, tx.Data[4:]); err != nil { + return nil, fmt.Errorf("failed get recipient and amount from tx") + } + + recipient, ok := args["to"].(ecommon.Address) + if !ok { + return nil, fmt.Errorf("invalid usdc address") + } + + amount, ok := args["value"].(*big.Int) + if !ok { + return nil, fmt.Errorf("invalid amount") + } + + return &erc20tx{ + to: recipient, + amount: amount, + token: *tx.To, + }, nil +} + +func hexutilDecode(hexStr string) ([]byte, error) { + if !strings.HasPrefix(hexStr, "0x") { + hexStr = "0x" + hexStr + } + return hexutil.Decode(hexStr) +} diff --git a/plugin/fees/post_tx.go b/plugin/fees/post_tx.go new file mode 100644 index 0000000..6bd9bd6 --- /dev/null +++ b/plugin/fees/post_tx.go @@ -0,0 +1,100 @@ +package fees + +import ( + "context" + "fmt" + "sync" + + "github.com/ethereum/go-ethereum" + ecommon "github.com/ethereum/go-ethereum/common" + "github.com/google/uuid" + "github.com/hibiken/asynq" + "github.com/sirupsen/logrus" + "github.com/vultisig/plugin/internal/types" + "golang.org/x/sync/errgroup" + "golang.org/x/sync/semaphore" +) + +// Functions here handle the post-transaction logic. Once a transaction is broadcasted, we need to update the fee run and the fee + +func (fp *FeePlugin) HandlePostTx(ctx context.Context, task *asynq.Task) error { + // Get a list of all fee runs that are in a sent state + runs, err := fp.db.GetAllFeeRuns(ctx, types.FeeRunStateSent) + if err != nil { + fp.logger.WithError(err).Error("failed to get fee runs") + return fmt.Errorf("failed to get fee runs: %w", err) + } + + currentBlock, err := fp.ethClient.BlockNumber(ctx) + if err != nil { + fp.logger.WithError(err).Error("failed to get current block") + return fmt.Errorf("failed to get current block: %w", err) + } + + sem := semaphore.NewWeighted(int64(fp.config.Jobs.Post.MaxConcurrentJobs)) + var wg sync.WaitGroup + var eg errgroup.Group + for _, run := range runs { + wg.Add(1) + run = run + eg.Go(func() error { + defer wg.Done() + if err := sem.Acquire(context.Background(), 1); err != nil { + return fmt.Errorf("failed to acquire semaphore: %w", err) + } + defer sem.Release(1) + if run.TxHash == nil || run.Status == types.FeeRunStateDraft { + return nil + } + return fp.updateStatus(ctx, run, currentBlock) + }) + } + wg.Wait() + if err := eg.Wait(); err != nil { + return fmt.Errorf("failed to execute fee run status check: %w", err) + } + fp.logger.Info("Fee run status check completed") + return nil +} + +func (fp *FeePlugin) updateStatus(ctx context.Context, run types.FeeRun, currentBlock uint64) error { + if run.TxHash == nil || run.Status == types.FeeRunStateDraft { + return nil + } + fp.logger.WithFields(logrus.Fields{"run_id": run.ID}).Info("Beginning status check/update") + hash := ecommon.HexToHash(*run.TxHash) + + receipt, err := fp.ethClient.TransactionReceipt(ctx, hash) + if err == ethereum.NotFound { + // TODO rebroadcast logic + fp.logger.WithFields(logrus.Fields{"run_id": run.ID}).Info("tx not found on chain, rebroadcasting") + return nil + } + if receipt.Status == 1 { + if currentBlock > receipt.BlockNumber.Uint64()+fp.config.Jobs.Post.SuccessConfirmations { + fp.logger.WithFields(logrus.Fields{"run_id": run.ID}).Info("tx successful, setting to success") + + ids := []uuid.UUID{} + for _, fee := range run.Fees { + ids = append(ids, fee.ID) + } + + if err = fp.verifierApi.MarkFeeAsCollected(*run.TxHash, run.CreatedAt, ids...); err != nil { + return fmt.Errorf("failed to mark fee as collected on verifier: %w", err) + } + + // This is semi critical code as it could create a state mismatch between the verifier and the database. + if err = fp.db.SetFeeRunSuccess(ctx, run.ID); err != nil { + return fmt.Errorf("failed to set fee run success: %w", err) + } + } else { + fp.logger.WithFields(logrus.Fields{"run_id": run.ID}).Info("tx successful, but not enough confirmations, waiting for more") + return nil + } + } else { + // TODO failed tx logic + fp.logger.WithFields(logrus.Fields{"run_id": run.ID}).Info("tx failed, setting to failed") + return nil + } + return nil +} diff --git a/plugin/fees/transaction.go b/plugin/fees/transaction.go index 26cde70..d104b00 100644 --- a/plugin/fees/transaction.go +++ b/plugin/fees/transaction.go @@ -4,12 +4,14 @@ import ( "context" "crypto/sha256" "encoding/base64" + "errors" "fmt" "math/big" "strconv" "strings" "github.com/google/uuid" + "github.com/sirupsen/logrus" "github.com/vultisig/mobile-tss-lib/tss" "github.com/vultisig/plugin/common" "github.com/vultisig/recipes/chain" @@ -25,12 +27,15 @@ import ( gcommon "github.com/ethereum/go-ethereum/common" "github.com/ethereum/go-ethereum/common/hexutil" etypes "github.com/ethereum/go-ethereum/core/types" + "github.com/vultisig/plugin/internal/types" ) -func (fp *FeePlugin) ProposeTransactions(policy vtypes.PluginPolicy) ([]vtypes.PluginKeysignRequest, error) { +func (fp *FeePlugin) proposeTransactions(ctx context.Context, policy vtypes.PluginPolicy, run types.FeeRun) ([]vtypes.PluginKeysignRequest, error) { + + if policy.ID != run.PolicyID { + return nil, fmt.Errorf("policy id does not match run policy id") + } - // Set config, get encryption secret and then get the vault connected to the fee policy. - ctx := context.Background() vault, err := common.GetVaultFromPolicy(fp.vaultStorage, policy, fp.encryptionSecret) if err != nil { return nil, fmt.Errorf("failed to get vault from policy: %v", err) @@ -103,12 +108,7 @@ func (fp *FeePlugin) ProposeTransactions(policy vtypes.PluginPolicy) ([]vtypes.P return nil, fmt.Errorf("token address does not match usdc address") } - // Here we call the verifier api to get a list of fees that have the same public key as the signed policy document. - feeHistory, err := fp.verifierApi.GetPublicKeysFees(policy.PublicKey) - if err != nil { - return nil, fmt.Errorf("failed to get fee history: %v", err) - } - amount := feeHistory.FeesPendingCollection + amount := run.TotalAmount tx, err := fp.eth.MakeAnyTransfer(ctx, gcommon.HexToAddress(ethAddress), @@ -155,6 +155,10 @@ func (fp *FeePlugin) ProposeTransactions(policy vtypes.PluginPolicy) ([]vtypes.P return txs, nil } +func (fp *FeePlugin) ProposeTransactions(policy vtypes.PluginPolicy) ([]vtypes.PluginKeysignRequest, error) { + return nil, errors.New("not implemented") +} + func (fp *FeePlugin) initSign( ctx context.Context, req vtypes.PluginKeysignRequest, @@ -162,17 +166,6 @@ func (fp *FeePlugin) initSign( runId uuid.UUID, ) error { - if runId != uuid.Nil { - if len(req.Messages) != 1 { - return fmt.Errorf("multiple messages in key sign request, expected 1") - } - - err := fp.db.SetFeeRunSent(ctx, runId, uuid.Nil) // TODO this will be replaced imminently in an upcoming PR - if err != nil { - return fmt.Errorf("failed to update fee run: %w", err) - } - } - sigs, err := fp.signer.Sign(ctx, req) if err != nil { fp.logger.WithError(err).Error("Keysign failed") @@ -185,17 +178,54 @@ func (fp *FeePlugin) initSign( Error("expected only 1 message+sig per request for evm") return fmt.Errorf("failed to sign transaction: invalid signature count: %d", len(sigs)) } + var sig tss.KeysignResponse for _, s := range sigs { sig = s } - err = fp.SigningComplete(ctx, sig, req, pluginPolicy) + txBytes, txErr := hexutilDecode(req.Transaction) + r, rErr := hexutilDecode(sig.R) + s, sErr := hexutilDecode(sig.S) + v, vErr := hexutilDecode(sig.RecoveryID) + if txErr != nil || rErr != nil || sErr != nil || vErr != nil { + return fmt.Errorf("error decoding tx or sigs: %w", errors.Join(txErr, rErr, sErr, vErr)) + } + + txHash, err := getHash(txBytes, r, s, v, fp.config.ChainId) + if err != nil { + return fmt.Errorf("failed to get hash: %w", err) + } + + erc20tx, err := decodeTx(req.Transaction) + if err != nil { + fp.logger.WithError(err).Error("failed to decode tx") + return fmt.Errorf("failed to decode tx: %w", err) + } + + fp.logger.WithFields(logrus.Fields{ + "tx_hash": txHash.Hash().Hex(), + "tx_to": erc20tx.to.Hex(), + "tx_amount": erc20tx.amount.String(), + "tx_token": erc20tx.token.Hex(), + "public_key": pluginPolicy.PublicKey, + }).Info("fee collection transaction") + + tx, err := fp.eth.Send(ctx, txBytes, r, s, v) if err != nil { - fp.logger.WithError(err).Error("failed to complete signing process (broadcast tx)") - return fmt.Errorf("failed to complete signing process: %w", err) + fp.logger.WithError(err).WithField("tx_hex", req.Transaction).Error("fp.eth.Send") + return fmt.Errorf("failed to send transaction: %w", err) + } + + // This is exceptionally important, as if it errors, the transaction will internally be recorded as draft, even after it's been broadcasted + if err := fp.db.SetFeeRunSent(ctx, runId, tx.Hash().Hex()); err != nil { //TODO pass the real tx id + return fmt.Errorf("failed to set fee run sent: %w", err) } + + // Log successful transaction broadcast + fp.logger.WithField("hash", tx.Hash().Hex()).Info("fee collection transaction successfully broadcasted") return nil + } func (fp *FeePlugin) ValidateProposedTransactions(policy vtypes.PluginPolicy, txs []vtypes.PluginKeysignRequest) error { @@ -245,21 +275,5 @@ func (fp *FeePlugin) ValidateProposedTransactions(policy vtypes.PluginPolicy, tx } func (fp *FeePlugin) SigningComplete(ctx context.Context, signature tss.KeysignResponse, signRequest vtypes.PluginKeysignRequest, policy vtypes.PluginPolicy) error { - // Broadcast the signed transaction to the Ethereum network - - tx, err := fp.eth.Send( - ctx, - gcommon.FromHex(signRequest.Transaction), - gcommon.Hex2Bytes(signature.R), - gcommon.Hex2Bytes(signature.S), - gcommon.Hex2Bytes(signature.RecoveryID), - ) - if err != nil { - fp.logger.WithError(err).WithField("tx_hex", signRequest.Transaction).Error("fp.eth.Send") - return fmt.Errorf("failed to send transaction: %w", err) - } - - // Log successful transaction broadcast - fp.logger.WithField("hash", tx.Hash().Hex()).Info("fee collection transaction successfully broadcasted") - return nil + return fmt.Errorf("not implemented") } diff --git a/service/policy.go b/service/policy.go index 3b8eeb9..30a61c4 100644 --- a/service/policy.go +++ b/service/policy.go @@ -130,7 +130,7 @@ func (s *PolicyService) DeletePolicy(ctx context.Context, policyID uuid.UUID, si } func (s *PolicyService) GetPluginPolicies(ctx context.Context, pluginID vtypes.PluginID, publicKey string, onlyActive bool) ([]vtypes.PluginPolicy, error) { - return s.db.GetAllPluginPolicies(ctx, publicKey, pluginID, onlyActive) + return s.db.GetPluginPolicies(ctx, publicKey, pluginID, onlyActive) } func (s *PolicyService) GetPluginPolicy(ctx context.Context, policyID uuid.UUID) (*vtypes.PluginPolicy, error) { diff --git a/storage/db.go b/storage/db.go index 114578d..8a98027 100644 --- a/storage/db.go +++ b/storage/db.go @@ -16,13 +16,20 @@ type DatabaseStorage interface { Close() error GetPluginPolicy(ctx context.Context, id uuid.UUID) (*vtypes.PluginPolicy, error) - GetAllPluginPolicies(ctx context.Context, publicKey string, pluginID vtypes.PluginID, onlyActive bool) ([]vtypes.PluginPolicy, error) + GetPluginPolicies(ctx context.Context, publicKey string, pluginID vtypes.PluginID, onlyActive bool) ([]vtypes.PluginPolicy, error) + GetAllFeePolicies(ctx context.Context) ([]vtypes.PluginPolicy, error) DeletePluginPolicyTx(ctx context.Context, dbTx pgx.Tx, id uuid.UUID) error InsertPluginPolicyTx(ctx context.Context, dbTx pgx.Tx, policy vtypes.PluginPolicy) (*vtypes.PluginPolicy, error) UpdatePluginPolicyTx(ctx context.Context, dbTx pgx.Tx, policy vtypes.PluginPolicy) (*vtypes.PluginPolicy, error) - CreateFeeRun(ctx context.Context, policyId uuid.UUID, state types.FeeRunState, fees []verifierapi.FeeDto) (*types.FeeRun, error) - SetFeeRunSent(ctx context.Context, runId uuid.UUID, txId uuid.UUID) error + CreateFeeRun(ctx context.Context, policyId uuid.UUID, state types.FeeRunState, fees ...verifierapi.FeeDto) (*types.FeeRun, error) + SetFeeRunSent(ctx context.Context, runId uuid.UUID, txHash string) error + SetFeeRunSuccess(ctx context.Context, runId uuid.UUID) error + GetAllFeeRuns(ctx context.Context, statuses ...types.FeeRunState) ([]types.FeeRun, error) // If no statuses are provided, all fee runs are returned. + GetFees(ctx context.Context, feeIds ...uuid.UUID) ([]types.Fee, error) + GetPendingFeeRun(ctx context.Context, policyId uuid.UUID) (*types.FeeRun, error) + CreateFee(ctx context.Context, runId uuid.UUID, fee verifierapi.FeeDto) error + GetFeeRuns(ctx context.Context, state types.FeeRunState) ([]types.FeeRun, error) Pool() *pgxpool.Pool } diff --git a/storage/postgres/fees.go b/storage/postgres/fees.go index 1198c79..4156121 100644 --- a/storage/postgres/fees.go +++ b/storage/postgres/fees.go @@ -7,11 +7,12 @@ import ( "github.com/google/uuid" + "github.com/jackc/pgx/v5" "github.com/vultisig/plugin/internal/types" "github.com/vultisig/plugin/internal/verifierapi" ) -func (p *PostgresBackend) CreateFeeRun(ctx context.Context, policyId uuid.UUID, state types.FeeRunState, fees []verifierapi.FeeDto) (*types.FeeRun, error) { +func (p *PostgresBackend) CreateFeeRun(ctx context.Context, policyId uuid.UUID, state types.FeeRunState, fees ...verifierapi.FeeDto) (*types.FeeRun, error) { // Check policy id is valid query := `select plugin_id from plugin_policies where id = $1` policyrows := p.pool.QueryRow(ctx, query, policyId) @@ -48,7 +49,7 @@ func (p *PostgresBackend) CreateFeeRun(ctx context.Context, policyId uuid.UUID, } var run types.FeeRun - err = p.pool.QueryRow(ctx, `select id, status, created_at, updated_at, tx_id, policy_id, total_amount, fee_count from fee_run_with_totals where id = $1`, runId).Scan(&run.ID, &run.Status, &run.CreatedAt, &run.UpdatedAt, &run.TxID, &run.PolicyID, &run.TotalAmount, &run.FeeCount) + err = p.pool.QueryRow(ctx, `select id, status, created_at, updated_at, tx_hash, policy_id, total_amount, fee_count from fee_run_with_totals where id = $1`, runId).Scan(&run.ID, &run.Status, &run.CreatedAt, &run.UpdatedAt, &run.TxHash, &run.PolicyID, &run.TotalAmount, &run.FeeCount) if err != nil { return nil, fmt.Errorf("failed to get fee run (post commit): %s", err) } @@ -56,10 +57,138 @@ func (p *PostgresBackend) CreateFeeRun(ctx context.Context, policyId uuid.UUID, return &run, nil } -func (p *PostgresBackend) SetFeeRunSent(ctx context.Context, runId uuid.UUID, txId uuid.UUID) error { - _, err := p.pool.Exec(ctx, `update fee_run set status = $1, tx_id = $2 where id = $3`, types.FeeRunStateSent, txId, runId) +func (p *PostgresBackend) SetFeeRunSent(ctx context.Context, runId uuid.UUID, txHash string) error { + _, err := p.pool.Exec(ctx, `update fee_run set status = $1, tx_hash = $2 where id = $3`, types.FeeRunStateSent, txHash, runId) if err != nil { return fmt.Errorf("failed to update fee run: %w", err) } return nil } + +func (p *PostgresBackend) SetFeeRunSuccess(ctx context.Context, runId uuid.UUID) error { + _, err := p.pool.Exec(ctx, `update fee_run set status = $1 where id = $2`, types.FeeRunStateSuccess, runId) + if err != nil { + return fmt.Errorf("failed to update fee run: %w", err) + } + return nil +} + +func (p *PostgresBackend) GetAllFeeRuns(ctx context.Context, statuses ...types.FeeRunState) ([]types.FeeRun, error) { + + var rows pgx.Rows + var err error + + if len(statuses) == 0 { + query := `select id, status, created_at, updated_at, tx_hash, policy_id, total_amount, fee_count from fee_run_with_totals` + rows, err = p.pool.Query(ctx, query) + } else { + query := `select id, status, created_at, updated_at, tx_hash, policy_id, total_amount, fee_count from fee_run_with_totals where status = ANY($1)` + rows, err = p.pool.Query(ctx, query, statuses) + } + + if err != nil { + return nil, err + } + defer rows.Close() + rm := make(map[uuid.UUID]types.FeeRun) + for rows.Next() { + var run types.FeeRun + err := rows.Scan(&run.ID, &run.Status, &run.CreatedAt, &run.UpdatedAt, &run.TxHash, &run.PolicyID, &run.TotalAmount, &run.FeeCount) + if err != nil { + return nil, err + } + rm[run.ID] = run + } + + feesQuery := `select id, fee_run_id, amount from fee` + feesRows, err := p.pool.Query(ctx, feesQuery) + if err != nil { + return nil, err + } + defer feesRows.Close() + for feesRows.Next() { + var fee types.Fee + err := feesRows.Scan(&fee.ID, &fee.FeeRunID, &fee.Amount) + if err != nil { + return nil, err + } + if run, ok := rm[fee.FeeRunID]; !ok { + return nil, fmt.Errorf("fee run not found: %s", fee.FeeRunID) + } else { + run.Fees = append(run.Fees, fee) + rm[fee.FeeRunID] = run + } + } + + runs := make([]types.FeeRun, 0, len(rm)) + for _, run := range rm { + runs = append(runs, run) + } + return runs, nil +} + +func (p *PostgresBackend) GetFees(ctx context.Context, feeIds ...uuid.UUID) ([]types.Fee, error) { + query := `select id, fee_run_id, amount from fee where id = ANY($1)` + rows, err := p.pool.Query(ctx, query, feeIds) + if err != nil { + return nil, err + } + defer rows.Close() + + fees := []types.Fee{} + for rows.Next() { + var fee types.Fee + err := rows.Scan(&fee.ID, &fee.FeeRunID, &fee.Amount) + if err != nil { + return nil, err + } + fees = append(fees, fee) + } + return fees, nil +} + +func (p *PostgresBackend) GetPendingFeeRun(ctx context.Context, policyId uuid.UUID) (*types.FeeRun, error) { + query := `select id, status, created_at, updated_at, tx_hash, policy_id, total_amount, fee_count from fee_run_with_totals where status = 'draft' and policy_id = $1 order by created_at desc limit 1` + rows, err := p.pool.Query(ctx, query, policyId) + if err != nil { + return nil, err + } + defer rows.Close() + if !rows.Next() { + return nil, nil + } + var run types.FeeRun + err = rows.Scan(&run.ID, &run.Status, &run.CreatedAt, &run.UpdatedAt, &run.TxHash, &run.PolicyID, &run.TotalAmount, &run.FeeCount) + if err != nil { + return nil, err + } + return &run, nil +} + +func (p *PostgresBackend) CreateFee(ctx context.Context, runId uuid.UUID, fee verifierapi.FeeDto) error { + _, err := p.pool.Exec(ctx, `insert into fee (id, fee_run_id, amount) values ($1, $2, $3)`, fee.ID, runId, fee.Amount) + if err != nil { + return fmt.Errorf("failed to insert fee: %w", err) + } + return nil +} + +func (p *PostgresBackend) GetFeeRuns(ctx context.Context, state types.FeeRunState) ([]types.FeeRun, error) { + query := `select id, status, created_at, updated_at, tx_hash, policy_id, total_amount, fee_count from fee_run_with_totals where status = $1` + rows, err := p.pool.Query(ctx, query, state) + if err != nil { + return nil, err + } + defer rows.Close() + runs := []types.FeeRun{} + + for rows.Next() { + var run types.FeeRun + err := rows.Scan(&run.ID, &run.Status, &run.CreatedAt, &run.UpdatedAt, &run.TxHash, &run.PolicyID, &run.TotalAmount, &run.FeeCount) + if err != nil { + return nil, err + } + runs = append(runs, run) + } + return runs, nil +} diff --git a/storage/postgres/migrations/plugin/20250630152230_fee_runs.sql b/storage/postgres/migrations/plugin/20250630152230_fee_runs.sql index 277a75d..dedfe34 100644 --- a/storage/postgres/migrations/plugin/20250630152230_fee_runs.sql +++ b/storage/postgres/migrations/plugin/20250630152230_fee_runs.sql @@ -6,7 +6,7 @@ CREATE TABLE IF NOT EXISTS fee_run ( status VARCHAR(50) NOT NULL DEFAULT 'draft' CHECK (status IN ('draft', 'sent', 'completed', 'failed')), created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(), updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(), - tx_id UUID REFERENCES tx_indexer(id) ON DELETE SET NULL, + tx_hash VARCHAR(66), policy_id UUID NOT NULL REFERENCES plugin_policies(id) ON DELETE CASCADE ); @@ -25,13 +25,13 @@ SELECT fr.status, fr.created_at, fr.updated_at, - fr.tx_id, + fr.tx_hash, fr.policy_id, COALESCE(SUM(fi.amount), 0) as total_amount, COUNT(fi.id) as fee_count FROM fee_run fr LEFT JOIN fee fi ON fr.id = fi.fee_run_id -GROUP BY fr.id, fr.status, fr.created_at, fr.updated_at, fr.tx_id, fr.policy_id; +GROUP BY fr.id, fr.status, fr.created_at, fr.updated_at, fr.tx_hash, fr.policy_id; -- Create indexes for better performance CREATE INDEX IF NOT EXISTS idx_fee_run_status ON fee_run(status); diff --git a/storage/postgres/policy.go b/storage/postgres/policy.go index 14f0591..07b1e5b 100644 --- a/storage/postgres/policy.go +++ b/storage/postgres/policy.go @@ -38,7 +38,39 @@ func (p *PostgresBackend) GetPluginPolicy(ctx context.Context, id uuid.UUID) (*v return &policy, nil } -func (p *PostgresBackend) GetAllPluginPolicies(ctx context.Context, publicKey string, pluginID vtypes.PluginID, onlyActive bool) ([]vtypes.PluginPolicy, error) { +func (p *PostgresBackend) GetAllFeePolicies(ctx context.Context) ([]vtypes.PluginPolicy, error) { + query := `SELECT DISTINCT ON(public_key) id, public_key, plugin_id, plugin_version, policy_version, signature, active, recipe + FROM plugin_policies + WHERE plugin_id = 'vultisig-fees-feee' AND active = true + ORDER BY public_key, created_at DESC` + + rows, err := p.pool.Query(ctx, query) + if err != nil { + return nil, err + } + defer rows.Close() + var policies []vtypes.PluginPolicy = []vtypes.PluginPolicy{} + for rows.Next() { + var policy vtypes.PluginPolicy + err := rows.Scan( + &policy.ID, + &policy.PublicKey, + &policy.PluginID, + &policy.PluginVersion, + &policy.PolicyVersion, + &policy.Signature, + &policy.Active, + &policy.Recipe, + ) + if err != nil { + return nil, err + } + policies = append(policies, policy) + } + return policies, nil +} + +func (p *PostgresBackend) GetPluginPolicies(ctx context.Context, publicKey string, pluginID vtypes.PluginID, onlyActive bool) ([]vtypes.PluginPolicy, error) { if p.pool == nil { return nil, fmt.Errorf("database pool is nil") diff --git a/storage/postgres/schema/schema.sql b/storage/postgres/schema/schema.sql index 1c52ce3..f14b9f8 100644 --- a/storage/postgres/schema/schema.sql +++ b/storage/postgres/schema/schema.sql @@ -71,7 +71,7 @@ CREATE TABLE "fee_run" ( "status" character varying(50) DEFAULT 'draft'::character varying NOT NULL, "created_at" timestamp with time zone DEFAULT "now"(), "updated_at" timestamp with time zone DEFAULT "now"(), - "tx_id" "uuid", + "tx_hash" character varying(66), "policy_id" "uuid" NOT NULL, CONSTRAINT "fee_run_status_check" CHECK ((("status")::"text" = ANY ((ARRAY['draft'::character varying, 'sent'::character varying, 'completed'::character varying, 'failed'::character varying])::"text"[]))) ); @@ -81,13 +81,13 @@ CREATE VIEW "fee_run_with_totals" AS "fr"."status", "fr"."created_at", "fr"."updated_at", - "fr"."tx_id", + "fr"."tx_hash", "fr"."policy_id", COALESCE("sum"("fi"."amount"), (0)::bigint) AS "total_amount", "count"("fi"."id") AS "fee_count" FROM ("fee_run" "fr" LEFT JOIN "fee" "fi" ON (("fr"."id" = "fi"."fee_run_id"))) - GROUP BY "fr"."id", "fr"."status", "fr"."created_at", "fr"."updated_at", "fr"."tx_id", "fr"."policy_id"; + GROUP BY "fr"."id", "fr"."status", "fr"."created_at", "fr"."updated_at", "fr"."tx_hash", "fr"."policy_id"; CREATE TABLE "plugin_policies" ( "id" "uuid" DEFAULT "gen_random_uuid"() NOT NULL, @@ -173,6 +173,3 @@ ALTER TABLE ONLY "fee" ALTER TABLE ONLY "fee_run" ADD CONSTRAINT "fee_run_policy_id_fkey" FOREIGN KEY ("policy_id") REFERENCES "plugin_policies"("id") ON DELETE CASCADE; -ALTER TABLE ONLY "fee_run" - ADD CONSTRAINT "fee_run_tx_id_fkey" FOREIGN KEY ("tx_id") REFERENCES "tx_indexer"("id") ON DELETE SET NULL; -