Skip to content

Commit

Permalink
add support for address generator in mini-wallet
Browse files Browse the repository at this point in the history
We keep the core logic for watching individual addresses, and just
generate more when needed.
  • Loading branch information
pm47 committed Jan 24, 2024
1 parent f41d471 commit ed405ff
Show file tree
Hide file tree
Showing 3 changed files with 135 additions and 27 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package fr.acinq.lightning.blockchain.electrum

import fr.acinq.bitcoin.*
import fr.acinq.lightning.SwapInParams
import fr.acinq.lightning.blockchain.electrum.WalletState.Companion.indexOrNull
import fr.acinq.lightning.utils.sum
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Job
Expand Down Expand Up @@ -89,6 +90,12 @@ data class WalletState(val addresses: Map<String, AddressState>, val parentTxs:
data object Single : AddressMeta()
data class Derived(val index: Int) : AddressMeta()
}

val AddressMeta.indexOrNull: Int?
get() = when (this) {
is AddressMeta.Single -> null
is AddressMeta.Derived -> this.index
}
}
}

Expand All @@ -99,6 +106,10 @@ private sealed interface WalletCommand {
data object ElectrumConnected : WalletCommand
data class ElectrumNotification(val msg: ElectrumResponse) : WalletCommand
data class AddAddress(val bitcoinAddress: String, val meta: WalletState.Companion.AddressMeta) : WalletCommand
data class AddAddressGenerator(val generator: AddressGenerator) : WalletCommand
data class GenerateAddress(val index: Int) : WalletCommand

class AddressGenerator(val generateAddress: (Int) -> String, val window: Int)
}
}

