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..fd51e9eed39e --- /dev/null +++ b/shenyu-common/src/main/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandle.java @@ -0,0 +1,111 @@ +/* + * 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; + + /** + * 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. + * + * @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; + } + + /** + * 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. + * + * @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..ddc61e24db16 --- /dev/null +++ b/shenyu-common/src/test/java/org/apache/shenyu/common/dto/convert/rule/SensitiveWordHandleTest.java @@ -0,0 +1,54 @@ +/* + * 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; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * 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 testFailClosed() { + SensitiveWordHandle handle = new SensitiveWordHandle(); + assertFalse(handle.isFailClosed()); + handle.setFailClosed(true); + assertTrue(handle.isFailClosed()); + } + + @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..6ce6f14335a2 --- /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,176 @@ +/* + * 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.HttpStatus; +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 a client whose request was rejected. It follows the waf plugin, + * which rejects with 403 as well. + */ + 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; + + 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(exchange, redisTemplate, handle) + .map(automaton -> automaton.search(body)) + .flatMap(matches -> { + if (matches.isEmpty()) { + return Mono.just(body); + } + // 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)); + }); + } + + /** + * 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 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())) { + 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()); + }); + } + // 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 + 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..cc491250e9fa --- /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,163 @@ +/* + * 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.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; +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. + * + *

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 + */ + 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; + if (!current.outputs.isEmpty()) { + matches.addAll(current.outputs); + } + } + 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; + // 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, and + * {@code outputs} holds every word matched when the automaton is in this state. + */ + private static final class TrieNode { + + private final Map children = new HashMap<>(); + + 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 new file mode 100644 index 000000000000..a718dbbb41de --- /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,168 @@ +/* + * 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 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 -> { + 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..6b7fe62d4fe2 --- /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,218 @@ +/* + * 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.clearInvocations; +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)); + // the client is told that the request was rejected, never which word matched + StepVerifier.create(exchange.getResponse().getBodyAsString()) + .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); + 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 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) + .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..767da72c7d1d --- /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,166 @@ +/* + * 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 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()); + 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()); + }); + } +}