|
|
@@ -5,103 +5,54 @@ import jakarta.servlet.FilterChain;
|
|
|
import jakarta.servlet.ServletException;
|
|
|
import jakarta.servlet.http.HttpServletRequest;
|
|
|
import jakarta.servlet.http.HttpServletResponse;
|
|
|
-import jakarta.servlet.http.HttpServletResponseWrapper;
|
|
|
import org.springframework.stereotype.Component;
|
|
|
import org.springframework.web.filter.OncePerRequestFilter;
|
|
|
|
|
|
import java.io.IOException;
|
|
|
-import java.time.LocalDate;
|
|
|
import java.util.Map;
|
|
|
import java.util.concurrent.ConcurrentHashMap;
|
|
|
import java.util.concurrent.atomic.AtomicInteger;
|
|
|
|
|
|
+/**
|
|
|
+ * 用户侧限流:仅限 chat POST 接口,每 IP 每分钟限制次数。
|
|
|
+ * 超限返回 429,不会进入 Controller 层,因此不会写缓存和 DB。
|
|
|
+ */
|
|
|
@Component
|
|
|
public class RateLimitFilter extends OncePerRequestFilter {
|
|
|
|
|
|
private final RateLimitProperties props;
|
|
|
- private final JwtUtil jwtUtil;
|
|
|
private final Map<String, AtomicInteger> minuteCounter = new ConcurrentHashMap<>();
|
|
|
- private final Map<String, AtomicInteger> dailyCounter = new ConcurrentHashMap<>();
|
|
|
- private volatile String currentDate = LocalDate.now().toString();
|
|
|
|
|
|
- public RateLimitFilter(RateLimitProperties props, JwtUtil jwtUtil) {
|
|
|
+ public RateLimitFilter(RateLimitProperties props) {
|
|
|
this.props = props;
|
|
|
- this.jwtUtil = jwtUtil;
|
|
|
}
|
|
|
|
|
|
@Override
|
|
|
protected void doFilterInternal(HttpServletRequest request,
|
|
|
HttpServletResponse response,
|
|
|
FilterChain filterChain) throws ServletException, IOException {
|
|
|
- // 日期切换时清空每日计数
|
|
|
- String today = LocalDate.now().toString();
|
|
|
- if (!currentDate.equals(today)) {
|
|
|
- dailyCounter.clear();
|
|
|
- currentDate = today;
|
|
|
- }
|
|
|
-
|
|
|
- String ip = request.getRemoteAddr();
|
|
|
boolean isChatApi = request.getRequestURI().contains("/chat/")
|
|
|
&& request.getMethod().equalsIgnoreCase("POST");
|
|
|
|
|
|
- // 仅 chat 接口走限流,药品库等接口跳过
|
|
|
if (isChatApi) {
|
|
|
- // 每 IP 每分钟限制
|
|
|
+ String ip = request.getRemoteAddr();
|
|
|
var ipCount = minuteCounter.computeIfAbsent(ip, k -> new AtomicInteger(0));
|
|
|
- if (ipCount.incrementAndGet() > props.getPerMinute()) {
|
|
|
+ int count = ipCount.incrementAndGet();
|
|
|
+ if (count > props.getPerMinute()) {
|
|
|
+ ipCount.decrementAndGet(); // 超限不计
|
|
|
writeRateLimitError(response, "请求过于频繁,请稍后再试");
|
|
|
return;
|
|
|
}
|
|
|
-
|
|
|
- // 每日限制
|
|
|
- String userId = getUserId(request);
|
|
|
- if (userId == null) userId = "ip:" + ip;
|
|
|
- var dailyCount = dailyCounter.computeIfAbsent(userId, k -> new AtomicInteger(0));
|
|
|
- if (dailyCount.get() >= props.getPerDay()) {
|
|
|
- writeRateLimitError(response, "今日提问次数已用完,请明天再提问");
|
|
|
- return;
|
|
|
- }
|
|
|
}
|
|
|
|
|
|
- // 包装 response 以捕获状态码
|
|
|
- var statusWrapper = new HttpServletResponseWrapper(response) {
|
|
|
- private int status = 200;
|
|
|
- @Override public void setStatus(int sc) { this.status = sc; super.setStatus(sc); }
|
|
|
- @Override public int getStatus() { return this.status; }
|
|
|
- };
|
|
|
-
|
|
|
- filterChain.doFilter(request, statusWrapper);
|
|
|
-
|
|
|
- // 只有成功响应(2xx)才计入每日次数
|
|
|
- if (isChatApi) {
|
|
|
- String userId = getUserId(request);
|
|
|
- if (userId == null) userId = "ip:" + ip;
|
|
|
- int status = statusWrapper.getStatus();
|
|
|
- if (status >= 200 && status < 300) {
|
|
|
- dailyCounter.computeIfAbsent(userId, k -> new AtomicInteger(0)).incrementAndGet();
|
|
|
- }
|
|
|
- }
|
|
|
+ filterChain.doFilter(request, response);
|
|
|
}
|
|
|
|
|
|
private static void writeRateLimitError(HttpServletResponse response, String message) throws IOException {
|
|
|
response.setStatus(429);
|
|
|
response.setContentType("application/json;charset=UTF-8");
|
|
|
- // 使用 detail 字段,前端 api/ai.js 的 request() 会优先读取 res.data.detail
|
|
|
response.getWriter().write(
|
|
|
"{\"error\":\"" + message + "\",\"detail\":\"" + message + "\",\"code\":429}"
|
|
|
);
|
|
|
}
|
|
|
-
|
|
|
- private String getUserId(HttpServletRequest request) {
|
|
|
- String header = request.getHeader("Authorization");
|
|
|
- if (header != null && header.startsWith("Bearer ")) {
|
|
|
- try {
|
|
|
- String token = header.substring(7);
|
|
|
- if (jwtUtil.validateToken(token)) {
|
|
|
- return jwtUtil.getSubject(token);
|
|
|
- }
|
|
|
- } catch (Exception ignored) {}
|
|
|
- }
|
|
|
- return null;
|
|
|
- }
|
|
|
-}
|
|
|
+}
|