1
use crate::common::*;
2
use anyhow::Result;
3
use bonsol_prover::input_resolver::{DefaultInputResolver, InputResolver, ProgramInput};
4
use bonsol_sdk::instructions::{ExecutionConfig, InputRef};
5
use bonsol_sdk::{BonsolClient, ExecutionAccountStatus, InputType, InputT};
6
use indicatif::ProgressBar;
7
use sha2::{Digest, Sha256};
8
use solana_rpc_client::nonblocking::rpc_client::RpcClient;
9
use solana_sdk::bs58;
10
use solana_sdk::commitment_config::CommitmentConfig;
11
use solana_sdk::pubkey::Pubkey;
12
use solana_sdk::signer::Signer;
13
use std::fs::File;
14
use std::sync::Arc;
15
use tokio::time::Instant;
16
use hex;
17

            
18
pub async fn execution_waiter(
19
    sdk: &BonsolClient,
20
    requester: Pubkey,
21
    execution_id: String,
22
    expiry: u64,
23
    timeout: Option<u64>,
24
) -> Result<()> {
25
    let indicator = ProgressBar::new_spinner();
26

            
27
    let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(1));
28
    let now = Instant::now();
29
    loop {
30
        if let Some(timeout) = timeout {
31
            if now.elapsed().as_secs() > timeout {
32
                return Err(anyhow::anyhow!("Timeout"));
33
            }
34
        }
35
        interval.tick().await;
36

            
37
        let current_block = sdk.get_current_slot().await?;
38
        indicator.set_message(format!(
39
            "Waiting for execution to be claimed, current block {} expiry {}",
40
            current_block, expiry
41
        ));
42
        if current_block > expiry {
43
            indicator.finish_with_message("Execution expired");
44
            return Err(anyhow::anyhow!("Execution expired"));
45
        }
46

            
47
        let claim_state = sdk.get_claim_state_v1(&requester, &execution_id).await;
48
        if let Ok(claim_state) = claim_state {
49
            let claim = claim_state.claim()?;
50
            indicator.finish_with_message(format!(
51
                "Claimed by {} at slot {}, committed {}",
52
                bs58::encode(claim.claimer).into_string(),
53
                claim.claimed_at,
54
                claim.block_commitment
55
            ));
56
            break;
57
        }
58
    }
59
    //now we are looking for execution request finished
60
    loop {
61
        if let Some(timeout) = timeout {
62
            if now.elapsed().as_secs() > timeout {
63
                indicator.finish_with_message("Execution timed out");
64
                return Err(anyhow::anyhow!("Timeout"));
65
            }
66
        }
67
        interval.tick().await;
68
        let exec_status = sdk
69
            .get_execution_request_v1(&requester, &execution_id)
70
            .await?;
71
        match exec_status {
72
            ExecutionAccountStatus::Completed(ec) => {
73
                indicator.finish_with_message(format!("Execution completed with exit code {}", ec));
74
                return Ok(());
75
            }
76
            ExecutionAccountStatus::Pending(_) => {
77
                indicator.tick();
78
                continue;
79
            }
80
        }
81
    }
