Skip to main content

kora_lib/rpc_server/
server.rs

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