Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -0,0 +1,143 @@
/*
* 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 com.nageoffer.ai.ragent.rag.core.security;

import cn.hutool.core.collection.CollUtil;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;

/**
* 瀹夊叏杩囨护鍣ㄩ粯璁ゅ疄鐜?
* <p>
* 閫氳繃鏋勯€犲嚱鏁版敞鍏?{@code List<SafetyRule>} 鈥斺€?Spring 鑷姩鏀堕泦鎵€鏈?
* 瀹炵幇浜?{@link SafetyRule} 鎺ュ彛鐨?Bean锛屾棤闇€鎵嬪姩娉ㄥ唽銆?
* 鍙傝€冿細{@code DefaultMcpToolRegistry} 鐨勮嚜鍔ㄥ彂鐜版ā寮忋€?
* <p>
* 瑙勫垯鎸?{@code @Order} 娉ㄨВ鐨勫€间粠灏忓埌澶т緷娆℃墽琛岋紝閬囧埌棣栦釜 BLOCK 鍗崇煭璺繑鍥炪€?
* 鎵€鏈?WARN 缁撴灉浼氳鏀堕泦锛屾渶缁堝悎骞跺埌娑堟伅鍒楄〃涓€?
*/
@Slf4j
@Service
@RequiredArgsConstructor
public class DefaultSafetyFilter implements SafetyFilter {

private final List<SafetyRule> rules;
private final SecurityProperties properties;
private final SafetyAuditLogger auditLogger;

@Override
public SafetyResult check(SafetyCheckContext context) {
// 鎬诲紑鍏冲叧闂?鈫?鍏ㄩ儴鏀捐
if (!properties.isEnabled()) {
return SafetyResult.pass();
}

if (CollUtil.isEmpty(rules)) {
log.debug("瀹夊叏瑙勫垯鍒楄〃涓虹┖锛岃烦杩囧畨鍏ㄦ鏌?);
return SafetyResult.pass();
}

List<SafetyResult> warnings = new ArrayList<>();

for (SafetyRule rule : rules) {
if (!rule.isEnabled()) {
continue;
}

SafetyResult result;
try {
result = rule.evaluate(context.getOriginalQuestion(), context);
} catch (Exception e) {
log.error("瀹夊叏瑙勫垯 {} 鎵ц寮傚父锛岄檷绾т负 PASS", rule.getId(), e);
continue;
}

if (result.isBlocked()) {
log.warn("瀹夊叏瑙勫垯闃绘柇: ruleId={}, ruleName={}, type={}, confidence={}",
rule.getId(), rule.getName(), result.getRuleType(), result.getConfidence());
auditLogger.log(SafetyAuditEvent.builder()
.ruleId(rule.getId())
.ruleName(rule.getName())
.ruleType(result.getRuleType())
.result(SafetyResult.Type.BLOCK)
.question(truncate(context.getOriginalQuestion()))
.userId(context.getUserId())
.conversationId(context.getConversationId())
.taskId(context.getTaskId())
.message(result.getMessage())
.confidence(result.getConfidence())
.build());
return result;
}

if (result.isWarned()) {
log.info("瀹夊叏瑙勫垯鍛婅: ruleId={}, ruleName={}, type={}, confidence={}",
rule.getId(), rule.getName(), result.getRuleType(), result.getConfidence());
warnings.add(result);
auditLogger.log(SafetyAuditEvent.builder()
.ruleId(rule.getId())
.ruleName(rule.getName())
.ruleType(result.getRuleType())
.result(SafetyResult.Type.WARN)
.question(truncate(context.getOriginalQuestion()))
.userId(context.getUserId())
.conversationId(context.getConversationId())
.taskId(context.getTaskId())
.message(result.getMessage())
.confidence(result.getConfidence())
.build());
}
}

// 娌℃湁 BLOCK锛屼絾瀛樺湪 WARN 鈫?鍚堝苟杩斿洖棣栦釜 WARN锛堟彁绀烘枃妗堬級
if (!warnings.isEmpty()) {
return SafetyResult.warn(
warnings.get(0).getRuleId(),
buildMergedWarnMessage(warnings),
warnings.get(0).getRuleType(),
warnings.stream().mapToDouble(SafetyResult::getConfidence).max().orElse(0.5)
);
}

return SafetyResult.pass();
}

private String buildMergedWarnMessage(List<SafetyResult> warnings) {
if (warnings.size() == 1) {
return warnings.get(0).getMessage();
}
StringBuilder sb = new StringBuilder();
sb.append("[瀹夊叏鎻愮ず] 鎮ㄧ殑鎻愰棶瑙﹀彂浜嗗椤瑰畨鍏ㄧ瓥鐣?(");
for (int i = 0; i < warnings.size(); i++) {
if (i > 0) sb.append("銆?);
sb.append(warnings.get(i).getRuleType().name());
}
sb.append(")锛屾湰娆″洖绛斿凡鑷姩杩藉姞瀹夊叏绾︽潫銆?);
return sb.toString();
}

private String truncate(String text) {
if (text == null) return null;
return text.length() <= 200 ? text : text.substring(0, 200) + "...";
}
}

Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
/*
* 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 com.nageoffer.ai.ragent.rag.core.security;

import com.nageoffer.ai.ragent.framework.context.UserContext;
import com.nageoffer.ai.ragent.framework.exception.ClientException;
import com.nageoffer.ai.ragent.rag.core.mcp.McpToolExecutor;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.Arrays;
import java.util.Set;

/**
* MCP 宸ュ叿鏉冮檺鎵ц鍣?
* <p>
* 鍦?{@code RetrievalEngine.executeSingleMcpTool()} 涓皟鐢紝
* 妫€鏌ュ綋鍓嶇敤鎴锋槸鍚﹀叿澶囪皟鐢ㄧ洰鏍囧伐鍏风殑鏉冮檺銆?
*/
@Slf4j
@Component
@RequiredArgsConstructor
public class McpPermissionEnforcer {

private final SecurityProperties properties;

private static final String ADMIN_ROLE = "admin";

/**
* 妫€鏌ュ綋鍓嶇敤鎴锋槸鍚︽湁鏉冮檺璋冪敤鎸囧畾宸ュ叿
*
* @param toolId 宸ュ叿 ID
* @param executor 宸ュ叿鎵ц鍣?
* @throws ClientException 鏉冮檺涓嶈冻鏃舵姏鍑?
*/
public void checkPermission(String toolId, McpToolExecutor executor) {
if (!properties.getMcp().isEnforceRoles()) {
return;
}

McpToolPermission perm = executor.getClass().getAnnotation(McpToolPermission.class);
if (perm == null) {
// 鏈爣璁?= 鏃犵壒娈婃潈闄愯姹?
return;
}

String currentRole = UserContext.getRole();
if (currentRole == null) {
log.warn("MCP 鏉冮檺鎷掔粷: toolId={}, 褰撳墠鐢ㄦ埛鏃犺鑹?, toolId);
throw new ClientException("褰撳墠鐢ㄦ埛鏃犺鑹叉潈闄愶紝鏃犳硶璋冪敤宸ュ叿: " + toolId);
}

// 绠$悊鍛樹紭鍏?
if (perm.adminOnly() && !ADMIN_ROLE.equalsIgnoreCase(currentRole)) {
log.warn("MCP 鏉冮檺鎷掔粷: toolId={}, 闇€瑕佺鐞嗗憳瑙掕壊, 褰撳墠瑙掕壊={}", toolId, currentRole);
throw new ClientException("宸ュ叿 " + toolId + " 浠呴檺绠$悊鍛樹娇鐢?);
}

// 鎸囧畾瑙掕壊鍖归厤
String[] requiredRoles = perm.requiredRoles();
if (requiredRoles.length > 0) {
Set<String> userRoles = Set.of(currentRole.toLowerCase().split(","));
boolean matched = Arrays.stream(requiredRoles)
.anyMatch(r -> userRoles.contains(r.toLowerCase()));
if (!matched) {
log.warn("MCP 鏉冮檺鎷掔粷: toolId={}, 闇€瑕佽鑹?{}, 褰撳墠瑙掕壊={}", toolId, Arrays.toString(requiredRoles), currentRole);
throw new ClientException("宸ュ叿 " + toolId + " 闇€瑕佽鑹? " + String.join(", ", requiredRoles));
}
}
}
}

Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
/*
* 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 com.nageoffer.ai.ragent.rag.core.security;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

/**
* MCP 宸ュ叿鏉冮檺娉ㄨВ
* <p>
* 鏍囪鍦?{@code McpToolExecutor} 瀹炵幇绫讳笂锛屽0鏄庤皟鐢ㄨ宸ュ叿鎵€闇€鐨勮鑹层€?
* 鏈爣璁扮殑宸ュ叿瑙嗕负鏃犻渶鐗规畩鏉冮檺锛堟墍鏈夌櫥褰曠敤鎴峰彲璋冪敤锛夈€?
*/
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.TYPE)
public @interface McpToolPermission {

/** 闇€瑕佺殑瑙掕壊鍒楄〃 */
String[] requiredRoles() default {};

/** 鏄惁浠呯鐞嗗憳鍙皟鐢?*/
boolean adminOnly() default false;
}

Original file line number Diff line number Diff line change
@@ -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 com.nageoffer.ai.ragent.rag.core.security;

/**
* 杈撳嚭瀹夊叏杩囨护涓婁笅鏂?鈥斺€?璺?chunk 绱Н妫€娴?
*/
public class OutputContext {

private final StringBuilder buffer = new StringBuilder();
private boolean terminated = false;
private final String taskId;

public OutputContext(String taskId) {
this.taskId = taskId;
}

public void append(String chunk) {
buffer.append(chunk);
}

public String getBuffer() {
return buffer.toString();
}

public int getBufferLength() {
return buffer.length();
}

/**
* 缂撳啿鍖鸿秴杩囦笂闄愭椂鎴柇鍓嶅崐閮ㄥ垎
*/
public void truncateHalf() {
int half = buffer.length() / 2;
buffer.delete(0, half);
}

public boolean isTerminated() {
return terminated;
}

public void markTerminated() {
this.terminated = true;
}

public String getTaskId() {
return taskId;
}
}

Loading