penumbra_wallet/
plan.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
use std::collections::BTreeMap;

use anyhow::Context;
use decaf377::Fq;
use rand_core::{CryptoRng, RngCore};
use tracing::instrument;

use penumbra_dex::swap_claim::SwapClaimPlan;
use penumbra_keys::keys::AddressIndex;
use penumbra_proto::view::v1::NotesRequest;
use penumbra_transaction::{TransactionParameters, TransactionPlan};
pub use penumbra_view::Planner;
use penumbra_view::{SpendableNoteRecord, ViewClient};

pub const SWEEP_COUNT: usize = 8;

#[instrument(skip(view, rng))]
pub async fn sweep<V, R>(view: &mut V, mut rng: R) -> anyhow::Result<Vec<TransactionPlan>>
where
    V: ViewClient,
    R: RngCore + CryptoRng,
{
    let mut plans = Vec::new();

    // First, find any un-claimed swaps and add `SwapClaim` plans for them.
    plans.extend(claim_unclaimed_swaps(view, &mut rng).await?);

    // Finally, sweep dust notes by spending them to their owner's address.
    // This will consolidate small-value notes into larger ones.
    plans.extend(sweep_notes(view, &mut rng).await?);

    Ok(plans)
}

#[instrument(skip(view, rng))]
async fn claim_unclaimed_swaps<V, R>(
    view: &mut V,
    mut rng: R,
) -> anyhow::Result<Vec<TransactionPlan>>
where
    V: ViewClient,
    R: RngCore + CryptoRng,
{
    let mut plans = Vec::new();
    // fetch all transactions
    // check if they contain Swap actions
    // if they do, check if the associated notes are unspent
    // if they are, decrypt the SwapCiphertext in the Swap action and construct a SwapClaim

    let app_params = view.app_params().await?;
    let epoch_duration = app_params.sct_params.epoch_duration;

    let unclaimed_swaps = view.unclaimed_swaps().await?;

    for swap in unclaimed_swaps {
        // We found an unclaimed swap, so we can claim it.
        let swap_plaintext = swap.swap;

        let output_data = swap.output_data;

        let mut plan = TransactionPlan {
            transaction_parameters: TransactionParameters {
                chain_id: app_params.clone().chain_id,
                fee: swap_plaintext.claim_fee.clone(),
                ..Default::default()
            },
            // The transaction doesn't need a memo, because it's to ourselves.
            memo: None,
            ..Default::default()
        };

        let action_plan = SwapClaimPlan {
            swap_plaintext,
            position: swap.position,
            output_data,
            epoch_duration,
            proof_blinding_r: Fq::rand(&mut rng),
            proof_blinding_s: Fq::rand(&mut rng),
        };
        plan.actions.push(action_plan.into());
        plans.push(plan);
    }

    Ok(plans)
}

#[instrument(skip(view, rng))]
async fn sweep_notes<V, R>(view: &mut V, mut rng: R) -> anyhow::Result<Vec<TransactionPlan>>
where
    V: ViewClient,
    R: RngCore + CryptoRng,
{
    let gas_prices = view.gas_prices().await?;

    let all_notes = view
        .notes(NotesRequest {
            ..Default::default()
        })
        .await?;

    let mut notes_by_addr_and_denom: BTreeMap<AddressIndex, BTreeMap<_, Vec<SpendableNoteRecord>>> =
        BTreeMap::new();

    for record in all_notes {
        notes_by_addr_and_denom
            .entry(record.address_index)
            .or_default()
            .entry(record.note.asset_id())
            .or_default()
            .push(record);
    }

    let mut plans = Vec::new();

    for (index, notes_by_denom) in notes_by_addr_and_denom {
        tracing::info!(?index, "processing address");

        for (asset_id, mut records) in notes_by_denom {
            tracing::debug!(?asset_id, "processing asset");

            // Sort notes by amount, ascending, so the biggest notes are at the end...
            records.sort_by(|a, b| a.note.value().amount.cmp(&b.note.value().amount));
            // ... so that when we use chunks_exact, we get SWEEP_COUNT sized
            // chunks, ignoring the biggest notes in the remainder.
            for group in records.chunks_exact(SWEEP_COUNT) {
                let mut planner = Planner::new(&mut rng);
                planner.set_gas_prices(gas_prices);

                for record in group {
                    planner.spend(record.note.clone(), record.position);
                }

                let plan = planner
                    .plan(view, index)
                    .await
                    .context("can't build sweep transaction")?;

                tracing::debug!(?plan);
                plans.push(plan);
            }
        }
    }

    Ok(plans)
}