feat(guardrail): add regex

.github/workflows/build.yaml

@@ -2,9 +2,9 @@ name: Build
 
 on:
   push:
-    branches: [ "main" ]
+    branches: ["main"]
   pull_request:
-    branches: [ "main" ]
+    branches: ["main"]
 
 env:
   CARGO_TERM_COLOR: always
@@ -32,6 +32,15 @@ jobs:
       - name: Setup Rust toolchain
         run: rustup toolchain install stable --profile minimal
 
+      - name: Install etcdctl
+        env:
+          ETCD_VER: v3.6.11
+        run: |
+          curl -fsSL "https://github.com/etcd-io/etcd/releases/download/${ETCD_VER}/etcd-${ETCD_VER}-linux-amd64.tar.gz" -o /tmp/etcd-${ETCD_VER}-linux-amd64.tar.gz
+          tar xzf /tmp/etcd-${ETCD_VER}-linux-amd64.tar.gz -C /tmp
+          sudo install /tmp/etcd-${ETCD_VER}-linux-amd64/etcdctl /usr/local/bin/etcdctl
+          etcdctl version
+
       - name: Setup environment
         run: sudo docker compose -f ci/docker-compose.yaml up -d
 

Cargo.toml

@@ -38,6 +38,7 @@ reqwest = { version = "0.13", default-features = false, features = [
   "stream",
   "native-tls",
 ] }
+regex = "1"
 serde = "1.0.228"
 serde_json = "1.0"
 thiserror = "2"
@@ -63,6 +64,7 @@ tokio.workspace = true
 reqwest.workspace = true
 http.workspace = true
 axum.workspace = true
+regex.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 async-trait.workspace = true

crates/aisix-guardrail/Cargo.toml

@@ -11,8 +11,12 @@ aws-sigv4.workspace = true
 aws-smithy-runtime-api.workspace = true
 http.workspace = true
 percent-encoding.workspace = true
+regex.workspace = true
 reqwest.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 thiserror.workspace = true
 utoipa.workspace = true
+
+[dev-dependencies]
+tokio.workspace = true

crates/aisix-guardrail/src/guardrails/mod.rs

@@ -1,13 +1,18 @@
 pub mod bedrock;
+pub mod regex;
 
-pub use bedrock::{BedrockGuardrailMeta, BedrockGuardrailRuntime};
+pub use self::{
+    bedrock::{BedrockGuardrailMeta, BedrockGuardrailRuntime},
+    regex::{RegexGuardrailMeta, RegexGuardrailRuntime},
+};
 
 pub mod identifiers {
-    use super::bedrock;
+    use super::{bedrock, regex};
 
     pub const BEDROCK: &str = bedrock::IDENTIFIER;
+    pub const REGEX: &str = regex::IDENTIFIER;
 }
 
 pub mod configs {
-    pub use super::bedrock::BedrockGuardrailConfig;
+    pub use super::{bedrock::BedrockGuardrailConfig, regex::RegexGuardrailConfig};
 }

crates/aisix-guardrail/src/guardrails/regex.rs

