Skip to main content

geode_client/
proto.rs

1//! Protobuf types generated from geode.proto via prost-build/tonic-build.
2//!
3//! This module re-exports the generated types and provides helper functions
4//! for QUIC transport framing (4-byte big-endian length prefix).
5
6// Include the prost-generated code.
7#[allow(clippy::large_enum_variant)]
8#[path = "generated/geode.rs"]
9mod generated;
10
11// Re-export all generated types at this module level.
12pub use generated::*;
13
14use prost::Message;
15
16// Note: We use fully-qualified `crate::error::Error` to avoid collision with
17// the generated `proto::Error` message type (from `message Error` in geode.proto).
18use crate::error::Result;
19
20// =============================================================================
21// QUIC framing helpers (length-prefixed protobuf)
22// =============================================================================
23
24/// Encode a QuicClientMessage to protobuf bytes with a 4-byte big-endian
25/// length prefix, as required by the QUIC transport.
26pub fn encode_with_length_prefix(msg: &QuicClientMessage) -> Vec<u8> {
27    let data = msg.encode_to_vec();
28    let length = data.len() as u32;
29    let mut result = Vec::with_capacity(4 + data.len());
30    result.extend(&length.to_be_bytes());
31    result.extend(data);
32    result
33}
34
35/// Decode a 4-byte big-endian length prefix from the start of a byte slice.
36pub fn decode_length_prefix(data: &[u8]) -> Result<u32> {
37    if data.len() < 4 {
38        return Err(crate::error::Error::protocol(
39            "Insufficient data for length prefix",
40        ));
41    }
42    Ok(u32::from_be_bytes([data[0], data[1], data[2], data[3]]))
43}
44
45/// Decode a QuicServerMessage from raw protobuf bytes (without length prefix).
46pub fn decode_quic_server_message(data: &[u8]) -> Result<QuicServerMessage> {
47    QuicServerMessage::decode(data)
48        .map_err(|e| crate::error::Error::protocol(format!("Protobuf decode error: {}", e)))
49}
50
51// =============================================================================
52// Tests
53// =============================================================================
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58
59    #[test]
60    fn test_encode_decode_hello_roundtrip() {
61        let req = HelloRequest {
62            username: "admin".to_string(),
63            password: "secret".to_string(),
64            tenant_id: Some("tenant1".to_string()),
65            client_name: String::new(),
66            client_version: String::new(),
67            wanted_conformance: String::new(),
68            graph: None,
69            role: None,
70        };
71        let msg = QuicClientMessage {
72            msg: Some(quic_client_message::Msg::Hello(req)),
73        };
74        let encoded = msg.encode_to_vec();
75        assert!(!encoded.is_empty());
76
77        // Decode the client message as a QuicClientMessage
78        let decoded = QuicClientMessage::decode(encoded.as_slice()).unwrap();
79        match decoded.msg {
80            Some(quic_client_message::Msg::Hello(hello)) => {
81                assert_eq!(hello.username, "admin");
82                assert_eq!(hello.password, "secret");
83                assert_eq!(hello.tenant_id, Some("tenant1".to_string()));
84            }
85            _ => panic!("Expected Hello variant"),
86        }
87    }
88
89    #[test]
90    fn test_hello_request_encodes_tenant_and_role() {
91        let req = HelloRequest {
92            username: "admin".to_string(),
93            password: "secret".to_string(),
94            tenant_id: Some("tenant1".to_string()),
95            client_name: String::new(),
96            client_version: String::new(),
97            wanted_conformance: String::new(),
98            graph: None,
99            role: Some("analyst".to_string()),
100        };
101        let encoded = req.encode_to_vec();
102        let decoded = HelloRequest::decode(encoded.as_slice()).unwrap();
103        assert_eq!(decoded.tenant_id, Some("tenant1".to_string()));
104        assert_eq!(decoded.role, Some("analyst".to_string()));
105
106        // Field 8 wire tag (8 << 3 | 2 == 0x42) must appear for the optional string role.
107        assert!(
108            encoded.windows(1).any(|w| w[0] == 0x42),
109            "expected wire tag for field 8 (role) in encoded bytes"
110        );
111    }
112
113    #[test]
114    fn test_encode_decode_execute_roundtrip() {
115        let params = vec![
116            Param {
117                name: "name".to_string(),
118                value: Some(Value {
119                    kind: Some(value::Kind::StringVal(StringValue {
120                        value: "Alice".to_string(),
121                        kind: 0,
122                    })),
123                }),
124            },
125            Param {
126                name: "age".to_string(),
127                value: Some(Value {
128                    kind: Some(value::Kind::IntVal(IntValue { value: 30, kind: 0 })),
129                }),
130            },
131        ];
132
133        let req = ExecuteRequest {
134            session_id: "session123".to_string(),
135            query: "MATCH (n) RETURN n".to_string(),
136            params,
137        };
138        let msg = QuicClientMessage {
139            msg: Some(quic_client_message::Msg::Execute(req)),
140        };
141        let encoded = msg.encode_to_vec();
142        assert!(!encoded.is_empty());
143
144        let decoded = QuicClientMessage::decode(encoded.as_slice()).unwrap();
145        match decoded.msg {
146            Some(quic_client_message::Msg::Execute(exec)) => {
147                assert_eq!(exec.session_id, "session123");
148                assert_eq!(exec.query, "MATCH (n) RETURN n");
149                assert_eq!(exec.params.len(), 2);
150            }
151            _ => panic!("Expected Execute variant"),
152        }
153    }
154
155    #[test]
156    fn test_encode_with_length_prefix() {
157        let msg = QuicClientMessage {
158            msg: Some(quic_client_message::Msg::Ping(PingRequest {})),
159        };
160        let encoded = encode_with_length_prefix(&msg);
161        // Should have 4-byte length prefix
162        assert!(encoded.len() >= 4);
163        let length = u32::from_be_bytes([encoded[0], encoded[1], encoded[2], encoded[3]]);
164        assert_eq!(length as usize, encoded.len() - 4);
165    }
166
167    #[test]
168    fn test_decode_length_prefix() {
169        let data = [0x00, 0x00, 0x00, 0x10];
170        let length = decode_length_prefix(&data).unwrap();
171        assert_eq!(length, 16);
172    }
173
174    #[test]
175    fn test_decode_length_prefix_insufficient_data() {
176        let data = [0x00, 0x00];
177        let result = decode_length_prefix(&data);
178        assert!(result.is_err());
179    }
180
181    #[test]
182    fn test_decode_hello_response() {
183        // Build a HelloResponse, encode it, then decode
184        let resp = HelloResponse {
185            success: true,
186            session_id: "sess123".to_string(),
187            error_message: String::new(),
188            capabilities: vec![],
189            password_reset_required: false,
190            graph: None,
191        };
192        let encoded = resp.encode_to_vec();
193        let decoded = HelloResponse::decode(encoded.as_slice()).unwrap();
194        assert!(decoded.success);
195        assert_eq!(decoded.session_id, "sess123");
196    }
197
198    #[test]
199    fn test_decode_ping_response() {
200        let resp = PingResponse { ok: true };
201        let encoded = resp.encode_to_vec();
202        let decoded = PingResponse::decode(encoded.as_slice()).unwrap();
203        assert!(decoded.ok);
204    }
205
206    #[test]
207    fn test_value_null() {
208        let val = Value {
209            kind: Some(value::Kind::NullVal(NullValue {})),
210        };
211        assert!(matches!(val.kind, Some(value::Kind::NullVal(_))));
212    }
213
214    #[test]
215    fn test_value_default() {
216        let val = Value::default();
217        assert!(val.kind.is_none());
218    }
219
220    #[test]
221    fn test_message_defaults() {
222        let client_msg = QuicClientMessage::default();
223        assert!(client_msg.msg.is_none());
224
225        let server_msg = QuicServerMessage::default();
226        assert!(server_msg.msg.is_none());
227    }
228}