mockforge_recorder/protocols/
grpc.rs

1//! gRPC recording helpers
2
3use crate::{models::*, recorder::Recorder};
4use chrono::Utc;
5use std::collections::HashMap;
6use tracing::debug;
7use uuid::Uuid;
8
9/// Record a gRPC request
10pub async fn record_grpc_request(
11    recorder: &Recorder,
12    service: &str,
13    method: &str,
14    metadata: &HashMap<String, String>,
15    message: Option<&[u8]>,
16    context: &crate::models::RequestContext,
17) -> Result<String, crate::RecorderError> {
18    let request_id = Uuid::new_v4().to_string();
19    let full_method = format!("{}/{}", service, method);
20
21    let (body_str, body_encoding) = if let Some(msg) = message {
22        // gRPC messages are protobuf, so always base64 encode
23        let encoded = base64::Engine::encode(&base64::engine::general_purpose::STANDARD, msg);
24        (Some(encoded), "base64".to_string())
25    } else {
26        (None, "utf8".to_string())
27    };
28
29    let request = RecordedRequest {
30        id: request_id.clone(),
31        protocol: Protocol::Grpc,
32        timestamp: Utc::now(),
33        method: full_method.clone(),
34        path: format!("/{}", full_method),
35        query_params: None,
36        headers: serde_json::to_string(&metadata)?,
37        body: body_str,
38        body_encoding,
39        client_ip: context.client_ip.clone(),
40        trace_id: context.trace_id.clone(),
41        span_id: context.span_id.clone(),
42        duration_ms: None,
43        status_code: None,
44        tags: None,
45    };
46
47    recorder.record_request(request).await?;
48    debug!("Recorded gRPC request: {} {}", request_id, full_method);
49
50    Ok(request_id)
51}
52
53/// Record a gRPC response
54pub async fn record_grpc_response(
55    recorder: &Recorder,
56    request_id: &str,
57    status_code: i32, // gRPC status code (0 = OK)
58    metadata: &HashMap<String, String>,
59    message: Option<&[u8]>,
60    duration_ms: i64,
61) -> Result<(), crate::RecorderError> {
62    let (body_str, body_encoding) = if let Some(msg) = message {
63        let encoded = base64::Engine::encode(&base64::engine::general_purpose::STANDARD, msg);
64        (Some(encoded), "base64".to_string())
65    } else {
66        (None, "utf8".to_string())
67    };
68
69    let size_bytes = message.map(|m| m.len()).unwrap_or(0) as i64;
70
71    let response = RecordedResponse {
72        request_id: request_id.to_string(),
73        status_code,
74        headers: serde_json::to_string(&metadata)?,
75        body: body_str,
76        body_encoding,
77        size_bytes,
78        timestamp: Utc::now(),
79    };
80
81    recorder.record_response(response).await?;
82    debug!(
83        "Recorded gRPC response: {} status={} duration={}ms",
84        request_id, status_code, duration_ms
85    );
86
87    Ok(())
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93    use crate::database::RecorderDatabase;
94
95    #[tokio::test]
96    async fn test_record_grpc_exchange() {
97        let db = RecorderDatabase::new_in_memory().await.unwrap();
98        let recorder = Recorder::new(db);
99
100        let metadata =
101            HashMap::from([("content-type".to_string(), "application/grpc".to_string())]);
102
103        let context = crate::models::RequestContext::new(Some("127.0.0.1"), None, None);
104        let request_id = record_grpc_request(
105            &recorder,
106            "helloworld.Greeter",
107            "SayHello",
108            &metadata,
109            Some(b"\x00\x00\x00\x00\x05hello"),
110            &context,
111        )
112        .await
113        .unwrap();
114
115        record_grpc_response(
116            &recorder,
117            &request_id,
118            0,
119            &metadata,
120            Some(b"\x00\x00\x00\x00\x05world"),
121            42,
122        )
123        .await
124        .unwrap();
125
126        // Verify it was recorded
127        let exchange = recorder.database().get_exchange(&request_id).await.unwrap();
128        assert!(exchange.is_some());
129
130        let exchange = exchange.unwrap();
131        assert_eq!(exchange.request.protocol, Protocol::Grpc);
132        assert_eq!(exchange.request.method, "helloworld.Greeter/SayHello");
133    }
134}