Expand Down Expand Up @@ -130,6 +141,9 @@ class ElectrumMiniWallet(
private val _walletStateFlow = MutableStateFlow(WalletState(emptyMap(), emptyMap()))
val walletStateFlow get() = _walletStateFlow.asStateFlow()

// generator, if used
private var addressGenerator: Pair<WalletCommand.Companion.AddressGenerator, Int>? = null

// all current meta associated to each address
private var addressMetas: Map<String, WalletState.Companion.AddressMeta> = emptyMap()

Expand All @@ -139,9 +153,15 @@ class ElectrumMiniWallet(
// the mailbox of this "actor"
private val mailbox: Channel<WalletCommand> = Channel(Channel.BUFFERED)

fun addAddress(bitcoinAddress: String, meta: WalletState.Companion.AddressMeta) {
fun addAddress(bitcoinAddress: String) {
launch {
mailbox.send(WalletCommand.Companion.AddAddress(bitcoinAddress, WalletState.Companion.AddressMeta.Single))
}
}

fun addAddressGenerator(generator: (Int) -> String, window: Int) {
launch {
mailbox.send(WalletCommand.Companion.AddAddress(bitcoinAddress, meta))
mailbox.send(WalletCommand.Companion.AddAddressGenerator(WalletCommand.Companion.AddressGenerator(generator, window)))
}
}

Expand Down Expand Up @@ -169,8 +189,17 @@ class ElectrumMiniWallet(
parentTxs.forEach { tx -> logger.mdcinfo { "received parent transaction with txid=${tx.txid}" } }
val nextAddressState = WalletState.Companion.AddressState(addressMeta, unspents)
val nextWalletState = this.copy(addresses = this.addresses + (bitcoinAddress to nextAddressState), parentTxs = this.parentTxs + parentTxs.associateBy { it.txid })
logger.mdcinfo { "${unspents.size} utxo(s) for address=$bitcoinAddress balance=${nextWalletState.totalBalance}" }
logger.mdcinfo { "${unspents.size} utxo(s) for address=$bitcoinAddress index=${addressMeta.indexOrNull ?: "n/a"} balance=${nextWalletState.totalBalance}" }
unspents.forEach { logger.debug { "utxo=${it.outPoint.txid}:${it.outPoint.index} amount=${it.value} sat" } }
when (val generator = addressGenerator) {
null -> {}
else -> if (addressMeta.indexOrNull == generator.second - 1 && nextAddressState.unspent.isNotEmpty()) {
logger.info { "requesting generation of next sequence of addresses" }
// if the window = 10 then the first series of address is 0 to 9 and addressGenerator.second = 10
// then when address 9 has utxos, we request generation up until (and excluding) index 10 + 10 = 20, this will generate addresses 10 to 19
mailbox.send(WalletCommand.Companion.GenerateAddress(addressMeta.indexOrNull!! + 1 + generator.first.window))
}
}
nextWalletState
}
}
Expand Down Expand Up @@ -225,13 +254,30 @@ class ElectrumMiniWallet(
is WalletCommand.Companion.AddAddress -> {
computeScriptHash(it.bitcoinAddress)?.let { scriptHash ->
if (!scriptHashes.containsKey(scriptHash)) {
logger.mdcinfo { "adding new address=${it.bitcoinAddress}" }
logger.mdcinfo { "adding new address=${it.bitcoinAddress} index=${it.meta.indexOrNull ?: "n/a"}" }
scriptHashes = scriptHashes + (scriptHash to it.bitcoinAddress)
addressMetas = addressMetas + (it.bitcoinAddress to it.meta)
subscribe(scriptHash, it.bitcoinAddress)
}
}
}
is WalletCommand.Companion.AddAddressGenerator -> {
if (addressGenerator == null) {
logger.mdcinfo { "adding new address generator" }
addressGenerator = it.generator to 0
mailbox.send(WalletCommand.Companion.GenerateAddress(it.generator.window))
}
}
is WalletCommand.Companion.GenerateAddress -> {
addressGenerator = addressGenerator?.let { (generator, currentIndex) ->
logger.info { "generating addresses with index $currentIndex to ${it.index - 1}" }
(currentIndex until it.index).forEach { addressIndex ->
mailbox.send(WalletCommand.Companion.AddAddress(generator.generateAddress(addressIndex), WalletState.Companion.AddressMeta.Derived(addressIndex)))
}
generator to maxOf(currentIndex, it.index)
}
}

}
}
}
Expand Down
27 changes: 15 additions & 12 deletions src/commonMain/kotlin/fr/acinq/lightning/io/Peer.kt
Original file line number Diff line number Diff line change
Expand Up @@ -190,17 +190,13 @@ class Peer(
}

val finalWallet = ElectrumMiniWallet(nodeParams.chainHash, watcher.client, scope, nodeParams.loggerFactory, name = "final")
val finalAddress: String = nodeParams.keyManager.finalOnChainWallet.address(addressIndex = 0L).also { finalWallet.addAddress(it, WalletState.Companion.AddressMeta.Single) }
val finalAddress: String = nodeParams.keyManager.finalOnChainWallet.address(addressIndex = 0L).also { finalWallet.addAddress(it) }

val swapInWallet = ElectrumMiniWallet(nodeParams.chainHash, watcher.client, scope, nodeParams.loggerFactory, name = "swap-in")
val legacySwapInAddress: String = nodeParams.keyManager.swapInOnChainWallet.legacySwapInProtocol.address(nodeParams.chain).also { swapInWallet.addAddress(it, WalletState.Companion.AddressMeta.Single) }
val legacySwapInAddress: String = nodeParams.keyManager.swapInOnChainWallet.legacySwapInProtocol.address(nodeParams.chain)
.also { swapInWallet.addAddress(it) }
val swapInAddressFlow = MutableStateFlow<String?>(null)
val swapInAddresses: List<String> = (0 until 100)
.map { nodeParams.keyManager.swapInOnChainWallet.getSwapInProtocol(it) }
.map { it.address(nodeParams.chain) }
.withIndex().onEach { swapInWallet.addAddress(it.value, WalletState.Companion.AddressMeta.Derived(it.index)) }
.map { it.value }
.toList()
.also { swapInWallet.addAddressGenerator({ index -> nodeParams.keyManager.swapInOnChainWallet.getSwapInProtocol(index).address(nodeParams.chain) }, 20) }

