diff --git a/libs/wasm-loader/src/lib.rs b/libs/wasm-loader/src/lib.rs index 838efe7..d8f3b09 100644 --- a/libs/wasm-loader/src/lib.rs +++ b/libs/wasm-loader/src/lib.rs @@ -16,9 +16,9 @@ use log::*; use serde::Deserialize; use std::fmt::Debug; use std::io::Read; -use std::{fs::File, path::Path}; +use std::{fs::File, net::TcpStream, path::Path}; use subrpcer::state; -use tungstenite::Message; +use tungstenite::{stream::MaybeTlsStream, Message, WebSocket}; const CODE: &str = "0x3a636f6465"; // :code in hex pub const CODE_BLOB_BOMB_LIMIT: usize = 50 * 1024 * 1024; @@ -36,24 +36,43 @@ pub struct WasmLoader { compression: Compression, } -impl WasmLoader { - /// Fetch the wasm blob from a node - fn fetch_wasm_from_rpc(reference: &OnchainBlock) -> Result { - #[derive(Deserialize)] - struct Response { - result: String, - } +#[derive(Deserialize)] +struct Response { + result: String, +} - fn map_err(r: std::result::Result, e: E2) -> std::result::Result - where - E1: Debug, - { - r.map_err(|e_| { - eprintln!("{e_:?}"); - e - }) +fn map_err(r: std::result::Result, e: E2) -> std::result::Result +where + E1: Debug, +{ + r.map_err(|e_| { + eprintln!("{e_:?}"); + e + }) +} + +impl WasmLoader { + fn read_wasm_hex_response_from_web_socket( + mut ws: WebSocket>, + url: &String, + ) -> Result { + match map_err(ws.read(), WasmLoaderError::WsClient(url.to_string()))? { + Message::Text(t) => { + Ok(serde_json::from_str::(&t).map(|r| r.result).expect("unexpected response from node")) + } + Message::Ping(_) => { + log::debug!("Got ping from node, retrying."); + Self::read_wasm_hex_response_from_web_socket(ws, url) + } + m => { + log::warn!("Got unexpected message {:?} from node, retrying.", m); + Self::read_wasm_hex_response_from_web_socket(ws, url) + } } + } + /// Fetch the wasm blob from a node + fn fetch_wasm_from_rpc(reference: &OnchainBlock) -> Result { let block_ref = reference.block_ref.as_ref(); let data = state::get_storage(0, CODE, block_ref); let wasm_hex = match &reference.endpoint { @@ -70,19 +89,7 @@ impl WasmLoader { ws.send(Message::Binary(serde_json::to_vec(&data).expect("invalid data"))), WasmLoaderError::WsClient(url.to_string()), )?; - - let mut wasm_hex = None; - - // One for Ping, one for response. - for _ in 0..2_u8 { - let Message::Text(t) = map_err(ws.read(), WasmLoaderError::WsClient(url.to_string()))? else { - continue; - }; - - wasm_hex = serde_json::from_str::(&t).map(|r| r.result).ok(); - } - - wasm_hex.expect("unexpected response from node") + Self::read_wasm_hex_response_from_web_socket(ws, url)? } }; let wasm = array_bytes::hex2bytes(wasm_hex).expect("Decoding bytes");