1use std::path::{Path, PathBuf};
19
20use prost_reflect::{DescriptorPool, DynamicMessage, MessageDescriptor};
21use waf_core::{Config, Decision, ParsedBody, Phase, RequestContext, ScoreItem, Severity, WafModule};
22
23use crate::rules;
24
25#[derive(Debug, Clone)]
28pub struct GrpcSchemaEnforcementConfig {
29 pub enabled: bool,
31 pub descriptor_path: PathBuf,
33 pub flag_unknown_fields: bool,
36 pub unknown_method_severity: Severity,
38 pub decode_error_severity: Severity,
40 pub unknown_field_severity: Severity,
42}
43
44impl Default for GrpcSchemaEnforcementConfig {
45 fn default() -> Self {
46 Self {
47 enabled: false,
48 descriptor_path: PathBuf::new(),
49 flag_unknown_fields: false,
50 unknown_method_severity: Severity::Critical,
51 decode_error_severity: Severity::Warning,
52 unknown_field_severity: Severity::Notice,
53 }
54 }
55}
56
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub enum GrpcValidity {
60 Valid,
62 UnknownMethod,
64 DecodeError,
66 UnknownField,
69}
70
71pub fn parse_descriptors(fds_bytes: &[u8]) -> Result<DescriptorPool, String> {
74 DescriptorPool::decode(fds_bytes).map_err(|e| format!("invalid FileDescriptorSet: {e}"))
75}
76
77pub fn load_descriptors(path: &Path) -> Result<DescriptorPool, String> {
80 let bytes = std::fs::read(path)
81 .map_err(|e| format!("cannot read descriptor file {}: {e}", path.display()))?;
82 parse_descriptors(&bytes)
83}
84
85fn resolve_input(pool: &DescriptorPool, path: &str) -> Option<MessageDescriptor> {
88 let (service, method) = path.strip_prefix('/')?.split_once('/')?;
89 let svc = pool.get_service_by_name(service)?;
90 let m = svc.methods().find(|m| m.name() == method)?;
91 Some(m.input())
92}
93
94pub fn validate_message(
98 pool: &DescriptorPool,
99 path: &str,
100 msg: &[u8],
101 flag_unknown: bool,
102) -> GrpcValidity {
103 let Some(input) = resolve_input(pool, path) else {
104 return GrpcValidity::UnknownMethod;
105 };
106 match DynamicMessage::decode(input.clone(), msg) {
107 Err(_) => GrpcValidity::DecodeError,
108 Ok(_) => {
109 if flag_unknown {
112 let known: Vec<u32> = input.fields().map(|f| f.number()).collect();
113 if has_unknown_field(&known, msg) {
114 return GrpcValidity::UnknownField;
115 }
116 }
117 GrpcValidity::Valid
118 }
119 }
120}
121
122fn read_varint(b: &[u8], i: &mut usize) -> Option<u64> {
124 let mut result: u64 = 0;
125 for k in 0..10 {
126 let byte = *b.get(*i)?;
127 *i += 1;
128 result |= ((byte & 0x7f) as u64) << (7 * k);
129 if byte & 0x80 == 0 {
130 return Some(result);
131 }
132 }
133 None
134}
135
136fn has_unknown_field(known: &[u32], msg: &[u8]) -> bool {
139 let mut i = 0;
140 while i < msg.len() {
141 let Some(tag) = read_varint(msg, &mut i) else { return false };
142 let field = (tag >> 3) as u32;
143 let wire = (tag & 7) as u8;
144 if field != 0 && !known.contains(&field) {
145 return true;
146 }
147 let ok = match wire {
148 0 => read_varint(msg, &mut i).is_some(), 1 => {
150 i += 8;
151 true
152 } 2 => match read_varint(msg, &mut i) {
154 Some(len) => {
155 i += len as usize;
156 true
157 }
158 None => false,
159 }, 5 => {
161 i += 4;
162 true
163 } _ => false, };
166 if !ok || i > msg.len() {
167 return false;
168 }
169 }
170 false
171}
172
173fn frames(body: &[u8]) -> Vec<&[u8]> {
178 let mut out = Vec::new();
179 let mut pos = 0;
180 while pos + 5 <= body.len() {
181 let flag = body[pos];
182 let len = u32::from_be_bytes([body[pos + 1], body[pos + 2], body[pos + 3], body[pos + 4]])
183 as usize;
184 pos += 5;
185 if pos + len > body.len() {
186 break; }
188 if flag == 0 {
189 out.push(&body[pos..pos + len]);
190 }
191 pos += len;
192 }
193 out
194}
195
196fn content_type(ctx: &RequestContext) -> &str {
197 ctx.normalized
198 .headers
199 .iter()
200 .find(|(k, _)| k == "content-type")
201 .map(|(_, v)| v.as_str())
202 .unwrap_or("")
203}
204
205pub struct GrpcSchemaModule {
207 pool: DescriptorPool,
208 enabled: bool,
209 flag_unknown_fields: bool,
210 unknown_method_severity: Severity,
211 decode_error_severity: Severity,
212 unknown_field_severity: Severity,
213}
214
215impl GrpcSchemaModule {
216 pub fn from_config(cfg: &GrpcSchemaEnforcementConfig) -> Result<Self, String> {
220 Ok(Self {
221 pool: load_descriptors(&cfg.descriptor_path)?,
222 enabled: cfg.enabled,
223 flag_unknown_fields: cfg.flag_unknown_fields,
224 unknown_method_severity: cfg.unknown_method_severity,
225 decode_error_severity: cfg.decode_error_severity,
226 unknown_field_severity: cfg.unknown_field_severity,
227 })
228 }
229
230 pub fn from_fds_bytes(fds_bytes: &[u8], enabled: bool) -> Result<Self, String> {
233 let cfg = GrpcSchemaEnforcementConfig::default();
234 Ok(Self {
235 pool: parse_descriptors(fds_bytes)?,
236 enabled,
237 flag_unknown_fields: cfg.flag_unknown_fields,
238 unknown_method_severity: cfg.unknown_method_severity,
239 decode_error_severity: cfg.decode_error_severity,
240 unknown_field_severity: cfg.unknown_field_severity,
241 })
242 }
243
244 #[cfg(test)]
246 fn from_fds_bytes_flagging(fds_bytes: &[u8]) -> Result<Self, String> {
247 Ok(Self { flag_unknown_fields: true, ..Self::from_fds_bytes(fds_bytes, true)? })
248 }
249}
250
251impl WafModule for GrpcSchemaModule {
252 fn id(&self) -> &str {
253 "schema-grpc"
254 }
255
256 fn phase(&self) -> Phase {
257 Phase::Body
258 }
259
260 fn structural(&self) -> bool {
261 true
262 }
263
264 fn init(&mut self, _cfg: &Config) {}
265
266 fn inspect(&self, ctx: &RequestContext) -> Decision {
267 if !self.enabled {
268 return Decision::Allow;
269 }
270 if !content_type(ctx).trim_start().starts_with("application/grpc") {
272 return Decision::Allow;
273 }
274 let ParsedBody::Raw(body) = &ctx.normalized.body else {
275 return Decision::Allow;
276 };
277 let msgs = frames(body);
278 if msgs.is_empty() {
279 return Decision::Allow;
280 }
281 let path = &ctx.path;
285 let mut unknown_method = false;
287 let mut decode_error = false;
288 let mut unknown_field = false;
289 for msg in &msgs {
290 match validate_message(&self.pool, path, msg, self.flag_unknown_fields) {
291 GrpcValidity::Valid => {}
292 GrpcValidity::UnknownMethod => unknown_method = true,
293 GrpcValidity::DecodeError => decode_error = true,
294 GrpcValidity::UnknownField => unknown_field = true,
295 }
296 }
297 let mut items = Vec::new();
298 if unknown_method {
299 items.push(ScoreItem {
300 rule_id: rules::GRPC_UNKNOWN_METHOD.id(),
301 severity: self.unknown_method_severity,
302 });
303 }
304 if decode_error {
305 items.push(ScoreItem {
306 rule_id: rules::GRPC_DECODE_ERROR.id(),
307 severity: self.decode_error_severity,
308 });
309 }
310 if unknown_field {
311 items.push(ScoreItem {
312 rule_id: rules::GRPC_UNKNOWN_FIELD.id(),
313 severity: self.unknown_field_severity,
314 });
315 }
316 if items.is_empty() {
317 Decision::Allow
318 } else {
319 Decision::Scores(items)
320 }
321 }
322}
323
324#[cfg(test)]
325mod tests {
326 use super::*;
327 use prost::Message;
328
329 fn fixture_fds() -> Vec<u8> {
332 use std::sync::atomic::{AtomicU64, Ordering};
335 static SEQ: AtomicU64 = AtomicU64::new(0);
336 let dir = std::env::temp_dir().join(format!(
337 "waf-grpc-{}-{}",
338 std::process::id(),
339 SEQ.fetch_add(1, Ordering::Relaxed)
340 ));
341 std::fs::create_dir_all(&dir).unwrap();
342 let proto = dir.join("greeter.proto");
343 std::fs::write(
344 &proto,
345 r#"
346 syntax = "proto3";
347 package test.v1;
348 message HelloRequest { string name = 1; }
349 message HelloReply { string message = 1; }
350 service Greeter { rpc SayHello(HelloRequest) returns (HelloReply); }
351 "#,
352 )
353 .unwrap();
354 let fds = protox::compile([&proto], [&dir]).expect("compile proto");
355 let _ = std::fs::remove_dir_all(&dir);
356 fds.encode_to_vec()
357 }
358
359 const SAYHELLO: &str = "/test.v1.Greeter/SayHello";
360
361 #[test]
365 #[ignore]
366 fn emit_fixture_descriptor() {
367 let out = std::env::var("WAF_FDS_OUT").expect("set WAF_FDS_OUT to the output .bin path");
368 std::fs::write(&out, fixture_fds()).unwrap();
369 eprintln!("wrote demo FileDescriptorSet to {out}");
370 }
371
372 #[test]
373 fn valid_descriptor_loads() {
374 assert!(parse_descriptors(&fixture_fds()).is_ok());
375 }
376
377 #[test]
378 fn invalid_descriptor_bytes_are_rejected() {
379 assert!(parse_descriptors(b"not a descriptor set").is_err());
380 }
381
382 #[test]
383 fn valid_message_against_known_method_is_valid() {
384 let pool = parse_descriptors(&fixture_fds()).unwrap();
385 let msg = [0x0a, 0x01, b'x'];
387 assert_eq!(validate_message(&pool, SAYHELLO, &msg, false), GrpcValidity::Valid);
388 }
389
390 #[test]
391 fn unknown_method_is_flagged() {
392 let pool = parse_descriptors(&fixture_fds()).unwrap();
393 assert_eq!(
394 validate_message(&pool, "/test.v1.Greeter/DoesNotExist", &[0x0a, 0x01, b'x'], false),
395 GrpcValidity::UnknownMethod
396 );
397 assert_eq!(
398 validate_message(&pool, "/test.v1.Unknown/SayHello", &[], false),
399 GrpcValidity::UnknownMethod
400 );
401 }
402
403 #[test]
404 fn undecodable_message_is_decode_error() {
405 let pool = parse_descriptors(&fixture_fds()).unwrap();
406 assert_eq!(validate_message(&pool, SAYHELLO, &[0x0a, 0xff], false), GrpcValidity::DecodeError);
408 }
409
410 const MSG_WITH_UNKNOWN_FIELD: [u8; 6] = [0x0a, 0x01, b'x', 0x2a, 0x01, b'y'];
415
416 #[test]
417 fn unknown_field_reported_only_when_flag_on() {
418 let pool = parse_descriptors(&fixture_fds()).unwrap();
419 assert_eq!(
421 validate_message(&pool, SAYHELLO, &MSG_WITH_UNKNOWN_FIELD, true),
422 GrpcValidity::UnknownField
423 );
424 assert_eq!(
426 validate_message(&pool, SAYHELLO, &MSG_WITH_UNKNOWN_FIELD, false),
427 GrpcValidity::Valid
428 );
429 }
430
431 #[test]
432 fn known_only_message_is_valid_even_with_flag_on() {
433 let pool = parse_descriptors(&fixture_fds()).unwrap();
434 assert_eq!(validate_message(&pool, SAYHELLO, &[0x0a, 0x01, b'x'], true), GrpcValidity::Valid);
436 }
437
438 #[test]
439 fn has_unknown_field_walks_top_level_tags() {
440 assert!(has_unknown_field(&[1], &MSG_WITH_UNKNOWN_FIELD)); assert!(!has_unknown_field(&[1, 5], &MSG_WITH_UNKNOWN_FIELD)); assert!(!has_unknown_field(&[1], &[0x0a, 0x01, b'x'])); }
444
445 fn temp_path(tag: &str) -> std::path::PathBuf {
446 let uniq = format!(
447 "waf-fds-{tag}-{}-{:?}.bin",
448 std::process::id(),
449 std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos()
450 );
451 std::env::temp_dir().join(uniq)
452 }
453
454 #[test]
455 fn load_descriptors_reads_and_validates_a_file() {
456 let path = temp_path("ok");
457 std::fs::write(&path, fixture_fds()).unwrap();
458 let loaded = load_descriptors(&path);
459 let _ = std::fs::remove_file(&path);
460 assert!(loaded.is_ok());
461 }
462
463 #[test]
464 fn load_descriptors_missing_file_is_err() {
465 assert!(load_descriptors(&temp_path("missing")).is_err());
466 }
467
468 #[test]
469 fn from_config_is_boot_fatal_on_invalid_descriptor() {
470 let path = temp_path("bad");
471 std::fs::write(&path, b"not a descriptor set").unwrap();
472 let cfg = GrpcSchemaEnforcementConfig { descriptor_path: path.clone(), ..Default::default() };
473 let built = GrpcSchemaModule::from_config(&cfg);
474 let _ = std::fs::remove_file(&path);
475 assert!(built.is_err(), "an invalid descriptor file must fail construction (fail-closed boot)");
476 }
477
478 #[test]
479 fn from_config_builds_from_a_valid_file() {
480 let path = temp_path("good");
481 std::fs::write(&path, fixture_fds()).unwrap();
482 let cfg = GrpcSchemaEnforcementConfig {
483 enabled: true,
484 descriptor_path: path.clone(),
485 ..Default::default()
486 };
487 let built = GrpcSchemaModule::from_config(&cfg);
488 let _ = std::fs::remove_file(&path);
489 assert!(built.is_ok());
490 }
491
492 #[test]
493 fn default_config_is_off_and_unknown_fields_disabled() {
494 let c = GrpcSchemaEnforcementConfig::default();
495 assert!(!c.enabled);
496 assert!(!c.flag_unknown_fields, "unknown-field flagging is opt-in (H3)");
497 }
498
499 #[test]
500 fn deframes_and_skips_compressed() {
501 let mut body = vec![0u8];
503 body.extend_from_slice(&3u32.to_be_bytes());
504 body.extend_from_slice(&[0x0a, 0x01, b'x']);
505 body.push(1u8); body.extend_from_slice(&2u32.to_be_bytes());
507 body.extend_from_slice(&[0xaa, 0xbb]);
508 let f = frames(&body);
509 assert_eq!(f.len(), 1, "compressed frame must be skipped");
510 assert_eq!(f[0], &[0x0a, 0x01, b'x']);
511 }
512
513 #[test]
514 fn frames_empty_and_truncated_are_safe() {
515 assert!(frames(&[]).is_empty());
516 let mut trunc = vec![0u8];
518 trunc.extend_from_slice(&9u32.to_be_bytes());
519 trunc.extend_from_slice(&[0x0a, 0x01]);
520 assert!(frames(&trunc).is_empty());
521 }
522
523 use waf_core::{Bytes, Normalized};
525
526 fn frame(msg: &[u8]) -> Vec<u8> {
527 let mut b = vec![0u8];
528 b.extend_from_slice(&(msg.len() as u32).to_be_bytes());
529 b.extend_from_slice(msg);
530 b
531 }
532
533 fn ctx(path: &str, ctype: &str, body: ParsedBody) -> RequestContext {
534 let normalized = Normalized {
535 path: path.to_string(),
536 headers: vec![("content-type".to_string(), ctype.to_string())],
537 body,
538 ..Default::default()
539 };
540 RequestContext {
541 client_ip: "127.0.0.1".parse().unwrap(),
542 request_id: "t".to_string(),
543 timestamp: std::time::SystemTime::now(),
544 method: "POST".to_string(),
545 path: path.to_string(),
546 raw_path: path.to_string(),
547 query: None,
548 http_version: "HTTP/2".to_string(),
549 headers: vec![],
550 cookies: vec![],
551 body: Bytes::new(),
552 normalized,
553 score: 0,
554 score_contributions: vec![],
555 }
556 }
557
558 fn module() -> GrpcSchemaModule {
559 GrpcSchemaModule::from_fds_bytes(&fixture_fds(), true).unwrap()
560 }
561
562 fn scored(d: &Decision) -> Vec<String> {
563 match d {
564 Decision::Scores(items) => items.iter().map(|i| i.rule_id.clone()).collect(),
565 _ => vec![],
566 }
567 }
568
569 #[test]
570 fn inspect_dedups_multi_frame_unknown_method() {
571 let mut body = frame(&[0x0a, 0x01, b'x']);
573 body.extend_from_slice(&frame(&[0x0a, 0x01, b'y']));
574 let c = ctx("/test.v1.Greeter/Nope", "application/grpc", ParsedBody::Raw(Bytes::from(body)));
575 assert_eq!(scored(&module().inspect(&c)), vec![rules::GRPC_UNKNOWN_METHOD.id()]);
576 }
577
578 #[test]
579 fn inspect_scores_decode_error_on_a_known_method_ignoring_valid_frames() {
580 let mut body = frame(&[0x0a, 0x01, b'x']);
582 body.extend_from_slice(&frame(&[0x0a, 0xff])); let c = ctx("/test.v1.Greeter/SayHello", "application/grpc", ParsedBody::Raw(Bytes::from(body)));
584 assert_eq!(scored(&module().inspect(&c)), vec![rules::GRPC_DECODE_ERROR.id()]);
585 }
586
587 #[test]
588 fn inspect_accepts_grpc_proto_content_type_variant() {
589 let c = ctx(
590 "/test.v1.Greeter/Nope",
591 "application/grpc+proto",
592 ParsedBody::Raw(Bytes::from(frame(&[0x0a, 0x01, b'x']))),
593 );
594 assert_eq!(scored(&module().inspect(&c)), vec![rules::GRPC_UNKNOWN_METHOD.id()]);
595 }
596
597 #[test]
598 fn inspect_ignores_non_grpc_content_type() {
599 let c = ctx(
600 "/test.v1.Greeter/Nope",
601 "application/json",
602 ParsedBody::Raw(Bytes::from(frame(&[0x0a, 0x01, b'x']))),
603 );
604 assert_eq!(module().inspect(&c), Decision::Allow);
605 }
606
607 #[test]
608 fn inspect_ignores_non_raw_body() {
609 let c = ctx(
611 "/test.v1.Greeter/Nope",
612 "application/grpc",
613 ParsedBody::JsonFlattened(vec![("a".to_string(), "b".to_string())]),
614 );
615 assert_eq!(module().inspect(&c), Decision::Allow);
616 }
617
618 #[test]
619 fn inspect_flags_unknown_field_only_with_flagging_module() {
620 let body = ParsedBody::Raw(Bytes::from(frame(&MSG_WITH_UNKNOWN_FIELD)));
621 let c = ctx("/test.v1.Greeter/SayHello", "application/grpc", body);
622
623 assert_eq!(module().inspect(&c), Decision::Allow);
625
626 let flagging = GrpcSchemaModule::from_fds_bytes_flagging(&fixture_fds()).unwrap();
628 assert_eq!(scored(&flagging.inspect(&c)), vec![rules::GRPC_UNKNOWN_FIELD.id()]);
629 }
630
631 #[test]
632 fn inspect_resolves_method_from_case_sensitive_raw_path() {
633 let mut c = ctx(
637 "/test.v1.Greeter/SayHello",
638 "application/grpc",
639 ParsedBody::Raw(Bytes::from(frame(&[0x0a, 0x01, b'x']))),
640 );
641 c.normalized.path = "/test.v1.greeter/sayhello".to_string(); assert_eq!(module().inspect(&c), Decision::Allow);
643 }
644}