private var swapInJob: Job? = null

Expand Down Expand Up @@ -233,6 +229,17 @@ class Peer(
logger.info { "${wallet.totalBalance} available on final wallet with ${wallet.utxos.size} utxos" }
}
}
launch {
// address rotation
swapInWallet.walletStateFlow
.onEach { it ->
// take the first unused address with the lowest index
swapInAddressFlow.value = it.addresses
.filter { it.value.unspent.isEmpty() && it.value.meta is WalletState.Companion.AddressMeta.Derived }
.minByOrNull { (it.value.meta as WalletState.Companion.AddressMeta.Derived).index }
?.key
}
}
launch {
// we don't restore closed channels
val bootChannels = db.channels.listLocalChannels().filterNot { it is Closed || it is LegacyWaitForFundingConfirmed }
Expand Down Expand Up @@ -462,10 +469,6 @@ class Peer(
swapInJob = launch {
swapInWallet.walletStateFlow
.filter { it.consistent }
.onEach {
// take the first unused address, or a random address if there are none
swapInAddressFlow.value = it.addresses.filter { it.value.unspent.isEmpty() }.map { it.key }.firstOrNull() ?: it.addresses.keys.random()
}
.combine(currentTipFlow.filterNotNull()) { walletState, currentTip -> Pair(walletState, currentTip.first) }
.combine(swapInFeeratesFlow.filterNotNull()) { (walletState, currentTip), feerate -> Triple(walletState, currentTip, feerate) }
.combine(nodeParams.liquidityPolicy) { (walletState, currentTip, feerate), policy -> TrySwapInFlow(currentTip, walletState, feerate, policy) }
Expand Down
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
package fr.acinq.lightning.blockchain.electrum

import fr.acinq.bitcoin.Bitcoin
import fr.acinq.bitcoin.Block
import fr.acinq.bitcoin.Transaction
import fr.acinq.bitcoin.TxId
import fr.acinq.bitcoin.*
import fr.acinq.lightning.Lightning.randomKey
import fr.acinq.lightning.SwapInParams
import fr.acinq.lightning.tests.utils.LightningTestSuite
import fr.acinq.lightning.tests.utils.runSuspendTest
Expand All @@ -22,7 +20,7 @@ class ElectrumMiniWalletTest : LightningTestSuite() {
fun `single address with no utxos`() = runSuspendTest(timeout = 15.seconds) {
val client = connectToMainnetServer()
val wallet = ElectrumMiniWallet(Block.LivenetGenesisBlock.hash, client, this, LoggerFactory.default)
wallet.addAddress("bc1qyjmhaptq78vh5j7tnzu7ujayd8sftjahphxppz", WalletState.Companion.AddressMeta.Single)
wallet.addAddress("bc1qyjmhaptq78vh5j7tnzu7ujayd8sftjahphxppz")

val walletState = wallet.walletStateFlow
.filter { it.addresses.size == 1 && it.consistent }
Expand All @@ -39,7 +37,7 @@ class ElectrumMiniWalletTest : LightningTestSuite() {
fun `single address with existing utxos`() = runSuspendTest(timeout = 15.seconds) {
val client = connectToMainnetServer()
val wallet = ElectrumMiniWallet(Block.LivenetGenesisBlock.hash, client, this, LoggerFactory.default)
wallet.addAddress("14xb2HATmkBzrHf4CR2hZczEtjYpTh92d2", WalletState.Companion.AddressMeta.Single)
wallet.addAddress("14xb2HATmkBzrHf4CR2hZczEtjYpTh92d2")

val walletState = wallet.walletStateFlow
.filter { it.addresses.size == 1 && it.consistent }
Expand Down Expand Up @@ -107,9 +105,9 @@ class ElectrumMiniWalletTest : LightningTestSuite() {
fun `multiple addresses`() = runSuspendTest(timeout = 15.seconds) {
val client = connectToMainnetServer()
val wallet = ElectrumMiniWallet(Block.LivenetGenesisBlock.hash, client, this, LoggerFactory.default)
wallet.addAddress("16MmJT8VqW465GEyckWae547jKVfMB14P8", WalletState.Companion.AddressMeta.Single)
wallet.addAddress("14xb2HATmkBzrHf4CR2hZczEtjYpTh92d2", WalletState.Companion.AddressMeta.Single)
wallet.addAddress("1NHFyu1uJ1UoDjtPjqZ4Et3wNCyMGCJ1qV", WalletState.Companion.AddressMeta.Single)
wallet.addAddress("16MmJT8VqW465GEyckWae547jKVfMB14P8")
wallet.addAddress("14xb2HATmkBzrHf4CR2hZczEtjYpTh92d2")
wallet.addAddress("1NHFyu1uJ1UoDjtPjqZ4Et3wNCyMGCJ1qV")

val walletState = wallet.walletStateFlow
.filter { it.parentTxs.size == 11 }
Expand Down Expand Up @@ -156,13 +154,74 @@ class ElectrumMiniWalletTest : LightningTestSuite() {
client.stop()
}

@Test
fun `multiple addresses with generator`() = runSuspendTest(timeout = 15.seconds) {
val client = connectToMainnetServer()
val wallet = ElectrumMiniWallet(Block.LivenetGenesisBlock.hash, client, this, LoggerFactory.default)
wallet.addAddressGenerator(
{
when (it) {
0 -> "16MmJT8VqW465GEyckWae547jKVfMB14P8"
1 -> "14xb2HATmkBzrHf4CR2hZczEtjYpTh92d2"
2 -> "1NHFyu1uJ1UoDjtPjqZ4Et3wNCyMGCJ1qV"
else -> Bitcoin.addressFromPublicKeyScript(Block.LivenetGenesisBlock.hash, Script.pay2pkh(randomKey().publicKey())).result!!
}
},
10
)

val walletState = wallet.walletStateFlow
.filter { it.parentTxs.size == 11 }
.first()

// this has been checked on the blockchain
assertEquals(4 + 6 + 1, walletState.utxos.size)
assertEquals(72_000_000.sat + 30_000_000.sat + 2_000_000.sat, walletState.totalBalance)
assertEquals(11, walletState.utxos.size)
// make sure txid is correct has electrum api is confusing
walletState.parentTxs.forEach { assertEquals(it.key, it.value.txid) }
assertContains(
walletState.utxos,
WalletState.Utxo(
previousTx = Transaction.read("0100000001758713310361270b5ec4cae9b0196cb84fdb2f174d29f9367ad341963fa83e56010000008b483045022100d7b8759aeffe9d829a5df062420eb25017d7341244e49cfede16136a0c0b8dd2022031b42048e66b1f82f7fa99a22954e2709269838ef587c20118e493ced0d63e21014104b9251638d1475b9c62e1cf03129c835bcd5ab843aa0016412e8b39e3f8f7188d3b59023bce2002a2e409ea070c7070392b65d9ae8c8631ae2672a8fbb4f62bbdffffffff02404b4c00000000001976a9143675767783fdf1922f57ab4bb783f3a88dfa609488ac404b4c00000000001976a9142b6ba7c9d796b75eef7942fc9288edd37c32f5c388ac00000000"),
outputIndex = 1,
blockHeight = 100_003,
addressMeta = WalletState.Companion.AddressMeta.Derived(1)
)
)

assertEquals(
expected = setOf(
Triple("16MmJT8VqW465GEyckWae547jKVfMB14P8", TxId("c1e943938e0bf2e9e6feefe22af0466514a58e9f7ed0f7ada6fd8e6dbeca0742") to 1, 39_000_000.sat),
Triple("16MmJT8VqW465GEyckWae547jKVfMB14P8", TxId("2cf392ecf573a638f01f72c276c3b097d05eb58f39e165eacc91b8a8df09fbd8") to 0, 12_000_000.sat),
Triple("16MmJT8VqW465GEyckWae547jKVfMB14P8", TxId("149a098d6261b7f9359a572d797c4a41b62378836a14093912618b15644ba402") to 1, 11_000_000.sat),
Triple("16MmJT8VqW465GEyckWae547jKVfMB14P8", TxId("2dd9cb7bcebb74b02efc85570a462f22a54a613235bee11d0a2c791342a26007") to 1, 10_000_000.sat),
Triple("14xb2HATmkBzrHf4CR2hZczEtjYpTh92d2", TxId("71b3dbaca67e9f9189dad3617138c19725ab541ef0b49c05a94913e9f28e3f4e") to 0, 5_000_000.sat),
Triple("14xb2HATmkBzrHf4CR2hZczEtjYpTh92d2", TxId("21d2eb195736af2a40d42107e6abd59c97eb6cffd4a5a7a7709e86590ae61987") to 0, 5_000_000.sat),
Triple("14xb2HATmkBzrHf4CR2hZczEtjYpTh92d2", TxId("74d681e0e03bafa802c8aa084379aa98d9fcd632ddc2ed9782b586ec87451f20") to 1, 5_000_000.sat),
Triple("14xb2HATmkBzrHf4CR2hZczEtjYpTh92d2", TxId("563ea83f9641d37a36f9294d172fdb4fb86c19b0e9cac45e0b27610331138775") to 0, 5_000_000.sat),
Triple("14xb2HATmkBzrHf4CR2hZczEtjYpTh92d2", TxId("971af80218684017722429be08548d1f30a2f1f220abc064380cbca5cabf7623") to 1, 5_000_000.sat),
Triple("14xb2HATmkBzrHf4CR2hZczEtjYpTh92d2", TxId("b1ec9c44009147f3cee26caba45abec2610c74df9751fad14074119b5314da21") to 0, 5_000_000.sat),
Triple("1NHFyu1uJ1UoDjtPjqZ4Et3wNCyMGCJ1qV", TxId("602839d82ac6c9aafd1a20fff5b23e11a99271e7cc238d2e48b352219b2b87ab") to 1, 2_000_000.sat),
),
actual = walletState.utxos.map {
val txOut = it.previousTx.txOut[it.outputIndex]
val address = Bitcoin.addressFromPublicKeyScript(Block.LivenetGenesisBlock.hash, txOut.publicKeyScript.toByteArray()).result!!
Triple(address, it.previousTx.txid to it.outputIndex, txOut.amount)
}.toSet()
)

wallet.stop()
client.stop()
}

@Test
fun `parallel wallets`() = runSuspendTest(timeout = 15.seconds) {
val client = connectToMainnetServer()
val wallet1 = ElectrumMiniWallet(Block.LivenetGenesisBlock.hash, client, this, LoggerFactory.default, name = "addr-16MmJT")
val wallet2 = ElectrumMiniWallet(Block.LivenetGenesisBlock.hash, client, this, LoggerFactory.default, name = "addr-14xb2H")
wallet1.addAddress("16MmJT8VqW465GEyckWae547jKVfMB14P8", WalletState.Companion.AddressMeta.Single)
wallet2.addAddress("14xb2HATmkBzrHf4CR2hZczEtjYpTh92d2", WalletState.Companion.AddressMeta.Single)
wallet1.addAddress("16MmJT8VqW465GEyckWae547jKVfMB14P8")
wallet2.addAddress("14xb2HATmkBzrHf4CR2hZczEtjYpTh92d2")

val walletState1 = wallet1.walletStateFlow.filter { it.parentTxs.size == 4 }.first()
val walletState2 = wallet2.walletStateFlow.filter { it.parentTxs.size == 6 }.first()
Expand Down

0 comments on commit ed405ff

Please sign in to comment.