Skip to main content

miden_note_transport_node/node/grpc/
mod.rs

1mod streaming;
2
3use std::collections::BTreeSet;
4use std::net::SocketAddr;
5use std::sync::Arc;
6use std::time::Duration;
7
8use chrono::Utc;
9use miden_note_transport_proto::miden_note_transport::miden_note_transport_server::MidenNoteTransportServer;
10use miden_note_transport_proto::miden_note_transport::{
11    FetchNotesRequest,
12    FetchNotesResponse,
13    SendNoteRequest,
14    SendNoteResponse,
15    StatsResponse,
16    StreamNotesRequest,
17    TransportNote,
18};
19use miden_protocol::utils::serde::Deserializable;
20use rand::Rng;
21use tokio::sync::mpsc;
22use tonic::Status;
23use tonic_web::GrpcWebLayer;
24use tower::limit::GlobalConcurrencyLimitLayer;
25use tower::timeout::TimeoutLayer;
26use tower_http::cors::{Any, CorsLayer};
27
28use self::streaming::{NoteStreamer, StreamerMessage, Sub, Subface};
29use crate::database::Database;
30use crate::metrics::MetricsGrpc;
31
32/// Upper bound on the number of tags a client may include in a single
33/// `fetch_notes` request. Guards against two concerns:
34///   - Server CPU: deduplicating `request_data.tags` via `BTreeSet` is `O(n log n)`; a client
35///     sending millions of tags can burn a worker.
36///   - SQLite `IN (...)`: the underlying driver caps bound variables at
37///     `SQLITE_MAX_VARIABLE_NUMBER` (32766 on recent builds, lower on older); blow that and the
38///     query errors. Well below the SQLite cap so we have headroom for future query-plan changes.
39///
40/// A realistic wallet tracks O(10) to O(100) tags; 128 is generous without
41/// being an attack surface.
42const MAX_TAGS_PER_FETCH_REQUEST: usize = 128;
43
44/// Miden Note Transport gRPC server
45pub struct GrpcServer {
46    database: Arc<Database>,
47    config: GrpcServerConfig,
48    streamer: StreamerCtx,
49    metrics: MetricsGrpc,
50}
51
52/// [`GrpcServer`] configuration
53#[derive(Clone, Debug)]
54pub struct GrpcServerConfig {
55    /// Server host
56    pub host: String,
57    /// Server port
58    pub port: u16,
59    /// Maximum note size to be stored
60    pub max_note_size: usize,
61    /// Maximum number of concurrent connections
62    pub max_connections: usize,
63    /// Connection timeout in seconds
64    pub request_timeout: usize,
65}
66
67/// Streaming task interface context
68pub(super) struct StreamerCtx {
69    tx: mpsc::Sender<StreamerMessage>,
70    handle: tokio::task::JoinHandle<()>,
71}
72
73impl Default for GrpcServerConfig {
74    fn default() -> Self {
75        Self {
76            host: "127.0.0.1".to_string(),
77            port: 57292,
78            max_note_size: 512_000,
79            max_connections: 4096,
80            request_timeout: 4,
81        }
82    }
83}
84
85impl GrpcServer {
86    /// gRPC server constructor
87    pub fn new(database: Arc<Database>, config: GrpcServerConfig, metrics: MetricsGrpc) -> Self {
88        let streamer = StreamerCtx::spawn(database.clone());
89        Self { database, config, streamer, metrics }
90    }
91
92    /// Convert into a service
93    pub fn into_service(self) -> MidenNoteTransportServer<Self> {
94        MidenNoteTransportServer::new(self)
95    }
96
97    /// gRPC server running-task
98    pub async fn serve(self) -> crate::Result<()> {
99        let (health_reporter, health_svc) = tonic_health::server::health_reporter();
100        health_reporter.set_serving::<MidenNoteTransportServer<Self>>().await;
101
102        let addr = format!("{}:{}", self.config.host, self.config.port)
103            .parse::<SocketAddr>()
104            .map_err(|e| crate::Error::Internal(format!("Invalid address: {e}")))?;
105
106        let cors = CorsLayer::new().allow_origin(Any).allow_headers(Any).allow_methods(Any);
107
108        tonic::transport::Server::builder()
109            .accept_http1(true)
110            .layer(cors)
111            .layer(GrpcWebLayer::new())
112            .layer(GlobalConcurrencyLimitLayer::new(self.config.max_connections))
113            .layer(TimeoutLayer::new(Duration::from_secs(self.config.request_timeout as u64)))
114            .add_service(health_svc)
115            .add_service(self.into_service())
116            .serve(addr)
117            .await
118            .map_err(|e| crate::Error::Internal(format!("Server error: {e}")))
119    }
120}
121
122impl StreamerCtx {
123    /// Spawn a [`NoteStreamer`] task
124    ///
125    /// Returns related context composed of the handle and `mpsc::Sender` `tx` for control messages.
126    pub(super) fn spawn(database: Arc<Database>) -> Self {
127        let (tx, rx) = mpsc::channel(128);
128        let handle = tokio::spawn(NoteStreamer::new(database, rx).stream());
129        Self { tx, handle }
130    }
131}
132
133#[tonic::async_trait]
134impl miden_note_transport_proto::miden_note_transport::miden_note_transport_server::MidenNoteTransport
135    for GrpcServer
136{
137    #[tracing::instrument(skip(self), fields(operation = "grpc.send_note.request"))]
138    async fn send_note(
139        &self,
140        request: tonic::Request<SendNoteRequest>,
141    ) -> Result<tonic::Response<SendNoteResponse>, tonic::Status> {
142        let request_data = request.into_inner();
143        let pnote = request_data.note.ok_or_else(|| Status::invalid_argument("Missing note"))?;
144
145        let timer = self.metrics.grpc_send_note_request((pnote.header.len() + pnote.details.len()) as u64);
146
147        // Validate note size
148        if pnote.details.len() > self.config.max_note_size {
149            return Err(Status::resource_exhausted(format!("Note too large ({})", pnote.details.len())));
150        }
151
152        // Convert protobuf request to internal types
153        let header = miden_protocol::note::NoteHeader::read_from_bytes(&pnote.header)
154            .map_err(|e| Status::invalid_argument(format!("Invalid header: {e:?}")))?;
155
156        // Create note for database
157        let note_for_db = crate::types::StoredNote {
158            header,
159            details: pnote.details,
160            created_at: Utc::now(),
161            // Ignored on INSERT: the DB assigns seq via AUTOINCREMENT.
162            seq: 0,
163        };
164
165        self.database
166            .store_note(&note_for_db)
167            .await.map_err(|e| tonic::Status::internal(format!("Failed to store note: {e:?}")))?;
168
169        timer.finish("ok");
170
171        Ok(tonic::Response::new(SendNoteResponse {}))
172    }
173
174    #[tracing::instrument(skip(self), fields(operation = "grpc.fetch_notes.request"))]
175    async fn fetch_notes(
176        &self,
177        request: tonic::Request<FetchNotesRequest>,
178    ) -> Result<tonic::Response<FetchNotesResponse>, tonic::Status> {
179        let timer = self.metrics.grpc_fetch_notes_request();
180
181        let request_data = request.into_inner();
182
183        // Reject requests with too many tags BEFORE any allocation /
184        // deduplication work. A client sending `[0u32; 1_000_000]` would
185        // otherwise force an O(n log n) BTreeSet build and then either blow
186        // through `SQLITE_MAX_VARIABLE_NUMBER` or return a pathologically
187        // expensive query plan.
188        if request_data.tags.len() > MAX_TAGS_PER_FETCH_REQUEST {
189            return Err(Status::invalid_argument(format!(
190                "Too many tags in fetch_notes request: {} (max {})",
191                request_data.tags.len(),
192                MAX_TAGS_PER_FETCH_REQUEST
193            )));
194        }
195
196        // Deduplicate incoming tags — the DB query is more efficient without repeats
197        // and the previous per-tag loop happened to dedupe via BTreeSet.
198        let tag_set: BTreeSet<_> = request_data.tags.into_iter().collect();
199        let tags: Vec<crate::types::NoteTag> = tag_set.into_iter().map(Into::into).collect();
200        let cursor = request_data.cursor;
201
202        // Single-snapshot fetch across ALL tags. Running per-tag queries back
203        // to back exposed a race where a concurrent INSERT could land between
204        // two per-tag queries and get leapfrogged when rcursor advanced past
205        // its seq on the next fetch. A single `tag IN (…)` query reads all
206        // matching rows in one consistent snapshot.
207        let stored_notes = self
208            .database
209            .fetch_notes_by_tags(&tags, cursor)
210            .await
211            .map_err(|e| tonic::Status::internal(format!("Failed to fetch notes: {e:?}")))?;
212
213        let mut rcursor = cursor;
214        for stored_note in &stored_notes {
215            let seq_cursor: u64 = stored_note
216                .seq
217                .try_into()
218                .map_err(|_| tonic::Status::internal("Negative seq in stored note"))?;
219            rcursor = rcursor.max(seq_cursor);
220        }
221
222        let proto_notes: Vec<_> = stored_notes.into_iter().map(TransportNote::from).collect();
223
224        timer.finish("ok");
225
226        let proto_notes_size = proto_notes.iter().map(|pnote| (pnote.header.len() + pnote.details.len()) as u64).sum();
227        self.metrics.grpc_fetch_notes_response(
228            proto_notes.len() as u64,
229            proto_notes_size,
230        );
231
232        Ok(tonic::Response::new(FetchNotesResponse { notes: proto_notes, cursor: rcursor }))
233    }
234
235    type StreamNotesStream = Sub;
236    #[tracing::instrument(skip(self), fields(operation = "grpc.stream_notes.request"))]
237    async fn stream_notes(
238        &self,
239        request: tonic::Request<StreamNotesRequest>,
240    ) -> Result<tonic::Response<Self::StreamNotesStream>, tonic::Status> {
241        let request_data = request.into_inner();
242        let tag = request_data.tag.into();
243        let id = rand::rng().random();
244        let (sub_tx, sub_rx) = mpsc::channel(32);
245        let sub = Sub::new(id, tag, sub_rx, self.streamer.tx.clone());
246        let subf = Subface::new(id, tag, sub_tx);
247        self.streamer.tx.try_send(StreamerMessage::AddSub(subf))
248                    .map_err(|e| tonic::Status::internal(format!("Failed sending internal streamer message: {e}")))?;
249
250        Ok(tonic::Response::new(sub))
251    }
252
253    #[tracing::instrument(skip(self), fields(operation = "grpc.stats.request"))]
254    async fn stats(
255        &self,
256        _request: tonic::Request<()>,
257    ) -> Result<tonic::Response<StatsResponse>, tonic::Status> {
258        let (total_notes, total_tags) = self
259            .database
260            .get_stats()
261            .await.map_err(|e| tonic::Status::internal(format!("Failed to get stats: {e:?}")))?;
262
263        let response = StatsResponse {
264            total_notes,
265            total_tags,
266            notes_per_tag: Vec::new(), // TODO: Implement notes_per_tag
267        };
268
269        Ok(tonic::Response::new(response))
270    }
271}
272
273impl Drop for StreamerCtx {
274    fn drop(&mut self) {
275        if let Err(e) = self.tx.try_send(StreamerMessage::Shutdown) {
276            tracing::error!("Streamer shutdown message sending failure: {e}");
277            self.handle.abort();
278        }
279    }
280}
281
282#[cfg(test)]
283mod tests {
284    use std::sync::Arc;
285
286    use miden_note_transport_proto::miden_note_transport::FetchNotesRequest;
287    use miden_note_transport_proto::miden_note_transport::miden_note_transport_server::MidenNoteTransport;
288
289    use super::*;
290    use crate::database::{Database, DatabaseConfig};
291    use crate::metrics::Metrics;
292
293    async fn test_server() -> GrpcServer {
294        let metrics = Metrics::default();
295        let db = Arc::new(
296            Database::connect(DatabaseConfig::default(), metrics.db.clone()).await.unwrap(),
297        );
298        GrpcServer::new(db, GrpcServerConfig::default(), metrics.grpc)
299    }
300
301    /// A client sending more tags than `MAX_TAGS_PER_FETCH_REQUEST` is rejected
302    /// with `InvalidArgument` BEFORE any `BTreeSet` or DB work. Guards against
303    /// both the O(n log n) dedup cost and the `SQLITE_MAX_VARIABLE_NUMBER`
304    /// ceiling.
305    #[tokio::test]
306    async fn test_fetch_notes_rejects_too_many_tags() {
307        let server = test_server().await;
308
309        let tags = vec![0u32; MAX_TAGS_PER_FETCH_REQUEST + 1];
310        let request = tonic::Request::new(FetchNotesRequest { tags, cursor: 0 });
311        let result = server.fetch_notes(request).await;
312
313        let status = result.expect_err("expected InvalidArgument");
314        assert_eq!(status.code(), tonic::Code::InvalidArgument);
315        assert!(
316            status.message().contains("Too many tags"),
317            "unexpected error message: {}",
318            status.message()
319        );
320    }
321
322    /// A client sending exactly `MAX_TAGS_PER_FETCH_REQUEST` tags is accepted.
323    /// (Using the same tag value many times is fine — the handler dedups via
324    /// `BTreeSet` before issuing the query.)
325    #[tokio::test]
326    async fn test_fetch_notes_accepts_max_tags_at_limit() {
327        let server = test_server().await;
328
329        let tags = vec![0u32; MAX_TAGS_PER_FETCH_REQUEST];
330        let request = tonic::Request::new(FetchNotesRequest { tags, cursor: 0 });
331        let result = server.fetch_notes(request).await;
332
333        let response = result.expect("request at the cap must succeed").into_inner();
334        assert_eq!(response.notes.len(), 0, "DB is empty, no notes returned");
335        assert_eq!(response.cursor, 0);
336    }
337}