@@ -0,0 +1,240 @@
+use std::convert::Infallible;
+
+use async_trait::async_trait;
+use regex::Regex;
+use serde::{Deserialize, Deserializer, Serialize, de};
+use utoipa::ToSchema;
+
+use crate::traits::{
+    GuardrailCheckPayload, GuardrailContentPart, GuardrailMessage, GuardrailMessageContent,
+    GuardrailMeta, GuardrailOutcome, GuardrailRuntime,
+};
+
+pub const IDENTIFIER: &str = "regex";
+
+#[derive(Debug, Clone, Serialize, ToSchema)]
+pub struct RegexGuardrailConfig {
+    pub pattern: String,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub block_reason: Option<String>,
+
+    #[serde(skip)]
+    #[schema(ignore)]
+    compiled_pattern: Regex,
+}
+
+impl RegexGuardrailConfig {
+    pub fn new(
+        pattern: impl Into<String>,
+        block_reason: Option<String>,
+    ) -> Result<Self, regex::Error> {
+        let pattern = pattern.into();
+        let compiled_pattern = Regex::new(&pattern)?;
+
+        Ok(Self {
+            pattern,
+            block_reason,
+            compiled_pattern,
+        })
+    }
+
+    pub fn compiled_pattern(&self) -> &Regex {
+        &self.compiled_pattern
+    }
+}
+
+impl<'de> Deserialize<'de> for RegexGuardrailConfig {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        #[derive(Deserialize)]
+        struct RawRegexGuardrailConfig {
+            pattern: String,
+            #[serde(default)]
+            block_reason: Option<String>,
+        }
+
+        let raw = RawRegexGuardrailConfig::deserialize(deserializer)?;
+
+        Self::new(raw.pattern, raw.block_reason)
+            .map_err(|error| de::Error::custom(format!("invalid regex guardrail pattern: {error}")))
+    }
+}
+
+#[derive(Debug, Clone, Copy, Default)]
+pub struct RegexGuardrailMeta;
+
+impl GuardrailMeta for RegexGuardrailMeta {
+    fn name(&self) -> &'static str {
+        IDENTIFIER
+    }
+}
+
+#[derive(Debug, Clone, Copy, Default)]
+pub struct RegexGuardrailRuntime;
+
+impl RegexGuardrailRuntime {
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+impl GuardrailMeta for RegexGuardrailRuntime {
+    fn name(&self) -> &'static str {
+        IDENTIFIER
+    }
+}
+
+#[async_trait]
+impl GuardrailRuntime<RegexGuardrailConfig> for RegexGuardrailRuntime {
+    type Error = Infallible;
+
+    async fn check(
+        &self,
+        payload: &GuardrailCheckPayload,
+        config: &RegexGuardrailConfig,
+    ) -> Result<GuardrailOutcome, Self::Error> {
+        if payload_matches(config.compiled_pattern(), payload) {
+            return Ok(GuardrailOutcome::Block {
+                reason: config
+                    .block_reason
+                    .clone()
+                    .unwrap_or_else(|| "regex guardrail blocked".into()),
+            });
+        }
+
+        Ok(GuardrailOutcome::Allow)
+    }
+}
+
+fn payload_matches(pattern: &Regex, payload: &GuardrailCheckPayload) -> bool {
+    let messages = match payload {
+        GuardrailCheckPayload::Input(payload) => &payload.messages,
+        GuardrailCheckPayload::Output(payload) => &payload.messages,
+    };
+
+    messages
+        .iter()
+        .any(|message| message_matches(pattern, message))
+}
+
+fn message_matches(pattern: &Regex, message: &GuardrailMessage) -> bool {
+    match &message.content {
+        Some(GuardrailMessageContent::Text(text)) => pattern.is_match(text),
+        Some(GuardrailMessageContent::Parts(parts)) => parts.iter().any(|part| match part {
+            GuardrailContentPart::Text { text } => pattern.is_match(text),
+            GuardrailContentPart::ImageUrl { .. } => false,
+        }),
+        None => false,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use serde_json::json;
+
+    use super::{RegexGuardrailConfig, RegexGuardrailRuntime};
+    use crate::traits::{
+        GuardrailCheckPayload, GuardrailContentPart, GuardrailImageUrl, GuardrailMessage,
+        GuardrailMessageContent, GuardrailOutcome, GuardrailRole, GuardrailRuntime,
+        InputGuardrailPayload,
+    };
+
+    fn config(pattern: &str) -> RegexGuardrailConfig {
+        RegexGuardrailConfig::new(pattern, Some("matched blocked content".into())).unwrap()
+    }
+
+    fn runtime() -> RegexGuardrailRuntime {
+        RegexGuardrailRuntime::new()
+    }
+
+    fn input_payload(content: GuardrailMessageContent) -> GuardrailCheckPayload {
+        GuardrailCheckPayload::Input(InputGuardrailPayload {
+            messages: vec![GuardrailMessage {
+                role: GuardrailRole::User,
+                content: Some(content),
+                name: None,
+                tool_calls: None,
+                tool_call_id: None,
+            }],
+        })
+    }
+
+    #[tokio::test]
+    async fn blocks_when_plain_text_matches_pattern() {
+        let outcome = runtime()
+            .check(
+                &input_payload(GuardrailMessageContent::Text(
+                    "my secret token is 12345".into(),
+                )),
+                &config(r"secret token"),
+            )
+            .await
+            .unwrap();
+
+        assert_eq!(
+            outcome,
+            GuardrailOutcome::Block {
+                reason: "matched blocked content".into(),
+            }
+        );
+    }
+
+    #[tokio::test]
+    async fn allows_when_no_message_text_matches_pattern() {
+        let outcome = runtime()
+            .check(
+                &input_payload(GuardrailMessageContent::Text("hello world".into())),
+                &config(r"secret token"),
+            )
+            .await
+            .unwrap();
+
+        assert_eq!(outcome, GuardrailOutcome::Allow);
+    }
+
+    #[tokio::test]
+    async fn matches_text_parts_and_ignores_non_text_parts() {
+        let outcome = runtime()
+            .check(
+                &input_payload(GuardrailMessageContent::Parts(vec![
+                    GuardrailContentPart::ImageUrl {
+                        image_url: GuardrailImageUrl {
+                            url: "https://example.com/cat.png".into(),
+                            detail: Some("high".into()),
+                        },
+                    },
+                    GuardrailContentPart::Text {
+                        text: "contains credit card 4111111111111111".into(),
+                    },
+                ])),
+                &config(r"\b\d{16}\b"),
+            )
+            .await
+            .unwrap();
+
+        assert_eq!(
+            outcome,
+            GuardrailOutcome::Block {
+                reason: "matched blocked content".into(),
+            }
+        );
+    }
+
+    #[test]
+    fn deserialize_rejects_invalid_patterns() {
+        let error = serde_json::from_value::<RegexGuardrailConfig>(json!({
+            "pattern": "[",
+            "block_reason": "matched blocked content"
+        }))
+        .unwrap_err();
+
+        assert!(
+            error
+                .to_string()
+                .contains("invalid regex guardrail pattern")
+        );
+    }
+}

src/config/entities/guardrails-schema.json

@@ -5,7 +5,7 @@
     "name": { "type": "string", "minLength": 1 },
     "type": {
       "type": "string",
-      "enum": ["bedrock"]
+      "enum": ["bedrock", "regex"]
     },
     "config": { "type": "object" }
   },
@@ -22,6 +22,17 @@
           "config": { "$ref": "#/$defs/bedrock" }
         }
       }
+    },
+    {
+      "if": {
+        "properties": { "type": { "const": "regex" } },
+        "required": ["type"]
+      },
+      "then": {
+        "properties": {
+          "config": { "$ref": "#/$defs/regex" }
+        }
+      }
     }
   ],
   "$defs": {
@@ -44,6 +55,15 @@
         "endpoint": { "type": "string", "minLength": 1 }
       },
       "additionalProperties": false
+    },
+    "regex": {
+      "type": "object",
+      "required": ["pattern"],
+      "properties": {
+        "pattern": { "type": "string", "minLength": 1 },
+        "block_reason": { "type": "string", "minLength": 1 }
+      },
+      "additionalProperties": false
     }
   }
 }