iota_rest_kv/
server.rs

1// Copyright (c) 2025 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4//! This module includes helper wrappers for building and starting a REST API
5//! server.
6
7use std::{net::SocketAddr, sync::Arc};
8
9use anyhow::Result;
10use axum::{Router, response::IntoResponse, routing::get};
11use tokio_util::sync::CancellationToken;
12
13use crate::{
14    RestApiConfig,
15    errors::ApiError,
16    kv_store_client::KvStoreClient,
17    routes::{health, kv_store},
18};
19
20/// A wrapper which builds the components needed for the REST API server and
21/// provides a simple way to start it.
22pub struct Server {
23    router: Router,
24    server_address: SocketAddr,
25    token: CancellationToken,
26}
27
28impl Server {
29    /// Create a new Server instance.
30    ///
31    /// Based on the config, it instantiates the [`KvStoreClient`] and
32    /// constructs the [`Router`].
33    pub async fn new(config: RestApiConfig, token: CancellationToken) -> Result<Self> {
34        let kv_store_client = KvStoreClient::new(config.kv_store_config).await?;
35
36        let shared_state = Arc::new(kv_store_client);
37
38        let router = Router::new()
39            .route("/health", get(health::health))
40            .route("/:item_type/:key", get(kv_store::data_as_bytes))
41            .with_state(shared_state)
42            .fallback(fallback);
43
44        Ok(Self {
45            router,
46            token,
47            server_address: config.server_address,
48        })
49    }
50
51    /// Start the server, this method is blocking.
52    pub async fn serve(self) -> Result<()> {
53        let listener = tokio::net::TcpListener::bind(self.server_address)
54            .await
55            .expect("failed to bind to socket");
56
57        tracing::info!("listening on: {}", self.server_address);
58
59        axum::serve(listener, self.router)
60            .with_graceful_shutdown(async move {
61                self.token.cancelled().await;
62                tracing::info!("shutdown signal received.");
63            })
64            .await
65            .inspect_err(|e| tracing::error!("server encountered an error: {e}"))
66            .map_err(Into::into)
67    }
68}
69
70/// Handles requests to routes that are not defined in the API.
71///
72/// This fallback handler is called when the requested URL path does not match
73/// any of the defined routes. It returns a `404 Not Found` error, indicating
74/// that the requested resource could not be found. This can happen if the user
75/// enters an incorrect URL or if the requested resource (identified by a
76/// [`Key`](iota_storage::http_key_value_store::Key)) cannot be extracted from
77/// the request.
78async fn fallback() -> impl IntoResponse {
79    ApiError::NotFound
80}
iota_rest_kv/server.rs

iota_rest_kv/
server.rs
