use std::path::{Path, PathBuf};
use prost_reflect::{DescriptorPool, DynamicMessage, MessageDescriptor};
use waf_core::{Config, Decision, ParsedBody, Phase, RequestContext, ScoreItem, Severity, WafModule};
use crate::rules;
#[derive(Debug, Clone)]
pub struct GrpcSchemaEnforcementConfig {
pub enabled: bool,
pub descriptor_path: PathBuf,
pub flag_unknown_fields: bool,
pub unknown_method_severity: Severity,
pub decode_error_severity: Severity,
pub unknown_field_severity: Severity,
}
impl Default for GrpcSchemaEnforcementConfig {
fn default() -> Self {
Self {
enabled: false,
descriptor_path: PathBuf::new(),
flag_unknown_fields: false,
unknown_method_severity: Severity::Critical,
decode_error_severity: Severity::Warning,
unknown_field_severity: Severity::Notice,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GrpcValidity {
Valid,
UnknownMethod,
DecodeError,
UnknownField,
}
pub fn parse_descriptors(fds_bytes: &[u8]) -> Result<DescriptorPool, String> {
DescriptorPool::decode(fds_bytes).map_err(|e| format!("invalid FileDescriptorSet: {e}"))
}
pub fn load_descriptors(path: &Path) -> Result<DescriptorPool, String> {
let bytes = std::fs::read(path)
.map_err(|e| format!("cannot read descriptor file {}: {e}", path.display()))?;
parse_descriptors(&bytes)
}
fn resolve_input(pool: &DescriptorPool, path: &str) -> Option<MessageDescriptor> {
let (service, method) = path.strip_prefix('/')?.split_once('/')?;
let svc = pool.get_service_by_name(service)?;
let m = svc.methods().find(|m| m.name() == method)?;
Some(m.input())
}
pub fn validate_message(
pool: &DescriptorPool,
path: &str,
msg: &[u8],
flag_unknown: bool,
) -> GrpcValidity {
let Some(input) = resolve_input(pool, path) else {
return GrpcValidity::UnknownMethod;
};
match DynamicMessage::decode(input.clone(), msg) {
Err(_) => GrpcValidity::DecodeError,
Ok(_) => {
if flag_unknown {
let known: Vec<u32> = input.fields().map(|f| f.number()).collect();
if has_unknown_field(&known, msg) {
return GrpcValidity::UnknownField;
}
}
GrpcValidity::Valid
}
}
}
fn read_varint(b: &[u8], i: &mut usize) -> Option<u64> {
let mut result: u64 = 0;
for k in 0..10 {
let byte = *b.get(*i)?;
*i += 1;
result |= ((byte & 0x7f) as u64) << (7 * k);
if byte & 0x80 == 0 {
return Some(result);
}
}
None
}
fn has_unknown_field(known: &[u32], msg: &[u8]) -> bool {
let mut i = 0;
while i < msg.len() {
let Some(tag) = read_varint(msg, &mut i) else { return false };
let field = (tag >> 3) as u32;
let wire = (tag & 7) as u8;
if field != 0 && !known.contains(&field) {
return true;
}
let ok = match wire {
0 => read_varint(msg, &mut i).is_some(), 1 => {
i += 8;
true
} 2 => match read_varint(msg, &mut i) {
Some(len) => {
i += len as usize;
true
}
None => false,
}, 5 => {
i += 4;
true
} _ => false, };
if !ok || i > msg.len() {
return false;
}
}
false
}
fn frames(body: &[u8]) -> Vec<&[u8]> {
let mut out = Vec::new();
let mut pos = 0;
while pos + 5 <= body.len() {
let flag = body[pos];
let len = u32::from_be_bytes([body[pos + 1], body[pos + 2], body[pos + 3], body[pos + 4]])
as usize;
pos += 5;
if pos + len > body.len() {
break; }
if flag == 0 {
out.push(&body[pos..pos + len]);
}
pos += len;
}
out
}
fn content_type(ctx: &RequestContext) -> &str {
ctx.normalized
.headers
.iter()
.find(|(k, _)| k == "content-type")
.map(|(_, v)| v.as_str())
.unwrap_or("")
}
pub struct GrpcSchemaModule {
pool: DescriptorPool,
enabled: bool,
flag_unknown_fields: bool,
unknown_method_severity: Severity,
decode_error_severity: Severity,
unknown_field_severity: Severity,
}
impl GrpcSchemaModule {
pub fn from_config(cfg: &GrpcSchemaEnforcementConfig) -> Result<Self, String> {
Ok(Self {
pool: load_descriptors(&cfg.descriptor_path)?,
enabled: cfg.enabled,
flag_unknown_fields: cfg.flag_unknown_fields,
unknown_method_severity: cfg.unknown_method_severity,
decode_error_severity: cfg.decode_error_severity,
unknown_field_severity: cfg.unknown_field_severity,
})
}
pub fn from_fds_bytes(fds_bytes: &[u8], enabled: bool) -> Result<Self, String> {
let cfg = GrpcSchemaEnforcementConfig::default();
Ok(Self {
pool: parse_descriptors(fds_bytes)?,
enabled,
flag_unknown_fields: cfg.flag_unknown_fields,
unknown_method_severity: cfg.unknown_method_severity,
decode_error_severity: cfg.decode_error_severity,
unknown_field_severity: cfg.unknown_field_severity,
})
}
#[cfg(test)]
fn from_fds_bytes_flagging(fds_bytes: &[u8]) -> Result<Self, String> {
Ok(Self { flag_unknown_fields: true, ..Self::from_fds_bytes(fds_bytes, true)? })
}
}
impl WafModule for GrpcSchemaModule {
fn id(&self) -> &str {
"schema-grpc"
}
fn phase(&self) -> Phase {
Phase::Body
}
fn structural(&self) -> bool {
true
}
fn init(&mut self, _cfg: &Config) {}
fn inspect(&self, ctx: &RequestContext) -> Decision {
if !self.enabled {
return Decision::Allow;
}
if !content_type(ctx).trim_start().starts_with("application/grpc") {
return Decision::Allow;
}
let ParsedBody::Raw(body) = &ctx.normalized.body else {
return Decision::Allow;
};
let msgs = frames(body);
if msgs.is_empty() {
return Decision::Allow;
}
let path = &ctx.path;
let mut unknown_method = false;
let mut decode_error = false;
let mut unknown_field = false;
for msg in &msgs {
match validate_message(&self.pool, path, msg, self.flag_unknown_fields) {
GrpcValidity::Valid => {}
GrpcValidity::UnknownMethod => unknown_method = true,
GrpcValidity::DecodeError => decode_error = true,
GrpcValidity::UnknownField => unknown_field = true,
}
}
let mut items = Vec::new();
if unknown_method {
items.push(ScoreItem {
rule_id: rules::GRPC_UNKNOWN_METHOD.id(),
severity: self.unknown_method_severity,
});
}
if decode_error {
items.push(ScoreItem {
rule_id: rules::GRPC_DECODE_ERROR.id(),
severity: self.decode_error_severity,
});
}
if unknown_field {
items.push(ScoreItem {
rule_id: rules::GRPC_UNKNOWN_FIELD.id(),
severity: self.unknown_field_severity,
});
}
if items.is_empty() {
Decision::Allow
} else {
Decision::Scores(items)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use prost::Message;
fn fixture_fds() -> Vec<u8> {
use std::sync::atomic::{AtomicU64, Ordering};
static SEQ: AtomicU64 = AtomicU64::new(0);
let dir = std::env::temp_dir().join(format!(
"waf-grpc-{}-{}",
std::process::id(),
SEQ.fetch_add(1, Ordering::Relaxed)
));
std::fs::create_dir_all(&dir).unwrap();
let proto = dir.join("greeter.proto");
std::fs::write(
&proto,
r#"
syntax = "proto3";
package test.v1;
message HelloRequest { string name = 1; }
message HelloReply { string message = 1; }
service Greeter { rpc SayHello(HelloRequest) returns (HelloReply); }
"#,
)
.unwrap();
let fds = protox::compile([&proto], [&dir]).expect("compile proto");
let _ = std::fs::remove_dir_all(&dir);
fds.encode_to_vec()
}
const SAYHELLO: &str = "/test.v1.Greeter/SayHello";
#[test]
#[ignore]
fn emit_fixture_descriptor() {
let out = std::env::var("WAF_FDS_OUT").expect("set WAF_FDS_OUT to the output .bin path");
std::fs::write(&out, fixture_fds()).unwrap();
eprintln!("wrote demo FileDescriptorSet to {out}");
}
#[test]
fn valid_descriptor_loads() {
assert!(parse_descriptors(&fixture_fds()).is_ok());
}
#[test]
fn invalid_descriptor_bytes_are_rejected() {
assert!(parse_descriptors(b"not a descriptor set").is_err());
}
#[test]
fn valid_message_against_known_method_is_valid() {
let pool = parse_descriptors(&fixture_fds()).unwrap();
let msg = [0x0a, 0x01, b'x'];
assert_eq!(validate_message(&pool, SAYHELLO, &msg, false), GrpcValidity::Valid);
}
#[test]
fn unknown_method_is_flagged() {
let pool = parse_descriptors(&fixture_fds()).unwrap();
assert_eq!(
validate_message(&pool, "/test.v1.Greeter/DoesNotExist", &[0x0a, 0x01, b'x'], false),
GrpcValidity::UnknownMethod
);
assert_eq!(
validate_message(&pool, "/test.v1.Unknown/SayHello", &[], false),
GrpcValidity::UnknownMethod
);
}
#[test]
fn undecodable_message_is_decode_error() {
let pool = parse_descriptors(&fixture_fds()).unwrap();
assert_eq!(validate_message(&pool, SAYHELLO, &[0x0a, 0xff], false), GrpcValidity::DecodeError);
}
const MSG_WITH_UNKNOWN_FIELD: [u8; 6] = [0x0a, 0x01, b'x', 0x2a, 0x01, b'y'];
#[test]
fn unknown_field_reported_only_when_flag_on() {
let pool = parse_descriptors(&fixture_fds()).unwrap();
assert_eq!(
validate_message(&pool, SAYHELLO, &MSG_WITH_UNKNOWN_FIELD, true),
GrpcValidity::UnknownField
);
assert_eq!(
validate_message(&pool, SAYHELLO, &MSG_WITH_UNKNOWN_FIELD, false),
GrpcValidity::Valid
);
}
#[test]
fn known_only_message_is_valid_even_with_flag_on() {
let pool = parse_descriptors(&fixture_fds()).unwrap();
assert_eq!(validate_message(&pool, SAYHELLO, &[0x0a, 0x01, b'x'], true), GrpcValidity::Valid);
}
#[test]
fn has_unknown_field_walks_top_level_tags() {
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'])); }
fn temp_path(tag: &str) -> std::path::PathBuf {
let uniq = format!(
"waf-fds-{tag}-{}-{:?}.bin",
std::process::id(),
std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos()
);
std::env::temp_dir().join(uniq)
}
#[test]
fn load_descriptors_reads_and_validates_a_file() {
let path = temp_path("ok");
std::fs::write(&path, fixture_fds()).unwrap();
let loaded = load_descriptors(&path);
let _ = std::fs::remove_file(&path);
assert!(loaded.is_ok());
}
#[test]
fn load_descriptors_missing_file_is_err() {
assert!(load_descriptors(&temp_path("missing")).is_err());
}
#[test]
fn from_config_is_boot_fatal_on_invalid_descriptor() {
let path = temp_path("bad");
std::fs::write(&path, b"not a descriptor set").unwrap();
let cfg = GrpcSchemaEnforcementConfig { descriptor_path: path.clone(), ..Default::default() };
let built = GrpcSchemaModule::from_config(&cfg);
let _ = std::fs::remove_file(&path);
assert!(built.is_err(), "an invalid descriptor file must fail construction (fail-closed boot)");
}
#[test]
fn from_config_builds_from_a_valid_file() {
let path = temp_path("good");
std::fs::write(&path, fixture_fds()).unwrap();
let cfg = GrpcSchemaEnforcementConfig {
enabled: true,
descriptor_path: path.clone(),
..Default::default()
};
let built = GrpcSchemaModule::from_config(&cfg);
let _ = std::fs::remove_file(&path);
assert!(built.is_ok());
}
#[test]
fn default_config_is_off_and_unknown_fields_disabled() {
let c = GrpcSchemaEnforcementConfig::default();
assert!(!c.enabled);
assert!(!c.flag_unknown_fields, "unknown-field flagging is opt-in (H3)");
}
#[test]
fn deframes_and_skips_compressed() {
let mut body = vec![0u8];
body.extend_from_slice(&3u32.to_be_bytes());
body.extend_from_slice(&[0x0a, 0x01, b'x']);
body.push(1u8); body.extend_from_slice(&2u32.to_be_bytes());
body.extend_from_slice(&[0xaa, 0xbb]);
let f = frames(&body);
assert_eq!(f.len(), 1, "compressed frame must be skipped");
assert_eq!(f[0], &[0x0a, 0x01, b'x']);
}
#[test]
fn frames_empty_and_truncated_are_safe() {
assert!(frames(&[]).is_empty());
let mut trunc = vec![0u8];
trunc.extend_from_slice(&9u32.to_be_bytes());
trunc.extend_from_slice(&[0x0a, 0x01]);
assert!(frames(&trunc).is_empty());
}
use waf_core::{Bytes, Normalized};
fn frame(msg: &[u8]) -> Vec<u8> {
let mut b = vec![0u8];
b.extend_from_slice(&(msg.len() as u32).to_be_bytes());
b.extend_from_slice(msg);
b
}
fn ctx(path: &str, ctype: &str, body: ParsedBody) -> RequestContext {
let normalized = Normalized {
path: path.to_string(),
headers: vec![("content-type".to_string(), ctype.to_string())],
body,
..Default::default()
};
RequestContext {
client_ip: "127.0.0.1".parse().unwrap(),
request_id: "t".to_string(),
timestamp: std::time::SystemTime::now(),
method: "POST".to_string(),
path: path.to_string(),
raw_path: path.to_string(),
query: None,
http_version: "HTTP/2".to_string(),
headers: vec![],
cookies: vec![],
body: Bytes::new(),
normalized,
score: 0,
score_contributions: vec![],
}
}
fn module() -> GrpcSchemaModule {
GrpcSchemaModule::from_fds_bytes(&fixture_fds(), true).unwrap()
}
fn scored(d: &Decision) -> Vec<String> {
match d {
Decision::Scores(items) => items.iter().map(|i| i.rule_id.clone()).collect(),
_ => vec![],
}
}
#[test]
fn inspect_dedups_multi_frame_unknown_method() {
let mut body = frame(&[0x0a, 0x01, b'x']);
body.extend_from_slice(&frame(&[0x0a, 0x01, b'y']));
let c = ctx("/test.v1.Greeter/Nope", "application/grpc", ParsedBody::Raw(Bytes::from(body)));
assert_eq!(scored(&module().inspect(&c)), vec![rules::GRPC_UNKNOWN_METHOD.id()]);
}
#[test]
fn inspect_scores_decode_error_on_a_known_method_ignoring_valid_frames() {
let mut body = frame(&[0x0a, 0x01, b'x']);
body.extend_from_slice(&frame(&[0x0a, 0xff])); let c = ctx("/test.v1.Greeter/SayHello", "application/grpc", ParsedBody::Raw(Bytes::from(body)));
assert_eq!(scored(&module().inspect(&c)), vec![rules::GRPC_DECODE_ERROR.id()]);
}
#[test]
fn inspect_accepts_grpc_proto_content_type_variant() {
let c = ctx(
"/test.v1.Greeter/Nope",
"application/grpc+proto",
ParsedBody::Raw(Bytes::from(frame(&[0x0a, 0x01, b'x']))),
);
assert_eq!(scored(&module().inspect(&c)), vec![rules::GRPC_UNKNOWN_METHOD.id()]);
}
#[test]
fn inspect_ignores_non_grpc_content_type() {
let c = ctx(
"/test.v1.Greeter/Nope",
"application/json",
ParsedBody::Raw(Bytes::from(frame(&[0x0a, 0x01, b'x']))),
);
assert_eq!(module().inspect(&c), Decision::Allow);
}
#[test]
fn inspect_ignores_non_raw_body() {
let c = ctx(
"/test.v1.Greeter/Nope",
"application/grpc",
ParsedBody::JsonFlattened(vec![("a".to_string(), "b".to_string())]),
);
assert_eq!(module().inspect(&c), Decision::Allow);
}
#[test]
fn inspect_flags_unknown_field_only_with_flagging_module() {
let body = ParsedBody::Raw(Bytes::from(frame(&MSG_WITH_UNKNOWN_FIELD)));
let c = ctx("/test.v1.Greeter/SayHello", "application/grpc", body);
assert_eq!(module().inspect(&c), Decision::Allow);
let flagging = GrpcSchemaModule::from_fds_bytes_flagging(&fixture_fds()).unwrap();
assert_eq!(scored(&flagging.inspect(&c)), vec![rules::GRPC_UNKNOWN_FIELD.id()]);
}
#[test]
fn inspect_resolves_method_from_case_sensitive_raw_path() {
let mut c = ctx(
"/test.v1.Greeter/SayHello",
"application/grpc",
ParsedBody::Raw(Bytes::from(frame(&[0x0a, 0x01, b'x']))),
);
c.normalized.path = "/test.v1.greeter/sayhello".to_string(); assert_eq!(module().inspect(&c), Decision::Allow);
}
}