|
| 1 | +package requester |
| 2 | + |
| 3 | +import ( |
| 4 | + "fmt" |
| 5 | + "sync" |
| 6 | + "time" |
| 7 | + |
| 8 | + flowsdk "github.com/onflow/flow-go-sdk" |
| 9 | + "github.com/onflow/flow-go-sdk/crypto" |
| 10 | + flowGo "github.com/onflow/flow-go/model/flow" |
| 11 | +) |
| 12 | + |
| 13 | +var ErrNoKeysAvailable = fmt.Errorf("no keys available") |
| 14 | + |
| 15 | +const accountKeyExpiry = 10 * time.Second |
| 16 | + |
| 17 | +type AccountKey struct { |
| 18 | + flowsdk.AccountKey |
| 19 | + |
| 20 | + mu sync.Mutex |
| 21 | + ks *KeyStore |
| 22 | + Address flowsdk.Address |
| 23 | + Signer crypto.Signer |
| 24 | + lastUsed time.Time |
| 25 | +} |
| 26 | + |
| 27 | +// Done unlocks a key after use and puts it back into the pool. |
| 28 | +func (k *AccountKey) Done() { |
| 29 | + k.mu.Lock() |
| 30 | + defer k.mu.Unlock() |
| 31 | + |
| 32 | + k.ks.availableKeys <- k |
| 33 | +} |
| 34 | + |
| 35 | +func (k *AccountKey) SetProposerPayerAndSign( |
| 36 | + tx *flowsdk.Transaction, |
| 37 | + account *flowsdk.Account, |
| 38 | +) error { |
| 39 | + if k.Address != account.Address { |
| 40 | + return fmt.Errorf( |
| 41 | + "expected address: %v, got address: %v", |
| 42 | + k.Address, |
| 43 | + account.Address, |
| 44 | + ) |
| 45 | + } |
| 46 | + if k.Index >= uint32(len(account.Keys)) { |
| 47 | + return fmt.Errorf( |
| 48 | + "key index: %d exceeds keys length: %d", |
| 49 | + k.Index, |
| 50 | + len(account.Keys), |
| 51 | + ) |
| 52 | + } |
| 53 | + seqNumber := account.Keys[k.Index].SequenceNumber |
| 54 | + |
| 55 | + return tx. |
| 56 | + SetProposalKey(k.Address, k.Index, seqNumber). |
| 57 | + SetPayer(k.Address). |
| 58 | + SignEnvelope(k.Address, k.Index, k.Signer) |
| 59 | +} |
| 60 | + |
| 61 | +func (k *AccountKey) expired() bool { |
| 62 | + return time.Since(k.lastUsed) > flowGo.DefaultTransactionExpiry |
| 63 | +} |
| 64 | + |
| 65 | +type KeyLock interface { |
| 66 | + LockKey(txID flowsdk.Identifier, key *AccountKey) |
| 67 | + UnlockKey(txID flowsdk.Identifier) |
| 68 | +} |
| 69 | + |
| 70 | +type KeyStore struct { |
| 71 | + availableKeys chan *AccountKey |
| 72 | + usedKeys map[flowsdk.Identifier]*AccountKey |
| 73 | + size int |
| 74 | +} |
| 75 | + |
| 76 | +var _ KeyLock = (*KeyStore)(nil) |
| 77 | + |
| 78 | +func NewKeyStore(keys []*AccountKey) *KeyStore { |
| 79 | + ks := &KeyStore{ |
| 80 | + usedKeys: map[flowsdk.Identifier]*AccountKey{}, |
| 81 | + } |
| 82 | + |
| 83 | + availableKeys := make(chan *AccountKey, len(keys)) |
| 84 | + for _, key := range keys { |
| 85 | + key.ks = ks |
| 86 | + availableKeys <- key |
| 87 | + } |
| 88 | + ks.size = len(keys) |
| 89 | + ks.availableKeys = availableKeys |
| 90 | + |
| 91 | + go ks.keyExpiryChecker() |
| 92 | + |
| 93 | + return ks |
| 94 | +} |
| 95 | + |
| 96 | +func (k *KeyStore) AvailableKeys() int { |
| 97 | + return k.size - len(k.usedKeys) |
| 98 | +} |
| 99 | + |
| 100 | +func (k *KeyStore) Take() (*AccountKey, error) { |
| 101 | + select { |
| 102 | + case key := <-k.availableKeys: |
| 103 | + return key, nil |
| 104 | + default: |
| 105 | + return nil, ErrNoKeysAvailable |
| 106 | + } |
| 107 | +} |
| 108 | + |
| 109 | +func (k *KeyStore) LockKey(txID flowsdk.Identifier, key *AccountKey) { |
| 110 | + key.mu.Lock() |
| 111 | + defer key.mu.Unlock() |
| 112 | + |
| 113 | + key.lastUsed = time.Now() |
| 114 | + k.usedKeys[txID] = key |
| 115 | +} |
| 116 | + |
| 117 | +func (k *KeyStore) UnlockKey(txID flowsdk.Identifier) { |
| 118 | + key, ok := k.usedKeys[txID] |
| 119 | + if ok && key != nil { |
| 120 | + key.Done() |
| 121 | + delete(k.usedKeys, txID) |
| 122 | + } |
| 123 | +} |
| 124 | + |
| 125 | +func (k *KeyStore) keyExpiryChecker() { |
| 126 | + for range time.Tick(accountKeyExpiry) { |
| 127 | + for txID, key := range k.usedKeys { |
| 128 | + if key.expired() { |
| 129 | + k.UnlockKey(txID) |
| 130 | + } |
| 131 | + } |
| 132 | + } |
| 133 | +} |
0 commit comments