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            let _ =
131                $module.register_async_method($method_name, |rpc_params, rpc_context| async move {
132                    let rpc = rpc_context.as_ref();
133                    let params = rpc_params.parse()?;
134                    rpc.$rpc_method(params).await.map_err(Into::into)
135                });
136        }
137    };
138}
139
140fn build_rpc_module(rpc: KoraRpc) -> Result<RpcModule<KoraRpc>, anyhow::Error> {
141    let mut module = RpcModule::new(rpc.clone());
142    let enabled_methods = &get_config()?.kora.enabled_methods;
143
144    register_method_if_enabled!(module, enabled_methods, liveness, "liveness", liveness);
145
146    register_method_if_enabled!(
147        module,
148        enabled_methods,
149        estimate_transaction_fee,
150        "estimateTransactionFee",
151        estimate_transaction_fee,
152        with_params
153    );
154    register_method_if_enabled!(
155        module,
156        enabled_methods,
157        get_supported_tokens,
158        "getSupportedTokens",
159        get_supported_tokens
160    );
161    register_method_if_enabled!(
162        module,
163        enabled_methods,
164        get_payer_signer,
165        "getPayerSigner",
166        get_payer_signer
167    );
168    register_method_if_enabled!(
169        module,
170        enabled_methods,
171        sign_transaction,
172        "signTransaction",
173        sign_transaction,
174        with_params
175    );
176    register_method_if_enabled!(
177        module,
178        enabled_methods,
179        sign_and_send_transaction,
180        "signAndSendTransaction",
181        sign_and_send_transaction,
182        with_params
183    );
184    register_method_if_enabled!(
185        module,
186        enabled_methods,
187        transfer_transaction,
188        "transferTransaction",
189        transfer_transaction,
190        with_params
191    );
192    register_method_if_enabled!(
193        module,
194        enabled_methods,
195        get_blockhash,
196        "getBlockhash",
197        get_blockhash
198    );
199    register_method_if_enabled!(module, enabled_methods, get_config, "getConfig", get_config);
200
201    Ok(module)
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207    use crate::{
208        config::EnabledMethods,
209        tests::{
210            common::setup_or_get_test_signer,
211            config_mock::{ConfigMockBuilder, KoraConfigBuilder},
212            rpc_mock::RpcMockBuilder,
213        },
214    };
215    use std::env;
216
217    #[test]
218    fn test_get_value_by_priority_env_var_takes_precedence() {
219        let env_var_name = "TEST_ENV_VAR_PRECEDENCE_UNIQUE";
220        env::set_var(env_var_name, "env_value");
221
222        let result = get_value_by_priority(env_var_name, Some("config_value".to_string()));
223        assert_eq!(result, Some("env_value".to_string()));
224
225        env::remove_var(env_var_name);
226    }
227
228    #[test]
229    fn test_get_value_by_priority_config_fallback() {
230        let env_var_name = "TEST_ENV_VAR_FALLBACK_UNIQUE_XYZ123";
231
232        let result = get_value_by_priority(env_var_name, Some("config_value".to_string()));
233        assert_eq!(result, Some("config_value".to_string()));
234    }
235
236    #[test]
237    fn test_get_value_by_priority_none_when_both_missing() {
238        let env_var_name = "TEST_ENV_VAR_MISSING_UNIQUE_ABC789";
239
240        let result = get_value_by_priority(env_var_name, None);
241        assert_eq!(result, None);
242    }
243
244    #[test]
245    fn test_build_rpc_module_all_methods_enabled() {
246        // Default is all methods enabled
247        let enabled_methods = EnabledMethods::default();
248
249        let kora_config = KoraConfigBuilder::new().with_enabled_methods(enabled_methods).build();
250        let _m = ConfigMockBuilder::new().with_kora(kora_config).build_and_setup();
251        let _ = setup_or_get_test_signer();
252
253        let rpc_client = RpcMockBuilder::new().build();
254        let kora_rpc = KoraRpc::new(rpc_client);
255
256        let result = build_rpc_module(kora_rpc);
257        assert!(result.is_ok(), "Failed to build RPC module with all methods enabled");
258
259        // Verify that the module has the expected methods
260        let module = result.unwrap();
261        let method_names: Vec<&str> = module.method_names().collect();
262        assert_eq!(method_names.len(), 9);
263        assert!(method_names.contains(&"liveness"));
264        assert!(method_names.contains(&"estimateTransactionFee"));
265        assert!(method_names.contains(&"getSupportedTokens"));
266        assert!(method_names.contains(&"getPayerSigner"));
267        assert!(method_names.contains(&"signTransaction"));
268        assert!(method_names.contains(&"signAndSendTransaction"));
269        assert!(method_names.contains(&"transferTransaction"));
270        assert!(method_names.contains(&"getBlockhash"));
271        assert!(method_names.contains(&"getConfig"));
272    }
273
274    #[test]
275    fn test_build_rpc_module_all_methods_disabled() {
276        // Setup config with all methods disabled
277        let enabled_methods = EnabledMethods {
278            estimate_transaction_fee: false,
279            get_supported_tokens: false,
280            get_payer_signer: false,
281            sign_transaction: false,
282            sign_and_send_transaction: false,
283            transfer_transaction: false,
284            get_blockhash: false,
285            get_config: false,
286            liveness: false,
287        };
288
289        let kora_config = KoraConfigBuilder::new().with_enabled_methods(enabled_methods).build();
290        let _m = ConfigMockBuilder::new().with_kora(kora_config).build_and_setup();
291        let _ = setup_or_get_test_signer();
292
293        // Create RPC module
294        let rpc_client = RpcMockBuilder::new().build();
295        let kora_rpc = KoraRpc::new(rpc_client);
296
297        // Build the module - should succeed even with no methods
298        let result = build_rpc_module(kora_rpc);
299        assert!(result.is_ok(), "Failed to build RPC module with all methods disabled");
300
301        assert_eq!(result.unwrap().method_names().count(), 0);
302    }
303
304    #[test]
305    fn test_build_rpc_module_selective_methods() {
306        // Setup config with only some methods enabled
307        let enabled_methods = EnabledMethods {
308            liveness: true,
309            get_config: true,
310            get_supported_tokens: true,
311            estimate_transaction_fee: false,
312            get_payer_signer: false,
313            sign_transaction: false,
314            sign_and_send_transaction: false,
315            transfer_transaction: false,
316            get_blockhash: false,
317        };
318
319        let kora_config = KoraConfigBuilder::new().with_enabled_methods(enabled_methods).build();
320        let _m = ConfigMockBuilder::new().with_kora(kora_config).build_and_setup();
321        let _ = setup_or_get_test_signer();
322
323        // Create RPC module
324        let rpc_client = RpcMockBuilder::new().build();
325        let kora_rpc = KoraRpc::new(rpc_client);
326
327        // Build the module
328        let result = build_rpc_module(kora_rpc);
329        assert!(result.is_ok(), "Failed to build RPC module with selective methods");
330
331        // Verify that only the expected methods are registered
332        let module = result.unwrap();
333        let method_names: Vec<&str> = module.method_names().collect();
334        assert_eq!(method_names.len(), 3);
335        assert!(method_names.contains(&"liveness"));
336        assert!(method_names.contains(&"getConfig"));
337        assert!(method_names.contains(&"getSupportedTokens"));
338    }
339}
kora_lib/rpc_server/server.rs

kora_lib/rpc_server/
server.rs
