From 9573aa1fbd68cbaf505d99a9b2873938809b6eea Mon Sep 17 00:00:00 2001 From: HY-love-sleep <583699747@qq.com> Date: Mon, 21 Sep 2026 11:20:58 +0800 Subject: [PATCH 1/2] feat: add the sensitive word plugin Reject a request whose body contains a word of the dictionary configured for the matched rule, so that a forbidden prompt never reaches the model. The dictionary is a redis set and it is compiled into an Aho-Corasick automaton, which scans the whole dictionary in a single pass and reports every matching word, nested and overlapping ones included. - the automaton is compiled on a bounded elastic thread and cached per rule dictionary, so rules pointing to different dictionaries never share an automaton and the gateway event loop is never blocked - a cached dictionary is read from redis again after refreshIntervalSeconds, and it is dropped immediately when the rule is configured again - a redis failure is logged and the request is passed through (fail open) - only the request body is inspected, so the plugin is not limited to AI routes - unit tests for the automaton, the plugin, the data handler, the rule handle and the starter --- shenyu-bootstrap/pom.xml | 8 + .../dto/convert/rule/SensitiveWordHandle.java | 86 +++++++++ .../shenyu/common/enums/PluginEnum.java | 5 + .../convert/rule/SensitiveWordHandleTest.java | 44 +++++ shenyu-plugin/shenyu-plugin-ai/pom.xml | 1 + .../shenyu-plugin-ai-sensitive-word/README.md | 60 +++++++ .../shenyu-plugin-ai-sensitive-word/pom.xml | 41 +++++ .../sensitive/word/SensitiveWordPlugin.java | 148 ++++++++++++++++ .../ai/sensitive/word/ac/AhoCorasick.java | 150 ++++++++++++++++ .../SensitiveWordPluginDataHandler.java | 161 +++++++++++++++++ .../word/SensitiveWordPluginTest.java | 167 ++++++++++++++++++ .../ai/sensitive/word/ac/AhoCorasickTest.java | 125 +++++++++++++ .../SensitiveWordPluginDataHandlerTest.java | 151 ++++++++++++++++ .../shenyu-spring-boot-starter-plugin/pom.xml | 1 + .../pom.xml | 35 ++++ .../SensitiveWordPluginConfiguration.java | 57 ++++++ .../main/resources/META-INF/spring.factories | 19 ++ .../main/resources/META-INF/spring.provides | 18 ++ ...ot.autoconfigure.AutoConfiguration.imports | 18 ++ .../SensitiveWordPluginConfigurationTest.java | 65 +++++++ 20 files changed, 1360 insertions(+) create mode 100644 shenyu-common/src/main/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandle.java create mode 100644 shenyu-common/src/test/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandleTest.java create mode 100644 shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/README.md create mode 100644 shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/pom.xml create mode 100644 shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPlugin.java create mode 100644 shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/ac/AhoCorasick.java create mode 100644 shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandler.java create mode 100644 shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPluginTest.java create mode 100644 shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/ac/AhoCorasickTest.java create mode 100644 shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandlerTest.java create mode 100644 shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/pom.xml create mode 100644 shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/springboot/starter/plugin/ai/sensitive/word/SensitiveWordPluginConfiguration.java create mode 100644 shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/resources/META-INF/spring.factories create mode 100644 shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/resources/META-INF/spring.provides create mode 100644 shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports create mode 100644 shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/springboot/starter/plugin/ai/sensitive/word/SensitiveWordPluginConfigurationTest.java diff --git a/shenyu-bootstrap/pom.xml b/shenyu-bootstrap/pom.xml index 4a0b181b6b90..caabd7d17f15 100644 --- a/shenyu-bootstrap/pom.xml +++ b/shenyu-bootstrap/pom.xml @@ -263,6 +263,14 @@ + + + org.apache.shenyu + shenyu-spring-boot-starter-plugin-ai-sensitive-word + ${project.version} + + + org.apache.shenyu diff --git a/shenyu-common/src/main/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandle.java b/shenyu-common/src/main/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandle.java new file mode 100644 index 000000000000..18cb1255f573 --- /dev/null +++ b/shenyu-common/src/main/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandle.java @@ -0,0 +1,86 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.shenyu.common.dto.convert.rule; + +/** + * The sensitive word rule handle, it tells the plugin where the dictionary lives and how long a + * loaded dictionary may be reused. + */ +public class SensitiveWordHandle { + + /** + * The default redis key holding the sensitive word set. + */ + public static final String DEFAULT_REDIS_KEY = "shenyu:sensitive:words"; + + /** + * The redis key of the sensitive word set, every rule may point to its own dictionary. + */ + private String redisKey = DEFAULT_REDIS_KEY; + + /** + * How long, in seconds, a loaded dictionary is reused before it is read from redis again. + * Zero or a negative value reads the dictionary on every request. + */ + private long refreshIntervalSeconds = 300L; + + /** + * get redis key. + * + * @return redis key + */ + public String getRedisKey() { + return redisKey; + } + + /** + * set redis key. + * + * @param redisKey redis key + */ + public void setRedisKey(final String redisKey) { + this.redisKey = redisKey; + } + + /** + * get refresh interval seconds. + * + * @return refresh interval in seconds + */ + public long getRefreshIntervalSeconds() { + return refreshIntervalSeconds; + } + + /** + * set refresh interval seconds. + * + * @param refreshIntervalSeconds refresh interval in seconds + */ + public void setRefreshIntervalSeconds(final long refreshIntervalSeconds) { + this.refreshIntervalSeconds = refreshIntervalSeconds; + } + + /** + * new default instance. + * + * @return the default handle + */ + public static SensitiveWordHandle newDefaultInstance() { + return new SensitiveWordHandle(); + } +} diff --git a/shenyu-common/src/main/java/org/apache/shenyu/common/enums/PluginEnum.java b/shenyu-common/src/main/java/org/apache/shenyu/common/enums/PluginEnum.java index 8501238fa405..0ad0f12febc3 100644 --- a/shenyu-common/src/main/java/org/apache/shenyu/common/enums/PluginEnum.java +++ b/shenyu-common/src/main/java/org/apache/shenyu/common/enums/PluginEnum.java @@ -287,6 +287,11 @@ public enum PluginEnum { */ AI_TOKEN_LIMITER(171, 0, "aiTokenLimiter"), + /** + * Sensitive-word plugin enum. + */ + SENSITIVE_WORD(197, 0, "sensitiveWord"), + /** * Mcp-server plugin enum. */ diff --git a/shenyu-common/src/test/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandleTest.java b/shenyu-common/src/test/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandleTest.java new file mode 100644 index 000000000000..9c007835b3b7 --- /dev/null +++ b/shenyu-common/src/test/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandleTest.java @@ -0,0 +1,44 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.shenyu.common.dto.convert.rule; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +/** + * Test cases for {@link SensitiveWordHandle}. + */ +public final class SensitiveWordHandleTest { + + @Test + public void testDefaultInstance() { + SensitiveWordHandle handle = SensitiveWordHandle.newDefaultInstance(); + assertEquals(SensitiveWordHandle.DEFAULT_REDIS_KEY, handle.getRedisKey()); + assertEquals(300L, handle.getRefreshIntervalSeconds()); + } + + @Test + public void testSetter() { + SensitiveWordHandle handle = new SensitiveWordHandle(); + handle.setRedisKey("custom:sensitive:words"); + handle.setRefreshIntervalSeconds(30L); + assertEquals("custom:sensitive:words", handle.getRedisKey()); + assertEquals(30L, handle.getRefreshIntervalSeconds()); + } +} diff --git a/shenyu-plugin/shenyu-plugin-ai/pom.xml b/shenyu-plugin/shenyu-plugin-ai/pom.xml index 22f032a8793c..7f910e45a617 100644 --- a/shenyu-plugin/shenyu-plugin-ai/pom.xml +++ b/shenyu-plugin/shenyu-plugin-ai/pom.xml @@ -34,6 +34,7 @@ shenyu-plugin-ai-request-transformer shenyu-plugin-ai-proxy shenyu-plugin-ai-response-transformer + shenyu-plugin-ai-sensitive-word diff --git a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/README.md b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/README.md new file mode 100644 index 000000000000..a5120053b010 --- /dev/null +++ b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/README.md @@ -0,0 +1,60 @@ +# shenyu-plugin-ai-sensitive-word + +The sensitive word plugin rejects a request whose body contains a word of the dictionary +configured for the matched rule. The dictionary lives in a **redis set**, so it can be maintained +by the operations team without redeploying shenyu. + +It is aimed at the content compliance of an AI gateway, where a prompt must not reach the model +with forbidden content, but it only inspects the request body, so it also applies to plain HTTP +routes. + +## How it works + +- The dictionary is read with `SMEMBERS ` and compiled into an + [Aho-Corasick](../main/java/org/apache/shenyu/plugin/ai/sensitive/word/ac/AhoCorasick.java) + automaton, which reports **every** matching word, including nested and overlapping ones, for + example both `中国` and `中国银行` for the text `中国银行`. +- The body is read with the shared `ServerWebExchangeUtils#rewriteRequestBody`, the whole path is + reactive and the automaton is compiled on a bounded elastic thread, so the gateway event loop is + never blocked. +- A compiled dictionary is reused for `refreshIntervalSeconds` and then read from redis again, so + a dictionary update takes effect within that interval. Updating the rule in the admin console + drops the cached dictionary immediately. +- If redis is unreachable the request is **passed through** (fail open, a warning is logged): a + broken dictionary must not take the traffic down. + +## Configuration + +Plugin level (`config` of the plugin, the redis client used to read the dictionaries): + +```json +{ + "url": "127.0.0.1:6379", + "password": "", + "database": 0, + "mode": "standalone", + "maxIdle": 8, + "minIdle": 0, + "maxActive": 8 +} +``` + +Rule level (`handle` of the rule): + +| field | type | default | description | +| --- | --- | --- | --- | +| `redisKey` | string | `shenyu:sensitive:words` | the redis set holding the dictionary of this rule | +| `refreshIntervalSeconds` | long | `300` | how long a compiled dictionary is reused, `0` reads it on every request | + +## Dictionary format + +A redis set of words, one word per member, for example: + +``` +SADD shenyu:sensitive:words "bad word 1" +SADD shenyu:sensitive:words "bad word 2" +``` + +The word list itself is **not** part of this repository: every deployment is expected to provide +its own dictionary, because the content of such a list depends on the country, the business and +the compliance rules that apply to it. diff --git a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/pom.xml b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/pom.xml new file mode 100644 index 000000000000..782c0b24d294 --- /dev/null +++ b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/pom.xml @@ -0,0 +1,41 @@ + + + + + + org.apache.shenyu + shenyu-plugin-ai + 2.7.2-SNAPSHOT + + 4.0.0 + shenyu-plugin-ai-sensitive-word + + + + org.apache.shenyu + shenyu-plugin-base + ${project.version} + + + org.apache.shenyu + shenyu-infra-redis + ${project.version} + + + + diff --git a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPlugin.java b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPlugin.java new file mode 100644 index 000000000000..0b6e3fce81cd --- /dev/null +++ b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPlugin.java @@ -0,0 +1,148 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.shenyu.plugin.ai.sensitive.word; + +import org.apache.shenyu.common.dto.RuleData; +import org.apache.shenyu.common.dto.SelectorData; +import org.apache.shenyu.common.dto.convert.rule.SensitiveWordHandle; +import org.apache.shenyu.common.enums.PluginEnum; +import org.apache.shenyu.plugin.api.ShenyuPluginChain; +import org.apache.shenyu.plugin.api.exception.ResponsiveException; +import org.apache.shenyu.plugin.api.utils.WebFluxResultUtils; +import org.apache.shenyu.plugin.base.AbstractShenyuPlugin; +import org.apache.shenyu.plugin.base.utils.CacheKeyUtils; +import org.apache.shenyu.plugin.base.utils.ServerWebExchangeUtils; +import org.apache.shenyu.plugin.ai.sensitive.word.ac.AhoCorasick; +import org.apache.shenyu.plugin.ai.sensitive.word.handler.SensitiveWordPluginDataHandler; +import org.apache.shenyu.plugin.ai.sensitive.word.handler.SensitiveWordPluginDataHandler.CachedDictionary; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.data.redis.core.ReactiveRedisTemplate; +import org.springframework.http.codec.HttpMessageReader; +import org.springframework.web.server.ServerWebExchange; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; + +import java.util.List; +import java.util.Objects; + +/** + * The sensitive word plugin, it rejects the request when its body contains a word of the dictionary + * configured for the matched rule. + * + *

The dictionary is read from a redis set, so it can be maintained outside of shenyu. A loaded + * dictionary is reused for the refresh interval of the rule handle, see {@link SensitiveWordHandle}. + */ +public class SensitiveWordPlugin extends AbstractShenyuPlugin { + + private static final Logger LOG = LoggerFactory.getLogger(SensitiveWordPlugin.class); + + /** + * The error code returned to the client when the request contains sensitive words. + */ + private static final int SENSITIVE_WORD_CODE = 1500; + + private final List> readers; + + public SensitiveWordPlugin(final List> readers) { + this.readers = readers; + } + + @Override + protected Mono doExecute(final ServerWebExchange exchange, + final ShenyuPluginChain chain, + final SelectorData selector, + final RuleData rule) { + SensitiveWordHandle handle = SensitiveWordPluginDataHandler.CACHED_HANDLE.get() + .obtainHandle(CacheKeyUtils.INST.getKey(rule)); + if (Objects.isNull(handle)) { + return chain.execute(exchange); + } + ReactiveRedisTemplate redisTemplate = SensitiveWordPluginDataHandler.REDIS_TEMPLATES.get() + .obtainHandle(SensitiveWordPluginDataHandler.PLUGIN_NAME); + if (Objects.isNull(redisTemplate)) { + LOG.warn("sensitive word plugin: the redis template is not initialized, skip the sensitive word check"); + return chain.execute(exchange); + } + return ServerWebExchangeUtils.rewriteRequestBody(exchange, readers, + body -> check(exchange, redisTemplate, handle, body)) + .flatMap(chain::execute) + .onErrorResume(error -> { + if (error instanceof ResponsiveException) { + return WebFluxResultUtils.failedResult((ResponsiveException) error); + } + return Mono.error(error); + }); + } + + private Mono check(final ServerWebExchange exchange, + final ReactiveRedisTemplate redisTemplate, + final SensitiveWordHandle handle, + final String body) { + return dictionary(redisTemplate, handle) + .map(automaton -> automaton.search(body)) + .flatMap(matches -> { + if (matches.isEmpty()) { + return Mono.just(body); + } + return Mono.error(new ResponsiveException(SENSITIVE_WORD_CODE, + String.format("The request contains sensitive words: %s", matches), exchange)); + }); + } + + /** + * Get the dictionary of the rule, it is read from redis when the cached one is missing or + * expired. The request is never blocked by the dictionary: a dictionary that cannot be read + * is treated as empty, so that a broken redis never stops the traffic. + * + * @param redisTemplate the redis template + * @param handle the rule handle + * @return the automaton of the dictionary + */ + private Mono dictionary(final ReactiveRedisTemplate redisTemplate, + final SensitiveWordHandle handle) { + String redisKey = Objects.isNull(handle.getRedisKey()) + ? SensitiveWordHandle.DEFAULT_REDIS_KEY : handle.getRedisKey(); + CachedDictionary cached = SensitiveWordPluginDataHandler.DICTIONARIES.get().obtainHandle(redisKey); + if (Objects.nonNull(cached) && !cached.isExpired(handle.getRefreshIntervalSeconds())) { + return Mono.just(cached.getAutomaton()); + } + return redisTemplate.opsForSet() + .members(redisKey) + .collectList() + // building the automaton is cpu bound, keep it away from the event loop + .map(AhoCorasick::of) + .subscribeOn(Schedulers.boundedElastic()) + .doOnNext(automaton -> SensitiveWordPluginDataHandler.DICTIONARIES.get() + .cachedHandle(redisKey, new CachedDictionary(automaton))) + .onErrorResume(error -> { + LOG.error("sensitive word plugin: cannot read the dictionary from redis key {}", redisKey, error); + return Mono.just(Objects.isNull(cached) ? AhoCorasick.empty() : cached.getAutomaton()); + }); + } + + @Override + public String named() { + return PluginEnum.SENSITIVE_WORD.getName(); + } + + @Override + public int getOrder() { + return PluginEnum.SENSITIVE_WORD.getCode(); + } +} diff --git a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/ac/AhoCorasick.java b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/ac/AhoCorasick.java new file mode 100644 index 000000000000..3c238ed78c10 --- /dev/null +++ b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/ac/AhoCorasick.java @@ -0,0 +1,150 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.shenyu.plugin.ai.sensitive.word.ac; + +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashSet; +import java.util.LinkedList; +import java.util.Map; +import java.util.Objects; +import java.util.Queue; +import java.util.Set; + +/** + * The Aho-Corasick multi pattern matching automaton. + * + *

An instance is built once for a dictionary and is immutable afterwards, so it can be shared + * by concurrent requests. The failure links are built with a breadth first traversal of the trie, + * which makes scanning a text linear in its length (times the length of the longest failure path). + * + *

Every word of the dictionary contained in the scanned text is reported, including nested and + * overlapping ones, for example both {@code 中国} and {@code 中国银行} are reported for the text + * {@code 中国银行}. + */ +public final class AhoCorasick { + + private static final AhoCorasick EMPTY_DICTIONARY = new AhoCorasick(); + + private final TrieNode root = new TrieNode(); + + private AhoCorasick() { + } + + /** + * Build an automaton for the given dictionary. Blank entries are ignored. + * + * @param words the dictionary words + * @return the automaton + */ + public static AhoCorasick of(final Collection words) { + if (Objects.isNull(words) || words.isEmpty()) { + return EMPTY_DICTIONARY; + } + AhoCorasick automaton = new AhoCorasick(); + automaton.insertAll(words); + automaton.buildFailureLinks(); + return automaton; + } + + /** + * An automaton without any word, it never matches. Useful as a fail open fallback when the + * dictionary cannot be loaded. + * + * @return an empty automaton + */ + public static AhoCorasick empty() { + return EMPTY_DICTIONARY; + } + + /** + * Find every dictionary word contained in the given text. + * + * @param text the text to scan + * @return the matched words, in the order they were found, each word is reported once + */ + public Set search(final String text) { + if (Objects.isNull(text) || text.isEmpty()) { + return Collections.emptySet(); + } + Set matches = new LinkedHashSet<>(); + TrieNode current = root; + for (int i = 0; i < text.length(); i++) { + char c = text.charAt(i); + while (current != root && !current.children.containsKey(c)) { + current = current.fail; + } + TrieNode next = current.children.get(c); + current = Objects.isNull(next) ? root : next; + // Walk the whole failure chain, otherwise nested and overlapping words are missed. + for (TrieNode node = current; node != root; node = node.fail) { + if (Objects.nonNull(node.word)) { + matches.add(node.word); + } + } + } + return matches; + } + + private void insertAll(final Collection words) { + for (String word : words) { + if (Objects.isNull(word) || word.trim().isEmpty()) { + continue; + } + insert(word.trim()); + } + } + + private void insert(final String word) { + TrieNode node = root; + for (char c : word.toCharArray()) { + node = node.children.computeIfAbsent(c, key -> new TrieNode()); + } + node.word = word; + } + + private void buildFailureLinks() { + Queue queue = new LinkedList<>(); + queue.add(root); + while (!queue.isEmpty()) { + TrieNode current = queue.poll(); + for (Map.Entry entry : current.children.entrySet()) { + TrieNode child = entry.getValue(); + TrieNode fail = current.fail; + while (Objects.nonNull(fail) && !fail.children.containsKey(entry.getKey())) { + fail = fail.fail; + } + child.fail = Objects.nonNull(fail) ? fail.children.get(entry.getKey()) : root; + queue.add(child); + } + } + } + + /** + * A trie node, {@code word} is not null only on the node that ends a dictionary word. + */ + private static final class TrieNode { + + private final Map children = new HashMap<>(); + + private TrieNode fail; + + private String word; + } +} diff --git a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandler.java b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandler.java new file mode 100644 index 000000000000..a80430849fdd --- /dev/null +++ b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandler.java @@ -0,0 +1,161 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.shenyu.plugin.ai.sensitive.word.handler; + +import org.apache.shenyu.common.dto.PluginData; +import org.apache.shenyu.common.dto.RuleData; +import org.apache.shenyu.common.dto.convert.rule.SensitiveWordHandle; +import org.apache.shenyu.common.enums.PluginEnum; +import org.apache.shenyu.common.utils.GsonUtils; +import org.apache.shenyu.infra.redis.RedisConfigProperties; +import org.apache.shenyu.infra.redis.RedisConnectionFactory; +import org.apache.shenyu.infra.redis.ShenyuReactiveRedisTemplate; +import org.apache.shenyu.infra.redis.serializer.ShenyuRedisSerializationContext; +import org.apache.shenyu.plugin.base.cache.CommonHandleCache; +import org.apache.shenyu.plugin.base.handler.PluginDataHandler; +import org.apache.shenyu.plugin.base.utils.BeanHolder; +import org.apache.shenyu.plugin.base.utils.CacheKeyUtils; +import org.apache.shenyu.plugin.ai.sensitive.word.ac.AhoCorasick; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.data.redis.core.ReactiveRedisTemplate; +import org.springframework.util.StringUtils; + +import java.util.Objects; +import java.util.Optional; +import java.util.function.Supplier; + +/** + * The sensitive word plugin data handler. + */ +public class SensitiveWordPluginDataHandler implements PluginDataHandler { + + /** + * The plugin name, it is also the cache key of the redis template. + */ + public static final String PLUGIN_NAME = PluginEnum.SENSITIVE_WORD.getName(); + + /** + * Cache the reactive redis template, the key is the plugin name. + */ + public static final Supplier>> REDIS_TEMPLATES = + new BeanHolder<>(CommonHandleCache::new); + + /** + * Cache the redis configuration the template was built from, the key is the plugin name. + */ + public static final Supplier> REDIS_PROPERTIES = + new BeanHolder<>(CommonHandleCache::new); + + /** + * Cache the rule handle, the key is the cache key of the rule. + */ + public static final Supplier> CACHED_HANDLE = + new BeanHolder<>(CommonHandleCache::new); + + /** + * Cache the dictionary automaton, the key is the redis key of the dictionary, so that rules + * pointing to different dictionaries never share the same automaton. + */ + public static final Supplier> DICTIONARIES = + new BeanHolder<>(CommonHandleCache::new); + + private static final Logger LOG = LoggerFactory.getLogger(SensitiveWordPluginDataHandler.class); + + @Override + public void handlerPlugin(final PluginData pluginData) { + if (Objects.isNull(pluginData) || !Boolean.TRUE.equals(pluginData.getEnabled())) { + return; + } + RedisConfigProperties redisConfig = GsonUtils.getInstance() + .fromJson(pluginData.getConfig(), RedisConfigProperties.class); + if (Objects.isNull(redisConfig) || !StringUtils.hasText(redisConfig.getUrl())) { + LOG.warn("sensitive word plugin: the redis configuration is missing, skip the redis initialization"); + return; + } + RedisConfigProperties cachedProperties = REDIS_PROPERTIES.get().obtainHandle(PLUGIN_NAME); + if (Objects.isNull(REDIS_TEMPLATES.get().obtainHandle(PLUGIN_NAME)) || !redisConfig.equals(cachedProperties)) { + RedisConnectionFactory connectionFactory = new RedisConnectionFactory(redisConfig); + ReactiveRedisTemplate redisTemplate = new ShenyuReactiveRedisTemplate<>( + connectionFactory.getLettuceConnectionFactory(), + ShenyuRedisSerializationContext.stringSerializationContext()); + REDIS_TEMPLATES.get().cachedHandle(PLUGIN_NAME, redisTemplate); + REDIS_PROPERTIES.get().cachedHandle(PLUGIN_NAME, redisConfig); + LOG.info("sensitive word plugin: cached the reactive redis template"); + } + } + + @Override + public void handlerRule(final RuleData ruleData) { + Optional.ofNullable(ruleData.getHandle()).ifPresent(json -> { + SensitiveWordHandle handle = GsonUtils.getInstance().fromJson(json, SensitiveWordHandle.class); + if (Objects.isNull(handle)) { + handle = SensitiveWordHandle.newDefaultInstance(); + } + CACHED_HANDLE.get().cachedHandle(CacheKeyUtils.INST.getKey(ruleData), handle); + // The rule changed, drop the cached dictionary so that it is read from redis again. + DICTIONARIES.get().removeHandle(handle.getRedisKey()); + }); + } + + @Override + public void removeRule(final RuleData ruleData) { + CACHED_HANDLE.get().removeHandle(CacheKeyUtils.INST.getKey(ruleData)); + } + + @Override + public String pluginNamed() { + return PLUGIN_NAME; + } + + /** + * A dictionary automaton together with the time it was loaded, used to refresh a dictionary + * whose refresh interval elapsed. + */ + public static final class CachedDictionary { + + private final AhoCorasick automaton; + + private final long loadTime; + + public CachedDictionary(final AhoCorasick automaton) { + this.automaton = automaton; + this.loadTime = System.currentTimeMillis(); + } + + /** + * Whether this dictionary is older than the refresh interval of the rule. + * + * @param refreshIntervalSeconds the refresh interval in seconds + * @return true when the dictionary must be read from redis again + */ + public boolean isExpired(final long refreshIntervalSeconds) { + return refreshIntervalSeconds <= 0L + || System.currentTimeMillis() - loadTime >= refreshIntervalSeconds * 1000L; + } + + /** + * get the automaton. + * + * @return the automaton + */ + public AhoCorasick getAutomaton() { + return automaton; + } + } +} diff --git a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPluginTest.java b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPluginTest.java new file mode 100644 index 000000000000..911c7b0880fb --- /dev/null +++ b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPluginTest.java @@ -0,0 +1,167 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.shenyu.plugin.ai.sensitive.word; + +import org.apache.shenyu.common.dto.RuleData; +import org.apache.shenyu.common.dto.convert.rule.SensitiveWordHandle; +import org.apache.shenyu.common.enums.PluginEnum; +import org.apache.shenyu.plugin.ai.sensitive.word.handler.SensitiveWordPluginDataHandler; +import org.apache.shenyu.plugin.api.ShenyuPluginChain; +import org.apache.shenyu.plugin.api.result.DefaultShenyuResult; +import org.apache.shenyu.plugin.api.result.ShenyuResult; +import org.apache.shenyu.plugin.api.utils.SpringBeanUtils; +import org.apache.shenyu.plugin.base.utils.CacheKeyUtils; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.springframework.context.ConfigurableApplicationContext; +import org.springframework.data.redis.core.ReactiveRedisTemplate; +import org.springframework.data.redis.core.ReactiveSetOperations; +import org.springframework.http.MediaType; +import org.springframework.http.codec.ServerCodecConfigurer; +import org.springframework.mock.http.server.reactive.MockServerHttpRequest; +import org.springframework.mock.web.server.MockServerWebExchange; +import org.springframework.web.server.ServerWebExchange; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Test cases for {@link SensitiveWordPlugin}. + */ +public final class SensitiveWordPluginTest { + + private static final String REDIS_KEY = "test:sensitive:words"; + + private SensitiveWordPlugin plugin; + + private ShenyuPluginChain chain; + + private RuleData ruleData; + + @BeforeEach + public void setUp() { + ConfigurableApplicationContext applicationContext = mock(ConfigurableApplicationContext.class); + when(applicationContext.getBean(ShenyuResult.class)).thenReturn(new DefaultShenyuResult()); + SpringBeanUtils.getInstance().setApplicationContext(applicationContext); + + plugin = new SensitiveWordPlugin(ServerCodecConfigurer.create().getReaders()); + chain = mock(ShenyuPluginChain.class); + when(chain.execute(any(ServerWebExchange.class))).thenReturn(Mono.empty()); + ruleData = new RuleData(); + ruleData.setId("rule-1"); + ruleData.setName("rule-1"); + ruleData.setSelectorId("selector-1"); + ruleData.setPluginName(PluginEnum.SENSITIVE_WORD.getName()); + ruleData.setHandle("{\"redisKey\":\"" + REDIS_KEY + "\"}"); + SensitiveWordHandle handle = SensitiveWordHandle.newDefaultInstance(); + handle.setRedisKey(REDIS_KEY); + SensitiveWordPluginDataHandler.CACHED_HANDLE.get() + .cachedHandle(CacheKeyUtils.INST.getKey(ruleData), handle); + } + + @AfterEach + public void tearDown() { + SpringBeanUtils.getInstance().setApplicationContext(null); + SensitiveWordPluginDataHandler.CACHED_HANDLE.get().removeHandle(CacheKeyUtils.INST.getKey(ruleData)); + SensitiveWordPluginDataHandler.REDIS_TEMPLATES.get().removeHandle(SensitiveWordPluginDataHandler.PLUGIN_NAME); + SensitiveWordPluginDataHandler.DICTIONARIES.get().removeHandle(REDIS_KEY); + } + + @Test + public void testNamed() { + assertEquals(PluginEnum.SENSITIVE_WORD.getName(), plugin.named()); + } + + @Test + public void testGetOrder() { + assertEquals(PluginEnum.SENSITIVE_WORD.getCode(), plugin.getOrder()); + } + + @Test + public void testPassThroughWhenNoRuleHandle() { + SensitiveWordPluginDataHandler.CACHED_HANDLE.get().removeHandle(CacheKeyUtils.INST.getKey(ruleData)); + StepVerifier.create(plugin.doExecute(exchange("a clean request"), chain, null, ruleData)).verifyComplete(); + verify(chain).execute(any(ServerWebExchange.class)); + } + + @Test + public void testPassThroughWhenRedisIsNotInitialized() { + StepVerifier.create(plugin.doExecute(exchange("a clean request"), chain, null, ruleData)).verifyComplete(); + verify(chain).execute(any(ServerWebExchange.class)); + } + + @Test + public void testPassThroughWhenTheBodyIsClean() { + mockRedisDictionary("forbidden", "banned"); + StepVerifier.create(plugin.doExecute(exchange("a clean request"), chain, null, ruleData)).verifyComplete(); + verify(chain).execute(any(ServerWebExchange.class)); + } + + @Test + public void testRejectTheRequestContainingASensitiveWord() { + mockRedisDictionary("forbidden", "banned"); + MockServerWebExchange exchange = exchange("this request is banned"); + StepVerifier.create(plugin.doExecute(exchange, chain, null, ruleData)).verifyComplete(); + verify(chain, never()).execute(any(ServerWebExchange.class)); + StepVerifier.create(exchange.getResponse().getBodyAsString()) + .expectNextMatches(body -> body.contains("sensitive words") && body.contains("banned")) + .verifyComplete(); + } + + @Test + public void testPassThroughWhenRedisFails() { + ReactiveRedisTemplate redisTemplate = mock(ReactiveRedisTemplate.class); + ReactiveSetOperations setOperations = mock(ReactiveSetOperations.class); + when(redisTemplate.opsForSet()).thenReturn(setOperations); + when(setOperations.members(REDIS_KEY)).thenReturn(Flux.error(new IllegalStateException("redis is down"))); + cacheRedisTemplate(redisTemplate); + + // fail open: a broken redis must not stop the traffic + StepVerifier.create(plugin.doExecute(exchange("this request is banned"), chain, null, ruleData)).verifyComplete(); + verify(chain).execute(any(ServerWebExchange.class)); + } + + @SuppressWarnings("unchecked") + private void mockRedisDictionary(final String... words) { + ReactiveRedisTemplate redisTemplate = mock(ReactiveRedisTemplate.class); + ReactiveSetOperations setOperations = mock(ReactiveSetOperations.class); + when(redisTemplate.opsForSet()).thenReturn(setOperations); + when(setOperations.members(REDIS_KEY)).thenReturn(Flux.fromArray(words)); + cacheRedisTemplate(redisTemplate); + } + + @SuppressWarnings("unchecked") + private void cacheRedisTemplate(final ReactiveRedisTemplate redisTemplate) { + SensitiveWordPluginDataHandler.REDIS_TEMPLATES.get() + .cachedHandle(SensitiveWordPluginDataHandler.PLUGIN_NAME, redisTemplate); + } + + private MockServerWebExchange exchange(final String body) { + return MockServerWebExchange.from(MockServerHttpRequest.post("/ai/chat") + .contentType(MediaType.TEXT_PLAIN) + .body(body)); + } +} diff --git a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/ac/AhoCorasickTest.java b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/ac/AhoCorasickTest.java new file mode 100644 index 000000000000..ad6e8bcd4a39 --- /dev/null +++ b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/ac/AhoCorasickTest.java @@ -0,0 +1,125 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.shenyu.plugin.ai.sensitive.word.ac; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Test cases for {@link AhoCorasick}. + */ +public final class AhoCorasickTest { + + @Test + public void testMatchSingleWord() { + AhoCorasick automaton = AhoCorasick.of(Collections.singletonList("sensitive")); + assertThat(automaton.search("this is a sensitive word")).containsExactly("sensitive"); + } + + @Test + public void testMatchEveryWordOfTheText() { + AhoCorasick automaton = AhoCorasick.of(Arrays.asList("sensitive", "word", "spam")); + assertThat(automaton.search("a sensitive word and a spam")).containsExactlyInAnyOrder("sensitive", "word", "spam"); + } + + @Test + public void testMatchNestedWords() { + AhoCorasick automaton = AhoCorasick.of(Arrays.asList("abc", "bc", "c")); + assertThat(automaton.search("abc")).containsExactlyInAnyOrder("abc", "bc", "c"); + } + + @Test + public void testMatchOverlappingWords() { + AhoCorasick automaton = AhoCorasick.of(Arrays.asList("中国", "中国银行")); + assertThat(automaton.search("中国银行")).containsExactlyInAnyOrder("中国", "中国银行"); + } + + @Test + public void testMatchSuffixWords() { + AhoCorasick automaton = AhoCorasick.of(Arrays.asList("敏感词", "词")); + assertThat(automaton.search("这里有敏感词")).containsExactlyInAnyOrder("敏感词", "词"); + } + + @Test + public void testFailureLinksAreFollowed() { + AhoCorasick automaton = AhoCorasick.of(Arrays.asList("he", "she", "his", "hers")); + assertThat(automaton.search("ushers")).containsExactlyInAnyOrder("she", "he", "hers"); + } + + @Test + public void testMatchAtTextBoundaries() { + AhoCorasick automaton = AhoCorasick.of(Collections.singletonList("ab")); + assertThat(automaton.search("abxxab")).containsExactly("ab"); + assertThat(automaton.search("ab")).containsExactly("ab"); + assertThat(automaton.search("a")).isEmpty(); + } + + @Test + public void testWordIsReportedOnce() { + AhoCorasick automaton = AhoCorasick.of(Collections.singletonList("spam")); + assertThat(automaton.search("spam spam spam")).containsExactly("spam"); + } + + @Test + public void testMatchCjkText() { + AhoCorasick automaton = AhoCorasick.of(Arrays.asList("违规", "敏感")); + assertThat(automaton.search("这段内容违规而且敏感")).containsExactlyInAnyOrder("违规", "敏感"); + } + + @Test + public void testNoMatch() { + AhoCorasick automaton = AhoCorasick.of(Arrays.asList("sensitive", "spam")); + assertThat(automaton.search("a clean text")).isEmpty(); + } + + @Test + public void testEmptyDictionary() { + assertThat(AhoCorasick.of(Collections.emptyList()).search("anything")).isEmpty(); + assertThat(AhoCorasick.of(null).search("anything")).isEmpty(); + assertThat(AhoCorasick.empty().search("anything")).isEmpty(); + } + + @Test + public void testBlankAndNullWordsAreIgnored() { + AhoCorasick automaton = AhoCorasick.of(Arrays.asList(" ", "", null, " hit ")); + assertThat(automaton.search("a hit")).containsExactly("hit"); + assertThat(automaton.search("nothing")).isEmpty(); + } + + @Test + public void testNullAndEmptyText() { + AhoCorasick automaton = AhoCorasick.of(Collections.singletonList("sensitive")); + assertThat(automaton.search(null)).isEmpty(); + assertThat(automaton.search("")).isEmpty(); + } + + @Test + public void testLongText() { + StringBuilder text = new StringBuilder(); + for (int i = 0; i < 10000; i++) { + text.append("filler"); + } + text.append("sensitive"); + AhoCorasick automaton = AhoCorasick.of(Arrays.asList("sensitive", "filler")); + assertThat(automaton.search(text.toString())).containsExactlyInAnyOrder("filler", "sensitive"); + } +} diff --git a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandlerTest.java b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandlerTest.java new file mode 100644 index 000000000000..a13dd290e77f --- /dev/null +++ b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandlerTest.java @@ -0,0 +1,151 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.shenyu.plugin.ai.sensitive.word.handler; + +import org.apache.shenyu.common.dto.PluginData; +import org.apache.shenyu.common.dto.RuleData; +import org.apache.shenyu.common.dto.convert.rule.SensitiveWordHandle; +import org.apache.shenyu.common.enums.PluginEnum; +import org.apache.shenyu.plugin.base.utils.CacheKeyUtils; +import org.apache.shenyu.plugin.ai.sensitive.word.ac.AhoCorasick; +import org.apache.shenyu.plugin.ai.sensitive.word.handler.SensitiveWordPluginDataHandler.CachedDictionary; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Test cases for {@link SensitiveWordPluginDataHandler}. + */ +public final class SensitiveWordPluginDataHandlerTest { + + private SensitiveWordPluginDataHandler handler; + + @BeforeEach + public void setUp() { + handler = new SensitiveWordPluginDataHandler(); + } + + @AfterEach + public void tearDown() { + SensitiveWordPluginDataHandler.CACHED_HANDLE.get() + .removeHandle(CacheKeyUtils.INST.getKey("selector-1", "rule-1")); + SensitiveWordPluginDataHandler.DICTIONARIES.get().removeHandle("my:sensitive:words"); + SensitiveWordPluginDataHandler.REDIS_TEMPLATES.get().removeHandle(SensitiveWordPluginDataHandler.PLUGIN_NAME); + SensitiveWordPluginDataHandler.REDIS_PROPERTIES.get().removeHandle(SensitiveWordPluginDataHandler.PLUGIN_NAME); + } + + @Test + public void testPluginNamed() { + assertEquals(PluginEnum.SENSITIVE_WORD.getName(), handler.pluginNamed()); + } + + @Test + public void testHandlerRuleCachesTheHandle() { + RuleData ruleData = ruleData("{\"redisKey\":\"my:sensitive:words\",\"refreshIntervalSeconds\":10}"); + handler.handlerRule(ruleData); + SensitiveWordHandle handle = SensitiveWordPluginDataHandler.CACHED_HANDLE.get() + .obtainHandle(CacheKeyUtils.INST.getKey(ruleData)); + assertNotNull(handle); + assertEquals("my:sensitive:words", handle.getRedisKey()); + assertEquals(10L, handle.getRefreshIntervalSeconds()); + } + + @Test + public void testHandlerRuleWithoutRedisKeyFallsBackToTheDefaultOne() { + RuleData ruleData = ruleData("{}"); + handler.handlerRule(ruleData); + SensitiveWordHandle handle = SensitiveWordPluginDataHandler.CACHED_HANDLE.get() + .obtainHandle(CacheKeyUtils.INST.getKey(ruleData)); + assertNotNull(handle); + assertEquals(SensitiveWordHandle.DEFAULT_REDIS_KEY, handle.getRedisKey()); + assertEquals(300L, handle.getRefreshIntervalSeconds()); + } + + @Test + public void testHandlerRuleInvalidatesTheCachedDictionary() { + String redisKey = "my:sensitive:words"; + SensitiveWordPluginDataHandler.DICTIONARIES.get() + .cachedHandle(redisKey, new CachedDictionary(AhoCorasick.empty())); + handler.handlerRule(ruleData("{\"redisKey\":\"my:sensitive:words\"}")); + assertNull(SensitiveWordPluginDataHandler.DICTIONARIES.get().obtainHandle(redisKey)); + } + + @Test + public void testRemoveRuleRemovesTheHandle() { + RuleData ruleData = ruleData("{\"redisKey\":\"my:sensitive:words\"}"); + handler.handlerRule(ruleData); + assertNotNull(SensitiveWordPluginDataHandler.CACHED_HANDLE.get().obtainHandle(CacheKeyUtils.INST.getKey(ruleData))); + handler.removeRule(ruleData); + assertNull(SensitiveWordPluginDataHandler.CACHED_HANDLE.get().obtainHandle(CacheKeyUtils.INST.getKey(ruleData))); + } + + @Test + public void testHandlerPluginWithoutRedisConfiguration() { + PluginData pluginData = new PluginData(); + pluginData.setEnabled(true); + pluginData.setConfig("{}"); + handler.handlerPlugin(pluginData); + assertNull(SensitiveWordPluginDataHandler.REDIS_TEMPLATES.get() + .obtainHandle(SensitiveWordPluginDataHandler.PLUGIN_NAME)); + } + + @Test + public void testHandlerPluginDisabled() { + PluginData pluginData = new PluginData(); + pluginData.setEnabled(false); + pluginData.setConfig("{\"url\":\"127.0.0.1:6379\"}"); + handler.handlerPlugin(pluginData); + assertNull(SensitiveWordPluginDataHandler.REDIS_TEMPLATES.get() + .obtainHandle(SensitiveWordPluginDataHandler.PLUGIN_NAME)); + } + + @Test + public void testHandlerPluginCachesTheRedisTemplate() { + PluginData pluginData = new PluginData(); + pluginData.setEnabled(true); + pluginData.setConfig("{\"url\":\"127.0.0.1:6379\",\"database\":0,\"maxIdle\":8,\"minIdle\":0,\"maxActive\":8}"); + handler.handlerPlugin(pluginData); + assertNotNull(SensitiveWordPluginDataHandler.REDIS_TEMPLATES.get() + .obtainHandle(SensitiveWordPluginDataHandler.PLUGIN_NAME)); + assertNotNull(SensitiveWordPluginDataHandler.REDIS_PROPERTIES.get() + .obtainHandle(SensitiveWordPluginDataHandler.PLUGIN_NAME)); + } + + @Test + public void testCachedDictionaryExpires() { + CachedDictionary dictionary = new CachedDictionary(AhoCorasick.empty()); + assertTrue(dictionary.isExpired(0L)); + assertTrue(dictionary.isExpired(-1L)); + assertTrue(!dictionary.isExpired(300L)); + } + + private RuleData ruleData(final String handle) { + RuleData ruleData = new RuleData(); + ruleData.setId("rule-1"); + ruleData.setName("rule-1"); + ruleData.setSelectorId("selector-1"); + ruleData.setPluginName(PluginEnum.SENSITIVE_WORD.getName()); + ruleData.setHandle(handle); + return ruleData; + } +} diff --git a/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/pom.xml b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/pom.xml index e5192d9829c4..d096319ffb47 100644 --- a/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/pom.xml +++ b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/pom.xml @@ -78,5 +78,6 @@ shenyu-spring-boot-starter-plugin-ai-request-transformer shenyu-spring-boot-starter-plugin-ai-response-transformer shenyu-spring-boot-starter-plugin-mcp-server + shenyu-spring-boot-starter-plugin-ai-sensitive-word diff --git a/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/pom.xml b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/pom.xml new file mode 100644 index 000000000000..0da5b0acc640 --- /dev/null +++ b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/pom.xml @@ -0,0 +1,35 @@ + + + + + + org.apache.shenyu + shenyu-spring-boot-starter-plugin + 2.7.2-SNAPSHOT + + 4.0.0 + shenyu-spring-boot-starter-plugin-ai-sensitive-word + + + + org.apache.shenyu + shenyu-plugin-ai-sensitive-word + ${project.version} + + + diff --git a/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/springboot/starter/plugin/ai/sensitive/word/SensitiveWordPluginConfiguration.java b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/springboot/starter/plugin/ai/sensitive/word/SensitiveWordPluginConfiguration.java new file mode 100644 index 000000000000..e9a4d06a3744 --- /dev/null +++ b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/springboot/starter/plugin/ai/sensitive/word/SensitiveWordPluginConfiguration.java @@ -0,0 +1,57 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.shenyu.springboot.starter.plugin.ai.sensitive.word; + +import org.apache.shenyu.plugin.api.ShenyuPlugin; +import org.apache.shenyu.plugin.base.handler.PluginDataHandler; +import org.apache.shenyu.plugin.ai.sensitive.word.SensitiveWordPlugin; +import org.apache.shenyu.plugin.ai.sensitive.word.handler.SensitiveWordPluginDataHandler; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.http.codec.ServerCodecConfigurer; + +/** + * The sensitive word plugin configuration. + */ +@Configuration +@ConditionalOnProperty(value = "shenyu.plugins.sensitive-word.enabled", havingValue = "true", matchIfMissing = true) +public class SensitiveWordPluginConfiguration { + + /** + * the sensitive word plugin. + * + * @param configurer the server codec configurer + * @return the sensitive word plugin + */ + @Bean + public ShenyuPlugin sensitiveWordPlugin(final ServerCodecConfigurer configurer) { + return new SensitiveWordPlugin(configurer.getReaders()); + } + + /** + * the sensitive word plugin data handler, it caches the redis client, the rule handles and the + * dictionaries. + * + * @return the sensitive word plugin data handler + */ + @Bean + public PluginDataHandler sensitiveWordPluginDataHandler() { + return new SensitiveWordPluginDataHandler(); + } +} diff --git a/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/resources/META-INF/spring.factories b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/resources/META-INF/spring.factories new file mode 100644 index 000000000000..1349512313a4 --- /dev/null +++ b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/resources/META-INF/spring.factories @@ -0,0 +1,19 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +org.springframework.boot.autoconfigure.EnableAutoConfiguration=\ +org.apache.shenyu.springboot.starter.plugin.ai.sensitive.word.SensitiveWordPluginConfiguration diff --git a/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/resources/META-INF/spring.provides b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/resources/META-INF/spring.provides new file mode 100644 index 000000000000..19277d9808a4 --- /dev/null +++ b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/resources/META-INF/spring.provides @@ -0,0 +1,18 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +provides: shenyu-spring-boot-starter-plugin-ai-sensitive-word diff --git a/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports new file mode 100644 index 000000000000..e802915ebfed --- /dev/null +++ b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports @@ -0,0 +1,18 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +org.apache.shenyu.springboot.starter.plugin.ai.sensitive.word.SensitiveWordPluginConfiguration diff --git a/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/springboot/starter/plugin/ai/sensitive/word/SensitiveWordPluginConfigurationTest.java b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/springboot/starter/plugin/ai/sensitive/word/SensitiveWordPluginConfigurationTest.java new file mode 100644 index 000000000000..16b4302f22f5 --- /dev/null +++ b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/springboot/starter/plugin/ai/sensitive/word/SensitiveWordPluginConfigurationTest.java @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.shenyu.springboot.starter.plugin.ai.sensitive.word; + +import org.apache.shenyu.common.enums.PluginEnum; +import org.apache.shenyu.plugin.api.ShenyuPlugin; +import org.apache.shenyu.plugin.base.handler.PluginDataHandler; +import org.junit.jupiter.api.Test; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Configuration; +import org.springframework.http.codec.support.DefaultServerCodecConfigurer; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +/** + * Test case for {@link SensitiveWordPluginConfiguration}. + */ +@Configuration +@EnableConfigurationProperties +public class SensitiveWordPluginConfigurationTest { + + @Test + public void testSensitiveWordPlugin() { + new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(SensitiveWordPluginConfiguration.class, DefaultServerCodecConfigurer.class)) + .withBean(SensitiveWordPluginConfigurationTest.class) + .withPropertyValues("debug=true") + .run(context -> { + ShenyuPlugin plugin = context.getBean("sensitiveWordPlugin", ShenyuPlugin.class); + assertNotNull(plugin); + assertThat(plugin.named()).isEqualTo(PluginEnum.SENSITIVE_WORD.getName()); + }); + } + + @Test + public void testSensitiveWordPluginDataHandler() { + new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(SensitiveWordPluginConfiguration.class, DefaultServerCodecConfigurer.class)) + .withBean(SensitiveWordPluginConfigurationTest.class) + .withPropertyValues("debug=true") + .run(context -> { + PluginDataHandler handler = context.getBean("sensitiveWordPluginDataHandler", PluginDataHandler.class); + assertNotNull(handler); + assertThat(handler.pluginNamed()).isEqualTo(PluginEnum.SENSITIVE_WORD.getName()); + }); + } +} From 06f2b610973d7b46a064d8d3220d11f380be4efb Mon Sep 17 00:00:00 2001 From: HY-love-sleep <583699747@qq.com> Date: Mon, 21 Sep 2026 15:08:48 +0800 Subject: [PATCH 2/2] fix: address the review of the sensitive word plugin - hop off the netty event loop on both dictionary paths: the cached path emitted on the subscribing thread, so the body scan ran on the event loop of the request being filtered - precompute the words matched at every trie node while building the failure links, which makes the scan linear in the length of the text instead of walking the failure chain for every character - never echo the matched words back to the caller: the client gets a generic rejection message, the matches are logged server side only, and the rejection code follows the waf plugin (403) - add the rule level failClosed option: when the dictionary cannot be read and no cached dictionary is available, such a rule rejects the request instead of passing it through - release the cached redis template in removePlugin - tests for the failClosed paths, the stale dictionary fallback and removePlugin --- .../dto/convert/rule/SensitiveWordHandle.java | 25 +++++++ .../convert/rule/SensitiveWordHandleTest.java | 10 +++ .../sensitive/word/SensitiveWordPlugin.java | 66 +++++++++++++------ .../ai/sensitive/word/ac/AhoCorasick.java | 25 +++++-- .../SensitiveWordPluginDataHandler.java | 7 ++ .../word/SensitiveWordPluginTest.java | 53 ++++++++++++++- .../SensitiveWordPluginDataHandlerTest.java | 15 +++++ 7 files changed, 175 insertions(+), 26 deletions(-) diff --git a/shenyu-common/src/main/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandle.java b/shenyu-common/src/main/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandle.java index 18cb1255f573..fd51e9eed39e 100644 --- a/shenyu-common/src/main/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandle.java +++ b/shenyu-common/src/main/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandle.java @@ -39,6 +39,13 @@ public class SensitiveWordHandle { */ private long refreshIntervalSeconds = 300L; + /** + * Whether the request must be rejected when the dictionary is unavailable. It defaults to + * false, which means such a request is passed through: a broken redis must not take the + * traffic down. Deployments with a hard compliance requirement can opt into blocking. + */ + private boolean failClosed; + /** * get redis key. * @@ -75,6 +82,24 @@ public void setRefreshIntervalSeconds(final long refreshIntervalSeconds) { this.refreshIntervalSeconds = refreshIntervalSeconds; } + /** + * whether the request must be rejected when the dictionary is unavailable. + * + * @return true when the request must be rejected + */ + public boolean isFailClosed() { + return failClosed; + } + + /** + * set whether the request must be rejected when the dictionary is unavailable. + * + * @param failClosed true to reject the request + */ + public void setFailClosed(final boolean failClosed) { + this.failClosed = failClosed; + } + /** * new default instance. * diff --git a/shenyu-common/src/test/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandleTest.java b/shenyu-common/src/test/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandleTest.java index 9c007835b3b7..ddc61e24db16 100644 --- a/shenyu-common/src/test/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandleTest.java +++ b/shenyu-common/src/test/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandleTest.java @@ -20,6 +20,8 @@ import org.junit.jupiter.api.Test; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; /** * Test cases for {@link SensitiveWordHandle}. @@ -33,6 +35,14 @@ public void testDefaultInstance() { assertEquals(300L, handle.getRefreshIntervalSeconds()); } + @Test + public void testFailClosed() { + SensitiveWordHandle handle = new SensitiveWordHandle(); + assertFalse(handle.isFailClosed()); + handle.setFailClosed(true); + assertTrue(handle.isFailClosed()); + } + @Test public void testSetter() { SensitiveWordHandle handle = new SensitiveWordHandle(); diff --git a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPlugin.java b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPlugin.java index 0b6e3fce81cd..6ce6f14335a2 100644 --- a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPlugin.java +++ b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPlugin.java @@ -33,6 +33,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.data.redis.core.ReactiveRedisTemplate; +import org.springframework.http.HttpStatus; import org.springframework.http.codec.HttpMessageReader; import org.springframework.web.server.ServerWebExchange; import reactor.core.publisher.Mono; @@ -53,9 +54,17 @@ public class SensitiveWordPlugin extends AbstractShenyuPlugin { private static final Logger LOG = LoggerFactory.getLogger(SensitiveWordPlugin.class); /** - * The error code returned to the client when the request contains sensitive words. + * The error code returned to a client whose request was rejected. It follows the waf plugin, + * which rejects with 403 as well. */ - private static final int SENSITIVE_WORD_CODE = 1500; + private static final int SENSITIVE_WORD_CODE = HttpStatus.FORBIDDEN.value(); + + /** + * The message returned to a client whose request was rejected. It never contains the matched + * words: echoing them back would confirm the dictionary to the caller and would re-emit + * forbidden content into the caller side logs. + */ + private static final String REJECT_MESSAGE = "Request rejected: sensitive content detected"; private final List> readers; @@ -94,14 +103,17 @@ private Mono check(final ServerWebExchange exchange, final ReactiveRedisTemplate redisTemplate, final SensitiveWordHandle handle, final String body) { - return dictionary(redisTemplate, handle) + return dictionary(exchange, redisTemplate, handle) .map(automaton -> automaton.search(body)) .flatMap(matches -> { if (matches.isEmpty()) { return Mono.just(body); } - return Mono.error(new ResponsiveException(SENSITIVE_WORD_CODE, - String.format("The request contains sensitive words: %s", matches), exchange)); + // The matches are logged server side only: returning them to the caller would + // turn the gateway into an oracle of the dictionary and would push content that + // was just classified as forbidden into the caller side logs. + LOG.warn("sensitive word plugin: the request was rejected, matched words: {}", matches); + return Mono.error(new ResponsiveException(SENSITIVE_WORD_CODE, REJECT_MESSAGE, exchange)); }); } @@ -114,26 +126,42 @@ private Mono check(final ServerWebExchange exchange, * @param handle the rule handle * @return the automaton of the dictionary */ - private Mono dictionary(final ReactiveRedisTemplate redisTemplate, + private Mono dictionary(final ServerWebExchange exchange, + final ReactiveRedisTemplate redisTemplate, final SensitiveWordHandle handle) { String redisKey = Objects.isNull(handle.getRedisKey()) ? SensitiveWordHandle.DEFAULT_REDIS_KEY : handle.getRedisKey(); CachedDictionary cached = SensitiveWordPluginDataHandler.DICTIONARIES.get().obtainHandle(redisKey); + Mono automaton; if (Objects.nonNull(cached) && !cached.isExpired(handle.getRefreshIntervalSeconds())) { - return Mono.just(cached.getAutomaton()); + automaton = Mono.just(cached.getAutomaton()); + } else { + automaton = redisTemplate.opsForSet() + .members(redisKey) + .collectList() + // building the automaton is cpu bound, keep it away from the event loop + .map(AhoCorasick::of) + .subscribeOn(Schedulers.boundedElastic()) + .doOnNext(loaded -> SensitiveWordPluginDataHandler.DICTIONARIES.get() + .cachedHandle(redisKey, new CachedDictionary(loaded))) + .onErrorResume(error -> { + LOG.error("sensitive word plugin: cannot read the dictionary from redis key {}", redisKey, error); + if (Objects.nonNull(cached)) { + return Mono.just(cached.getAutomaton()); + } + if (handle.isFailClosed()) { + LOG.warn("sensitive word plugin: no dictionary is available and failClosed is set," + + " the request is rejected"); + return Mono.error(new ResponsiveException(SENSITIVE_WORD_CODE, REJECT_MESSAGE, exchange)); + } + // fail open: a broken redis must not take the traffic down + LOG.warn("sensitive word plugin: no dictionary is available, the request is passed through"); + return Mono.just(AhoCorasick.empty()); + }); } - return redisTemplate.opsForSet() - .members(redisKey) - .collectList() - // building the automaton is cpu bound, keep it away from the event loop - .map(AhoCorasick::of) - .subscribeOn(Schedulers.boundedElastic()) - .doOnNext(automaton -> SensitiveWordPluginDataHandler.DICTIONARIES.get() - .cachedHandle(redisKey, new CachedDictionary(automaton))) - .onErrorResume(error -> { - LOG.error("sensitive word plugin: cannot read the dictionary from redis key {}", redisKey, error); - return Mono.just(Objects.isNull(cached) ? AhoCorasick.empty() : cached.getAutomaton()); - }); + // The cached path emits on the thread that subscribes to it, which is the netty event loop + // of the request being filtered: hop once here so that the scan below never runs on it. + return automaton.publishOn(Schedulers.boundedElastic()); } @Override diff --git a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/ac/AhoCorasick.java b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/ac/AhoCorasick.java index 3c238ed78c10..cc491250e9fa 100644 --- a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/ac/AhoCorasick.java +++ b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/ac/AhoCorasick.java @@ -17,11 +17,13 @@ package org.apache.shenyu.plugin.ai.sensitive.word.ac; +import java.util.ArrayList; import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.LinkedHashSet; import java.util.LinkedList; +import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Queue; @@ -76,6 +78,9 @@ public static AhoCorasick empty() { /** * Find every dictionary word contained in the given text. * + *

The scan is linear in the length of the text: every node carries the words matched by + * itself and by its failure chain, so no failure link is walked here. + * * @param text the text to scan * @return the matched words, in the order they were found, each word is reported once */ @@ -92,11 +97,8 @@ public Set search(final String text) { } TrieNode next = current.children.get(c); current = Objects.isNull(next) ? root : next; - // Walk the whole failure chain, otherwise nested and overlapping words are missed. - for (TrieNode node = current; node != root; node = node.fail) { - if (Objects.nonNull(node.word)) { - matches.add(node.word); - } + if (!current.outputs.isEmpty()) { + matches.addAll(current.outputs); } } return matches; @@ -131,13 +133,22 @@ private void buildFailureLinks() { fail = fail.fail; } child.fail = Objects.nonNull(fail) ? fail.children.get(entry.getKey()) : root; + // Resolve the words matched at this node once, so that scanning a text never has + // to walk the failure chain: the words of the failure chain are the suffixes. + List outputs = new ArrayList<>(); + if (Objects.nonNull(child.word)) { + outputs.add(child.word); + } + outputs.addAll(child.fail.outputs); + child.outputs = outputs.isEmpty() ? Collections.emptyList() : outputs; queue.add(child); } } } /** - * A trie node, {@code word} is not null only on the node that ends a dictionary word. + * A trie node: {@code word} is not null only on the node that ends a dictionary word, and + * {@code outputs} holds every word matched when the automaton is in this state. */ private static final class TrieNode { @@ -146,5 +157,7 @@ private static final class TrieNode { private TrieNode fail; private String word; + + private List outputs = Collections.emptyList(); } } diff --git a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandler.java b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandler.java index a80430849fdd..a718dbbb41de 100644 --- a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandler.java +++ b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/main/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandler.java @@ -100,6 +100,13 @@ public void handlerPlugin(final PluginData pluginData) { } } + @Override + public void removePlugin(final PluginData pluginData) { + REDIS_TEMPLATES.get().removeHandle(PLUGIN_NAME); + REDIS_PROPERTIES.get().removeHandle(PLUGIN_NAME); + LOG.info("sensitive word plugin: released the cached redis template"); + } + @Override public void handlerRule(final RuleData ruleData) { Optional.ofNullable(ruleData.getHandle()).ifPresent(json -> { diff --git a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPluginTest.java b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPluginTest.java index 911c7b0880fb..6b7fe62d4fe2 100644 --- a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPluginTest.java +++ b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/SensitiveWordPluginTest.java @@ -43,6 +43,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.clearInvocations; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -126,11 +127,53 @@ public void testRejectTheRequestContainingASensitiveWord() { MockServerWebExchange exchange = exchange("this request is banned"); StepVerifier.create(plugin.doExecute(exchange, chain, null, ruleData)).verifyComplete(); verify(chain, never()).execute(any(ServerWebExchange.class)); + // the client is told that the request was rejected, never which word matched StepVerifier.create(exchange.getResponse().getBodyAsString()) - .expectNextMatches(body -> body.contains("sensitive words") && body.contains("banned")) + .expectNextMatches(body -> body.contains("sensitive content detected") && !body.contains("banned")) .verifyComplete(); } + @Test + public void testFailClosedRejectsTheRequestWhenTheDictionaryIsUnavailable() { + SensitiveWordPluginDataHandler.CACHED_HANDLE.get() + .cachedHandle(CacheKeyUtils.INST.getKey(ruleData), failClosedHandle()); + ReactiveRedisTemplate redisTemplate = mock(ReactiveRedisTemplate.class); + ReactiveSetOperations setOperations = mock(ReactiveSetOperations.class); + when(redisTemplate.opsForSet()).thenReturn(setOperations); + when(setOperations.members(REDIS_KEY)).thenReturn(Flux.error(new IllegalStateException("redis is down"))); + cacheRedisTemplate(redisTemplate); + + MockServerWebExchange exchange = exchange("a clean request"); + StepVerifier.create(plugin.doExecute(exchange, chain, null, ruleData)).verifyComplete(); + verify(chain, never()).execute(any(ServerWebExchange.class)); + StepVerifier.create(exchange.getResponse().getBodyAsString()) + .expectNextMatches(body -> body.contains("sensitive content detected")) + .verifyComplete(); + } + + @Test + public void testTheStaleDictionaryIsUsedWhenRedisFails() { + mockRedisDictionary("forbidden", "banned"); + MockServerWebExchange first = exchange("a clean request"); + StepVerifier.create(plugin.doExecute(first, chain, null, ruleData)).verifyComplete(); + // the first request legitimately went through, only what follows matters here + clearInvocations(chain); + + // failClosed + refreshIntervalSeconds = 0: the dictionary is read again on every request + SensitiveWordPluginDataHandler.CACHED_HANDLE.get() + .cachedHandle(CacheKeyUtils.INST.getKey(ruleData), failClosedHandle()); + ReactiveRedisTemplate brokenTemplate = mock(ReactiveRedisTemplate.class); + ReactiveSetOperations setOperations = mock(ReactiveSetOperations.class); + when(brokenTemplate.opsForSet()).thenReturn(setOperations); + when(setOperations.members(REDIS_KEY)).thenReturn(Flux.error(new IllegalStateException("redis is down"))); + cacheRedisTemplate(brokenTemplate); + + MockServerWebExchange exchange = exchange("this request is banned"); + StepVerifier.create(plugin.doExecute(exchange, chain, null, ruleData)).verifyComplete(); + // the stale dictionary is still enforced instead of letting the request through + verify(chain, never()).execute(any(ServerWebExchange.class)); + } + @Test public void testPassThroughWhenRedisFails() { ReactiveRedisTemplate redisTemplate = mock(ReactiveRedisTemplate.class); @@ -159,6 +202,14 @@ private void cacheRedisTemplate(final ReactiveRedisTemplate redi .cachedHandle(SensitiveWordPluginDataHandler.PLUGIN_NAME, redisTemplate); } + private SensitiveWordHandle failClosedHandle() { + SensitiveWordHandle handle = SensitiveWordHandle.newDefaultInstance(); + handle.setRedisKey(REDIS_KEY); + handle.setRefreshIntervalSeconds(0L); + handle.setFailClosed(true); + return handle; + } + private MockServerWebExchange exchange(final String body) { return MockServerWebExchange.from(MockServerHttpRequest.post("/ai/chat") .contentType(MediaType.TEXT_PLAIN) diff --git a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandlerTest.java b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandlerTest.java index a13dd290e77f..767da72c7d1d 100644 --- a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandlerTest.java +++ b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-sensitive-word/src/test/java/org/apache/shenyu/plugin/ai/sensitive/word/handler/SensitiveWordPluginDataHandlerTest.java @@ -131,6 +131,21 @@ public void testHandlerPluginCachesTheRedisTemplate() { .obtainHandle(SensitiveWordPluginDataHandler.PLUGIN_NAME)); } + @Test + public void testRemovePluginReleasesTheRedisTemplate() { + PluginData pluginData = new PluginData(); + pluginData.setEnabled(true); + pluginData.setConfig("{\"url\":\"127.0.0.1:6379\"}"); + handler.handlerPlugin(pluginData); + assertNotNull(SensitiveWordPluginDataHandler.REDIS_TEMPLATES.get() + .obtainHandle(SensitiveWordPluginDataHandler.PLUGIN_NAME)); + handler.removePlugin(pluginData); + assertNull(SensitiveWordPluginDataHandler.REDIS_TEMPLATES.get() + .obtainHandle(SensitiveWordPluginDataHandler.PLUGIN_NAME)); + assertNull(SensitiveWordPluginDataHandler.REDIS_PROPERTIES.get() + .obtainHandle(SensitiveWordPluginDataHandler.PLUGIN_NAME)); + } + @Test public void testCachedDictionaryExpires() { CachedDictionary dictionary = new CachedDictionary(AhoCorasick.empty());