1use crate::{
2 bundle::{
3 jito::{config::JitoConfig, constant::JITO_MOCK_BLOCK_ENGINE_URL, error::JitoError},
4 BundleError,
5 },
6 sanitize_error,
7 transaction::VersionedTransactionResolved,
8};
9use base64::{prelude::BASE64_STANDARD, Engine};
10use bincode::serialize;
11use reqwest::Client;
12use serde::{Deserialize, Serialize};
13use serde_json::{json, Value};
14
15pub enum JitoBundleClient {
16 Live(JitoClient),
17 Mock(JitoMockClient),
18}
19
20impl JitoBundleClient {
21 pub fn new(config: &JitoConfig) -> Self {
22 if config.block_engine_url == JITO_MOCK_BLOCK_ENGINE_URL {
23 Self::Mock(JitoMockClient::new())
24 } else {
25 Self::Live(JitoClient::new(config))
26 }
27 }
28
29 pub async fn send_bundle(
30 &self,
31 transactions: &[VersionedTransactionResolved],
32 ) -> Result<String, BundleError> {
33 match self {
34 Self::Live(client) => client.send_bundle(transactions).await,
35 Self::Mock(client) => client.send_bundle(transactions).await,
36 }
37 }
38
39 pub async fn get_bundle_statuses(
40 &self,
41 bundle_uuids: Vec<String>,
42 ) -> Result<Value, BundleError> {
43 match self {
44 Self::Live(client) => client.get_bundle_statuses(bundle_uuids).await,
45 Self::Mock(client) => client.get_bundle_statuses(bundle_uuids).await,
46 }
47 }
48}
49
50pub struct JitoClient {
51 block_engine_url: String,
52 client: Client,
53}
54
55#[derive(Debug, Serialize, Deserialize)]
56struct JsonRpcRequest {
57 jsonrpc: String,
58 id: u64,
59 method: String,
60 params: Value,
61}
62
63#[derive(Debug, Deserialize)]
64struct JsonRpcResponse {
65 result: Option<Value>,
66 error: Option<JsonRpcError>,
67}
68
69#[derive(Debug, Deserialize)]
70struct JsonRpcError {
71 message: String,
72}
73
74impl JitoClient {
75 pub fn new(config: &JitoConfig) -> Self {
76 Self { block_engine_url: config.block_engine_url.clone(), client: Client::new() }
77 }
78
79 async fn send_request(&self, method: &str, params: Value) -> Result<Value, JitoError> {
80 let url = format!("{}/api/v1/bundles", self.block_engine_url);
81
82 let request = JsonRpcRequest {
83 jsonrpc: "2.0".to_string(),
84 id: 1,
85 method: method.to_string(),
86 params,
87 };
88
89 let response = self
90 .client
91 .post(&url)
92 .header("Content-Type", "application/json")
93 .json(&request)
94 .send()
95 .await
96 .map_err(|e| {
97 JitoError::ApiError(format!("HTTP request failed: {}", sanitize_error!(e)))
98 })?;
99
100 let status = response.status();
101 if !status.is_success() {
102 return Err(JitoError::ApiError(format!("HTTP error: {}", status)));
103 }
104
105 let rpc_response: JsonRpcResponse = response.json().await.map_err(|e| {
106 JitoError::ApiError(format!("Failed to parse response: {}", sanitize_error!(e)))
107 })?;
108
109 if let Some(error) = rpc_response.error {
110 return Err(JitoError::ApiError(error.message));
111 }
112
113 rpc_response
114 .result
115 .ok_or_else(|| JitoError::ApiError("Empty response from Jito".to_string()))
116 }
117
118 pub async fn send_bundle(
119 &self,
120 transactions: &[VersionedTransactionResolved],
121 ) -> Result<String, BundleError> {
122 let mut encoded_txs = Vec::with_capacity(transactions.len());
123 for resolved in transactions {
124 let serialized = serialize(&resolved.transaction)
125 .map_err(|e| BundleError::SerializationError(e.to_string()))?;
126 encoded_txs.push(BASE64_STANDARD.encode(&serialized));
127 }
128
129 let params = json!([encoded_txs, {"encoding": "base64"}]);
130 let result = self.send_request("sendBundle", params).await?;
131
132 result.as_str().map(|s| s.to_string()).ok_or_else(|| {
133 JitoError::ApiError("Invalid bundle UUID in response".to_string()).into()
134 })
135 }
136
137 pub async fn get_bundle_statuses(
138 &self,
139 bundle_uuids: Vec<String>,
140 ) -> Result<Value, BundleError> {
141 let params = json!([bundle_uuids]);
142 Ok(self.send_request("getBundleStatuses", params).await?)
143 }
144}
145
146pub struct JitoMockClient;
147
148impl JitoMockClient {
149 pub fn new() -> Self {
150 Self
151 }
152
153 pub async fn send_bundle(
154 &self,
155 _transactions: &[VersionedTransactionResolved],
156 ) -> Result<String, BundleError> {
157 let random_id: u64 = rand::random();
158 Ok(format!("mock-bundle-{random_id}"))
159 }
160
161 pub async fn get_bundle_statuses(
162 &self,
163 bundle_uuids: Vec<String>,
164 ) -> Result<Value, BundleError> {
165 let mock_statuses: Vec<Value> = bundle_uuids
166 .iter()
167 .map(|uuid| {
168 json!({
169 "bundle_id": uuid,
170 "status": "Landed",
171 "landed_slot": 12345678
172 })
173 })
174 .collect();
175 Ok(json!({ "value": mock_statuses }))
176 }
177}
178
179impl Default for JitoMockClient {
180 fn default() -> Self {
181 Self::new()
182 }
183}
184
185#[cfg(test)]
186mod tests {
187 use super::*;
188 use crate::tests::transaction_mock::create_mock_resolved_transaction;
189 use mockito::{Matcher, Server};
190
191 #[tokio::test]
192 async fn test_send_bundle_success() {
193 let mut server = Server::new_async().await;
194
195 let mock = server
196 .mock("POST", "/api/v1/bundles")
197 .match_header("content-type", "application/json")
198 .match_body(Matcher::PartialJson(json!({"method": "sendBundle"})))
199 .with_status(200)
200 .with_header("content-type", "application/json")
201 .with_body(r#"{"jsonrpc":"2.0","id":1,"result":"bundle-uuid-12345"}"#)
202 .create();
203
204 let config = JitoConfig { block_engine_url: server.url() };
205 let client = JitoClient::new(&config);
206
207 let tx = create_mock_resolved_transaction();
208
209 let result = client.send_bundle(&[tx]).await;
210 mock.assert();
211 assert!(result.is_ok());
212 assert_eq!(result.unwrap(), "bundle-uuid-12345");
213 }
214
215 #[tokio::test]
216 async fn test_send_bundle_http_error() {
217 let mut server = Server::new_async().await;
218
219 let mock = server
220 .mock("POST", "/api/v1/bundles")
221 .with_status(500)
222 .with_body("Internal Server Error")
223 .create();
224
225 let config = JitoConfig { block_engine_url: server.url() };
226 let client = JitoClient::new(&config);
227
228 let tx = create_mock_resolved_transaction();
229
230 let result = client.send_bundle(&[tx]).await;
231 mock.assert();
232 assert!(result.is_err());
233 }
234
235 #[tokio::test]
236 async fn test_send_bundle_rpc_error() {
237 let mut server = Server::new_async().await;
238
239 let mock = server
240 .mock("POST", "/api/v1/bundles")
241 .with_status(200)
242 .with_header("content-type", "application/json")
243 .with_body(r#"{"jsonrpc":"2.0","id":1,"error":{"message":"Bundle rejected"}}"#)
244 .create();
245
246 let config = JitoConfig { block_engine_url: server.url() };
247 let client = JitoClient::new(&config);
248
249 let tx = create_mock_resolved_transaction();
250
251 let result = client.send_bundle(&[tx]).await;
252 mock.assert();
253 assert!(result.is_err());
254 }
255
256 #[tokio::test]
257 async fn test_get_bundle_statuses_success() {
258 let mut server = Server::new_async().await;
259
260 let mock = server
261 .mock("POST", "/api/v1/bundles")
262 .match_body(Matcher::PartialJson(json!({"method": "getBundleStatuses"})))
263 .with_status(200)
264 .with_header("content-type", "application/json")
265 .with_body(r#"{"jsonrpc":"2.0","id":1,"result":{"value":[{"bundle_id":"uuid1","status":"Landed"}]}}"#)
266 .create();
267
268 let config = JitoConfig { block_engine_url: server.url() };
269 let client = JitoClient::new(&config);
270
271 let result = client.get_bundle_statuses(vec!["uuid1".to_string()]).await;
272 mock.assert();
273 assert!(result.is_ok());
274 }
275
276 #[tokio::test]
277 async fn test_send_bundle_invalid_uuid_response() {
278 let mut server = Server::new_async().await;
279
280 let mock = server
281 .mock("POST", "/api/v1/bundles")
282 .with_status(200)
283 .with_header("content-type", "application/json")
284 .with_body(r#"{"jsonrpc":"2.0","id":1,"result":12345}"#)
285 .create();
286
287 let config = JitoConfig { block_engine_url: server.url() };
288 let client = JitoClient::new(&config);
289
290 let tx = create_mock_resolved_transaction();
291
292 let result = client.send_bundle(&[tx]).await;
293 mock.assert();
294 assert!(result.is_err());
295 }
296
297 #[tokio::test]
298 async fn test_send_bundle_empty_result() {
299 let mut server = Server::new_async().await;
300
301 let mock = server
302 .mock("POST", "/api/v1/bundles")
303 .with_status(200)
304 .with_header("content-type", "application/json")
305 .with_body(r#"{"jsonrpc":"2.0","id":1}"#)
306 .create();
307
308 let config = JitoConfig { block_engine_url: server.url() };
309 let client = JitoClient::new(&config);
310
311 let tx = create_mock_resolved_transaction();
312
313 let result = client.send_bundle(&[tx]).await;
314 mock.assert();
315 assert!(result.is_err());
316 }
317
318 #[tokio::test]
319 async fn test_send_bundle_multiple_transactions() {
320 let mut server = Server::new_async().await;
321
322 let mock = server
323 .mock("POST", "/api/v1/bundles")
324 .with_status(200)
325 .with_header("content-type", "application/json")
326 .with_body(r#"{"jsonrpc":"2.0","id":1,"result":"multi-tx-bundle-uuid"}"#)
327 .create();
328
329 let config = JitoConfig { block_engine_url: server.url() };
330 let client = JitoClient::new(&config);
331
332 let txs = vec![
333 create_mock_resolved_transaction(),
334 create_mock_resolved_transaction(),
335 create_mock_resolved_transaction(),
336 ];
337
338 let result = client.send_bundle(&txs).await;
339 mock.assert();
340 assert!(result.is_ok());
341 assert_eq!(result.unwrap(), "multi-tx-bundle-uuid");
342 }
343
344 #[tokio::test]
345 async fn test_send_bundle_malformed_json_response() {
346 let mut server = Server::new_async().await;
347
348 let mock = server
349 .mock("POST", "/api/v1/bundles")
350 .with_status(200)
351 .with_header("content-type", "application/json")
352 .with_body(r#"not valid json"#)
353 .create();
354
355 let config = JitoConfig { block_engine_url: server.url() };
356 let client = JitoClient::new(&config);
357
358 let tx = create_mock_resolved_transaction();
359
360 let result = client.send_bundle(&[tx]).await;
361 mock.assert();
362 assert!(result.is_err());
363 }
364
365 #[tokio::test]
366 async fn test_get_bundle_statuses_empty_uuids() {
367 let mut server = Server::new_async().await;
368
369 let mock = server
370 .mock("POST", "/api/v1/bundles")
371 .match_body(Matcher::PartialJson(
372 json!({"method": "getBundleStatuses", "params": [[]]}),
373 ))
374 .with_status(200)
375 .with_header("content-type", "application/json")
376 .with_body(r#"{"jsonrpc":"2.0","id":1,"result":{"value":[]}}"#)
377 .create();
378
379 let config = JitoConfig { block_engine_url: server.url() };
380 let client = JitoClient::new(&config);
381
382 let result = client.get_bundle_statuses(vec![]).await;
383 mock.assert();
384 assert!(result.is_ok());
385 }
386
387 #[tokio::test]
388 async fn test_get_bundle_statuses_multiple_uuids() {
389 let mut server = Server::new_async().await;
390
391 let mock = server
392 .mock("POST", "/api/v1/bundles")
393 .match_body(Matcher::PartialJson(json!({
394 "method": "getBundleStatuses",
395 "params": [["uuid1", "uuid2", "uuid3"]]
396 })))
397 .with_status(200)
398 .with_header("content-type", "application/json")
399 .with_body(r#"{"jsonrpc":"2.0","id":1,"result":{"value":[{"bundle_id":"uuid1","status":"Landed"},{"bundle_id":"uuid2","status":"Pending"}]}}"#)
400 .create();
401
402 let config = JitoConfig { block_engine_url: server.url() };
403 let client = JitoClient::new(&config);
404
405 let result = client
406 .get_bundle_statuses(vec![
407 "uuid1".to_string(),
408 "uuid2".to_string(),
409 "uuid3".to_string(),
410 ])
411 .await;
412 mock.assert();
413 assert!(result.is_ok());
414 }
415
416 #[tokio::test]
417 async fn test_jito_bundle_client_dispatches_to_mock() {
418 let config = JitoConfig { block_engine_url: JITO_MOCK_BLOCK_ENGINE_URL.to_string() };
419 let client = JitoBundleClient::new(&config);
420
421 assert!(matches!(client, JitoBundleClient::Mock(_)));
422
423 let tx = create_mock_resolved_transaction();
424 let result = client.send_bundle(&[tx]).await;
425
426 assert!(result.is_ok());
427 let uuid = result.unwrap();
428 assert!(uuid.starts_with("mock-bundle-"), "Expected mock UUID prefix");
429 }
430
431 #[tokio::test]
432 async fn test_jito_bundle_client_dispatches_to_real() {
433 let config = JitoConfig { block_engine_url: "https://example.com".to_string() };
434 let client = JitoBundleClient::new(&config);
435
436 assert!(matches!(client, JitoBundleClient::Live(_)));
437 }
438
439 #[tokio::test]
440 async fn test_mock_client_get_bundle_statuses() {
441 let config = JitoConfig { block_engine_url: JITO_MOCK_BLOCK_ENGINE_URL.to_string() };
442 let client = JitoBundleClient::new(&config);
443
444 let result = client.get_bundle_statuses(vec!["test-uuid".to_string()]).await;
445
446 assert!(result.is_ok());
447 let statuses = result.unwrap();
448 assert!(statuses["value"].is_array());
449 assert_eq!(statuses["value"][0]["status"], "Landed");
450 }
451}