kora_lib/rpc_server/
server.rs

1use crate::{
2    constant::{X_API_KEY, X_HMAC_SIGNATURE, X_TIMESTAMP},
3    metrics::run_metrics_server_if_required,
4    rpc_server::{
5        auth::{ApiKeyAuthLayer, HmacAuthLayer},
6        middleware_utils::MethodValidationLayer,
7        rpc::KoraRpc,
8    },
9    usage_limit::UsageTracker,
10};
11
12#[cfg(not(test))]
13use crate::state::get_config;
14
15#[cfg(test)]
16use crate::tests::config_mock::mock_state::get_config;
17use http::{header, Method};
18use jsonrpsee::{
19    server::{middleware::proxy_get_request::ProxyGetRequestLayer, ServerBuilder, ServerHandle},
20    RpcModule,
21};
22use std::{net::SocketAddr, time::Duration};
23use tokio::task::JoinHandle;
24use tower::limit::RateLimitLayer;
25use tower_http::cors::CorsLayer;
26
27pub struct ServerHandles {
28    pub rpc_handle: ServerHandle,
29    pub metrics_handle: Option<ServerHandle>,
30    pub balance_tracker_handle: Option<JoinHandle<()>>,
31}
32
33// We'll always prioritize the environment variable over the config value
34fn get_value_by_priority(env_var: &str, config_value: Option<String>) -> Option<String> {
35    std::env::var(env_var).ok().or(config_value)
36}
37
38pub async fn run_rpc_server(rpc: KoraRpc, port: u16) -> Result<ServerHandles, anyhow::Error> {
39    let addr = SocketAddr::from(([0, 0, 0, 0], port));
40    log::info!("RPC server started on {addr}, port {port}");
41
42    // Initialize usage limiter
43    if let Err(e) = UsageTracker::init_usage_limiter().await {
44        log::error!("Failed to initialize usage limiter: {e}");
45        return Err(anyhow::anyhow!("Usage limiter initialization failed: {e}"));
46    }
47
48    // Build middleware stack with tracing and CORS
49    let cors = CorsLayer::new()
50        .allow_origin(tower_http::cors::Any)
51        .allow_methods([Method::POST, Method::GET])
52        .allow_headers([
53            header::CONTENT_TYPE,
54            header::HeaderName::from_static(X_API_KEY),
55            header::HeaderName::from_static(X_HMAC_SIGNATURE),
56            header::HeaderName::from_static(X_TIMESTAMP),
57        ])
58        .max_age(Duration::from_secs(3600));
59
60    let config = get_config()?;
61
62    // Get the RPC client from KoraRpc to pass to metrics initialization
63    let rpc_client = rpc.get_rpc_client().clone();
64
65    let (metrics_handle, metrics_layers, balance_tracker_handle) =
66        run_metrics_server_if_required(port, rpc_client).await?;
67
68    // Build whitelist of allowed methods from enabled_methods config
69    let allowed_methods = config.kora.enabled_methods.get_enabled_method_names();
70
71    let middleware = tower::ServiceBuilder::new()
72        // Add metrics handler first (before other layers) so it can intercept /metrics
73        .layer(ProxyGetRequestLayer::new("/liveness", "liveness")?)
74        .layer(RateLimitLayer::new(config.kora.rate_limit, Duration::from_secs(1)))
75        // Add metrics handler layer for Prometheus metrics
76        .option_layer(
77            metrics_layers.as_ref().and_then(|layers| layers.metrics_handler_layer.clone()),
78        )
79        .layer(cors)
80        // Method validation layer -  to fail fast
81        .layer(MethodValidationLayer::new(allowed_methods.clone()))
82        // Add metrics collection layer
83        .option_layer(metrics_layers.as_ref().and_then(|layers| layers.http_metrics_layer.clone()))
84        // Add authentication layer for API key if configured
85        .option_layer(
86            (get_value_by_priority("KORA_API_KEY", config.kora.auth.api_key.clone()))
87                .map(ApiKeyAuthLayer::new),
88        )
89        // Add authentication layer for HMAC if configured
90        .option_layer(
91            (get_value_by_priority("KORA_HMAC_SECRET", config.kora.auth.hmac_secret.clone()))
92                .map(|secret| HmacAuthLayer::new(secret, config.kora.auth.max_timestamp_age)),
93        );
94
95    // Configure and build the server with HTTP support
96    let server = ServerBuilder::default()
97        .max_request_body_size(config.kora.max_request_body_size as u32)
98        .set_middleware(middleware)
99        .http_only() // Explicitly enable HTTP
100        .build(addr)
101        .await?;
102
103    let rpc_module = build_rpc_module(rpc)?;
104
105    // Start the RPC server
106    let rpc_handle = server
107        .start(rpc_module)
108        .map_err(|e| anyhow::anyhow!("Failed to start RPC server: {}", e))?;
109
110    Ok(ServerHandles { rpc_handle, metrics_handle, balance_tracker_handle })
111}
112
113macro_rules! register_method_if_enabled {
114    // For methods without parameters
115    ($module:expr, $enabled_methods:expr, $field:ident, $method_name:expr, $rpc_method:ident) => {
116        if $enabled_methods.$field {
117            let _ = $module.register_async_method(
118                $method_name,
119                |_rpc_params, rpc_context| async move {
120                    let rpc = rpc_context.as_ref();
121                    rpc.$rpc_method().await.map_err(Into::into)
122                },
123            );
124        }
125    };
126
127    // For methods with parameters
128    ($module:expr, $enabled_methods:expr, $field:ident, $method_name:expr, $rpc_method:ident, with_params) => {
129        if $enabled_methods.$field {
130            #[allow(deprecated)]
131            let _ =
132                $module.register_async_method($method_name, |rpc_params, rpc_context| async move {
133                    let rpc = rpc_context.as_ref();
134                    let params = rpc_params.parse()?;
135                    #[allow(deprecated)]
136                    rpc.$rpc_method(params).await.map_err(Into::into)
137                });
138        }
139    };
140}
141
142fn build_rpc_module(rpc: KoraRpc) -> Result<RpcModule<KoraRpc>, anyhow::Error> {
143    let mut module = RpcModule::new(rpc.clone());
144    let enabled_methods = &get_config()?.kora.enabled_methods;
145
146    register_method_if_enabled!(module, enabled_methods, liveness, "liveness", liveness);
147
148    register_method_if_enabled!(
149        module,
150        enabled_methods,
151        estimate_transaction_fee,
152        "estimateTransactionFee",
153        estimate_transaction_fee,
154        with_params
155    );
156    register_method_if_enabled!(
157        module,
158        enabled_methods,
159        get_supported_tokens,
160        "getSupportedTokens",
161        get_supported_tokens
162    );
163    register_method_if_enabled!(
164        module,
165        enabled_methods,
166        get_payer_signer,
167        "getPayerSigner",
168        get_payer_signer
169    );
170    register_method_if_enabled!(
171        module,
172        enabled_methods,
173        sign_transaction,
174        "signTransaction",
175        sign_transaction,
176        with_params
177    );
178    register_method_if_enabled!(
179        module,
180        enabled_methods,
181        sign_and_send_transaction,
182        "signAndSendTransaction",
183        sign_and_send_transaction,
184        with_params
185    );
186    register_method_if_enabled!(
187        module,
188        enabled_methods,
189        transfer_transaction,
190        "transferTransaction",
191        transfer_transaction,
192        with_params
193    );
194    register_method_if_enabled!(
195        module,
196        enabled_methods,
197        get_blockhash,
198        "getBlockhash",
199        get_blockhash
200    );
201    register_method_if_enabled!(module, enabled_methods, get_config, "getConfig", get_config);
202    register_method_if_enabled!(module, enabled_methods, get_version, "getVersion", get_version);
203    register_method_if_enabled!(
204        module,
205        enabled_methods,
206        sign_bundle,
207        "signBundle",
208        sign_bundle,
209        with_params
210    );
211    register_method_if_enabled!(
212        module,
213        enabled_methods,
214        sign_and_send_bundle,
215        "signAndSendBundle",
216        sign_and_send_bundle,
217        with_params
218    );
219
220    Ok(module)
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226    use crate::{
227        config::EnabledMethods,
228        tests::{
229            common::setup_or_get_test_signer,
230            config_mock::{ConfigMockBuilder, KoraConfigBuilder},
231            rpc_mock::RpcMockBuilder,
232        },
233    };
234    use std::env;
235
236    #[test]
237    fn test_get_value_by_priority_env_var_takes_precedence() {
238        let env_var_name = "TEST_ENV_VAR_PRECEDENCE_UNIQUE";
239        env::set_var(env_var_name, "env_value");
240
241        let result = get_value_by_priority(env_var_name, Some("config_value".to_string()));
242        assert_eq!(result, Some("env_value".to_string()));
243
244        env::remove_var(env_var_name);
245    }
246
247    #[test]
248    fn test_get_value_by_priority_config_fallback() {
249        let env_var_name = "TEST_ENV_VAR_FALLBACK_UNIQUE_XYZ123";
250
251        let result = get_value_by_priority(env_var_name, Some("config_value".to_string()));
252        assert_eq!(result, Some("config_value".to_string()));
253    }
254
255    #[test]
256    fn test_get_value_by_priority_none_when_both_missing() {
257        let env_var_name = "TEST_ENV_VAR_MISSING_UNIQUE_ABC789";
258
259        let result = get_value_by_priority(env_var_name, None);
260        assert_eq!(result, None);
261    }
262
263    #[test]
264    fn test_build_rpc_module_all_methods_enabled() {
265        // Default is all methods enabled
266        let enabled_methods = EnabledMethods::default();
267
268        let kora_config = KoraConfigBuilder::new().with_enabled_methods(enabled_methods).build();
269        let _m = ConfigMockBuilder::new().with_kora(kora_config).build_and_setup();
270        let _ = setup_or_get_test_signer();
271
272        let rpc_client = RpcMockBuilder::new().build();
273        let kora_rpc = KoraRpc::new(rpc_client);
274
275        let result = build_rpc_module(kora_rpc);
276        assert!(result.is_ok(), "Failed to build RPC module with all methods enabled");
277
278        // Verify that the module has the expected methods
279        let module = result.unwrap();
280        let method_names: Vec<&str> = module.method_names().collect();
281        assert_eq!(method_names.len(), 10);
282        assert!(method_names.contains(&"liveness"));
283        assert!(method_names.contains(&"estimateTransactionFee"));
284        assert!(method_names.contains(&"getSupportedTokens"));
285        assert!(method_names.contains(&"getPayerSigner"));
286        assert!(method_names.contains(&"signTransaction"));
287        assert!(method_names.contains(&"signAndSendTransaction"));
288        assert!(method_names.contains(&"transferTransaction"));
289        assert!(method_names.contains(&"getBlockhash"));
290        assert!(method_names.contains(&"getConfig"));
291        assert!(method_names.contains(&"getVersion"));
292        // Note: signBundle is NOT included by default (opt-in via enabled_methods.sign_bundle)
293    }
294
295    #[test]
296    fn test_build_rpc_module_all_methods_disabled() {
297        // Setup config with all methods disabled
298        let enabled_methods = EnabledMethods {
299            estimate_transaction_fee: false,
300            get_supported_tokens: false,
301            get_payer_signer: false,
302            sign_transaction: false,
303            sign_and_send_transaction: false,
304            transfer_transaction: false,
305            get_blockhash: false,
306            get_config: false,
307            get_version: false,
308            liveness: false,
309            sign_and_send_bundle: false,
310            sign_bundle: false,
311        };
312
313        let kora_config = KoraConfigBuilder::new().with_enabled_methods(enabled_methods).build();
314        let _m = ConfigMockBuilder::new().with_kora(kora_config).build_and_setup();
315        let _ = setup_or_get_test_signer();
316
317        // Create RPC module
318        let rpc_client = RpcMockBuilder::new().build();
319        let kora_rpc = KoraRpc::new(rpc_client);
320
321        // Build the module - should succeed even with no methods
322        let result = build_rpc_module(kora_rpc);
323        assert!(result.is_ok(), "Failed to build RPC module with all methods disabled");
324
325        assert_eq!(result.unwrap().method_names().count(), 0);
326    }
327
328    #[test]
329    fn test_build_rpc_module_selective_methods() {
330        // Setup config with only some methods enabled
331        let enabled_methods = EnabledMethods {
332            liveness: true,
333            get_config: true,
334            get_supported_tokens: true,
335            estimate_transaction_fee: false,
336            get_payer_signer: false,
337            sign_transaction: false,
338            sign_and_send_transaction: false,
339            transfer_transaction: false,
340            get_blockhash: false,
341            get_version: false,
342            sign_and_send_bundle: false,
343            sign_bundle: false,
344        };
345
346        let kora_config = KoraConfigBuilder::new().with_enabled_methods(enabled_methods).build();
347        let _m = ConfigMockBuilder::new().with_kora(kora_config).build_and_setup();
348        let _ = setup_or_get_test_signer();
349
350        // Create RPC module
351        let rpc_client = RpcMockBuilder::new().build();
352        let kora_rpc = KoraRpc::new(rpc_client);
353
354        // Build the module
355        let result = build_rpc_module(kora_rpc);
356        assert!(result.is_ok(), "Failed to build RPC module with selective methods");
357
358        // Verify that only the expected methods are registered
359        let module = result.unwrap();
360        let method_names: Vec<&str> = module.method_names().collect();
361        assert_eq!(method_names.len(), 3);
362        assert!(method_names.contains(&"liveness"));
363        assert!(method_names.contains(&"getConfig"));
364        assert!(method_names.contains(&"getSupportedTokens"));
365    }
366}
kora_lib/rpc_server/server.rs

kora_lib/rpc_server/
server.rs