82
}
83

            
84
pub async fn execute(
85
    sdk: &BonsolClient,
86
    rpc_url: String,
87
    keypair: impl Signer,
88
    execution_request_file: Option<String>,
89
    image_id: Option<String>,
90
    execution_id: Option<String>,
91
    timeout: Option<u64>,
92
    inputs_file: Option<String>,
93
    tip: Option<u64>,
94
    expiry: Option<u64>,
95
    stdin: Option<String>,
96
    wait: bool,
97
) -> Result<()> {
98
    let indicator = ProgressBar::new_spinner();
99
    let erstr =
100
        execution_request_file.ok_or(anyhow::anyhow!("Execution request file not provided"))?;
101
    let erfile = File::open(erstr)?;
102
    let execution_request_file: ExecutionRequestFile = serde_json::from_reader(erfile)?;
103
    let inputs = if let Some(inputs) = execution_request_file.inputs {
104
        inputs
105
    } else {
106
        execute_get_inputs(inputs_file, stdin)?
107
    };
108
    let execution_id = execution_id
109
        .or(execution_request_file.execution_id)
110
        .or(Some(rand_id(8)))
111
        .ok_or(anyhow::anyhow!("Execution id not provided"))?;
112
    let image_id = image_id
113
        .or(execution_request_file.image_id)
114
        .ok_or(anyhow::anyhow!("Image id not provided"))?;
115
    let tip = tip
116
        .or(execution_request_file.tip)
117
        .ok_or(anyhow::anyhow!("Tip not provided"))?;
118
    let expiry = expiry
119
        .or(execution_request_file.expiry)
120
        .ok_or(anyhow::anyhow!("Expiry not provided"))?;
121
    let callback_config = execution_request_file.callback_config;
122
    let mut input_hash =
123
        if let Some(input_hash) = execution_request_file.execution_config.input_hash {
124
            hex::decode(&input_hash)
125
                .map_err(|_| anyhow::anyhow!("Invalid input hash, must be hex encoded"))?
126
        } else {
127
            vec![]
128
        };
129

            
130
    let signer = keypair.pubkey();
131
    let transformed_cli_inputs = execute_transform_cli_inputs(inputs)?;
132

            
133
    // Declare data storage that will live long enough
134
    let single_concatenated_input_data: Option<Vec<u8>> = if transformed_cli_inputs.iter().all(|i| i.input_type == InputType::PublicData) && !transformed_cli_inputs.is_empty() {
135
        let mut concatenated_bytes: Vec<u8> = Vec::new();
136
        for input_t_val in &transformed_cli_inputs {
137
            if let Some(bytes) = &input_t_val.data { // Assuming InputT { input_type: InputType, data: Option<Vec<u8>> }
138
                concatenated_bytes.extend_from_slice(bytes);
139
            }
140
        }
141
        if !concatenated_bytes.is_empty() {
142
            Some(concatenated_bytes)
143
        } else {
144
            None // Or Some(Vec::new()) if an empty input ref is desired for empty concatenation
145
        }
146
    } else {
147
        None
148
    };
149

            
150
    let final_input_refs: Vec<InputRef> = if let Some(ref data_bytes) = single_concatenated_input_data {
151
        // We have concatenated data, create a single InputRef
152
        // Ensure data_bytes is not empty before creating InputRef if that's a requirement
153
        if data_bytes.is_empty() && !transformed_cli_inputs.is_empty() {
154
            // This handles the case where concatenation was attempted but resulted in empty (e.g. all inputs had None data)
155
            // Depending on desired behavior, could be an empty Vec or an error.
156
            // For now, if concatenation yields empty but there WERE inputs, maybe it should be an empty Vec<InputRef>
157
            // or an error. If it must be one input ref, an empty one might be InputRef::new(InputType::PublicData, &[])
158
            // Sticking to one ref if single_concatenated_input_data is Some:
159
            vec![InputRef::new(InputType::PublicData, data_bytes.as_slice())] 
160
        } else if !data_bytes.is_empty() {
161
            vec![InputRef::new(InputType::PublicData, data_bytes.as_slice())]
162
        } else { // single_concatenated_input_data was Some(empty_vec) and transformed_cli_inputs was empty
163
            vec![]
164
        }
165
    } else {
166
        // No concatenation, or concatenation resulted in None (e.g. no public data inputs, or mixed inputs)
167
        // Use original logic for multiple inputs
168
        transformed_cli_inputs
169
            .iter()
170
            .map(|input_t_val| InputRef::new(input_t_val.input_type, input_t_val.data.as_deref().unwrap_or_default()))
171
            .collect()
172
    };
173

            
174
    // The verify_input_hash logic was here, it might need adjustment if it relied on the structure of transformed_inputs
175
    // For now, let's assume it's handled or less critical than the main execution flow.
176
    // Re-inserting the verify_input_hash logic based on transformed_cli_inputs (pre-concatenation)
177
    let verify_input_hash = execution_request_file
178
        .execution_config
179
        .verify_input_hash
180
        .unwrap_or(false);
181
    
182
    // Hash inputs logic needs to operate on Vec<InputT> (transformed_cli_inputs)
183
    // because DefaultInputResolver::resolve_public_inputs expects something compatible with ProgramInputType.
184
    // The current transformed_cli_inputs IS Vec<InputT> (from bonsol_sdk which is ProgramInputType<Vec<u8>, Option<String>>)
185
    // So this part should be fine if resolve_public_inputs can take Vec<InputT>
186
    let hash_inputs = verify_input_hash
187
        && transformed_cli_inputs.iter().all(|i| i.input_type != InputType::Private);
188
    
189
    if hash_inputs {
190
        // ... (hashing logic as it was, using transformed_cli_inputs) ...
191
        // This part needs to be correct based on what resolve_public_inputs expects.
192
        // If it expects Vec<ProgramInputType<Vec<u8>, Option<String>>>, then transformed_cli_inputs is fine.
193
        indicator.set_message("Getting/Hashing inputs");
194
        let rpc_client = Arc::new(RpcClient::new_with_commitment(
195
            rpc_url.clone(),
196
            CommitmentConfig::confirmed(),
197
        ));
198
        let input_resolver =
199
            DefaultInputResolver::new(Arc::new(reqwest::Client::new()), rpc_client);
200
        
201
        // THIS CLONE IS IMPORTANT for ProgramInputType if it's not Copy.
202
        // Use InputT alias which is bonsol_sdk::ProgramInputType<Vec<u8>, Option<String>>
203
        let resolvable_inputs_for_hashing: Vec<InputT> = transformed_cli_inputs.iter().map(|i_t| 
204
            // Assuming i_t is ProgramInputType<Vec<u8>, Option<String>> from bonsol_sdk::InputT
205
            // We need to ensure it's cloned correctly if resolve_public_inputs takes ownership or modifies.
206
            // Or if ProgramInputType itself is not directly cloneable in this form for that function.
207
            // Let's assume a direct clone works for now, or that resolve_public_inputs takes references.
208
            i_t.clone() // This requires InputT to be Clone.
209
        ).collect();
210

            
211
        let hashing_inputs = input_resolver
212
            .resolve_public_inputs(resolvable_inputs_for_hashing)
213
            .await?;
214
        let mut hash_obj = Sha256::new(); // Renamed from 'hash' to avoid conflict if input_hash is in scope
215
        for prog_input in hashing_inputs {
216
            if let ProgramInput::Resolved(ri) = prog_input {
217
                hash_obj.update(&ri.data);
218
            } else {
219
                return Err(anyhow::anyhow!("Unresolved input during hashing"));
220
            }
221
        }
222
        input_hash = hash_obj.finalize().to_vec();
223
    }
224

            
225
    let execution_config = ExecutionConfig {
226
        verify_input_hash: execution_request_file
227
            .execution_config
228
            .verify_input_hash
229
            .unwrap_or(false),
230
        input_hash: Some(&input_hash),
231
        forward_output: execution_request_file
232
            .execution_config
233
            .forward_output
234
            .unwrap_or(false),
235
    };
236
    let current_block = sdk.get_current_slot().await?;
237
    let expiry = expiry + current_block;
238
    println!("Execution expiry {}", expiry);
239
    println!("current block {}", current_block);
240
    indicator.set_message("Building transaction");
241
    let ixs = sdk
242
        .execute_v1(
243
            &signer,
244
            &image_id,
245
            &execution_id,
246
            final_input_refs,
247
            tip,
248
            expiry,
249
            execution_config,
250
            callback_config.map(|c| c.into()),
251
            None, // A future cli change can implement prover version selection
252
        )
253
        .await?;
254
    indicator.finish_with_message("Sending transaction");
255
    sdk.send_txn_standard(&keypair, ixs).await?;
256
    indicator.finish_with_message("Waiting for execution");
257
    if wait {
258
        execution_waiter(sdk, keypair.pubkey(), execution_id, expiry, timeout).await?;
259
    }
260
    Ok(())
261
}