From 61aeba9c650a7dbfab274f1976bafa0d05eed68e Mon Sep 17 00:00:00 2001 From: ikmkj Date: Tue, 3 Mar 2026 18:23:28 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=A2=9E=E5=BC=BA=E5=AE=89=E5=85=A8?= =?UTF-8?q?=E7=BB=84=E4=BB=B6=E5=B9=B6=E6=B7=BB=E5=8A=A0=E5=86=85=E5=AD=98?= =?UTF-8?q?=E4=BF=9D=E6=8A=A4=E6=9C=BA=E5=88=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. 重构 ReplayAttackInterceptor、RateLimitInterceptor、CaptchaUtil 和 LoginLockUtil，使用 LRUCache 和读写锁优化内存管理 2. 新增 MemoryProtector 类实现内存监控和保护机制 3. 为所有内存缓存组件添加容量限制和过期清理策略 4. 更新 .gitignore 文件配置 --- .gitignore | 1 + .../interceptor/RateLimitInterceptor.java | 97 ++++++++--- .../interceptor/ReplayAttackInterceptor.java | 78 +++++++-- .../test/bijihoudaun/util/CaptchaUtil.java | 127 +++++++++++--- .../test/bijihoudaun/util/LoginLockUtil.java | 161 +++++++++++++----- .../bijihoudaun/util/MemoryProtector.java | 124 ++++++++++++++ 6 files changed, 488 insertions(+), 100 deletions(-) create mode 100644 biji-houdaun/src/main/java/com/test/bijihoudaun/util/MemoryProtector.java diff --git a/.gitignore b/.gitignore index 5dbfd19..59f539c 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,7 @@ .idea target .trae +.trae/* *.iml .roo out diff --git a/biji-houdaun/src/main/java/com/test/bijihoudaun/interceptor/RateLimitInterceptor.java b/biji-houdaun/src/main/java/com/test/bijihoudaun/interceptor/RateLimitInterceptor.java index 25d5a5c..9871a68 100644 --- a/biji-houdaun/src/main/java/com/test/bijihoudaun/interceptor/RateLimitInterceptor.java +++ b/biji-houdaun/src/main/java/com/test/bijihoudaun/interceptor/RateLimitInterceptor.java @@ -3,19 +3,22 @@ package com.test.bijihoudaun.interceptor; import com.fasterxml.jackson.databind.ObjectMapper; import com.test.bijihoudaun.common.response.R; import com.test.bijihoudaun.common.response.ResultCode; +import com.test.bijihoudaun.util.MemoryProtector; import jakarta.servlet.http.HttpServletRequest; import jakarta.servlet.http.HttpServletResponse; +import lombok.extern.slf4j.Slf4j; import org.springframework.security.core.Authentication; import org.springframework.security.core.context.SecurityContextHolder; import org.springframework.web.servlet.HandlerInterceptor; +import java.util.LinkedHashMap; import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.ReentrantReadWriteLock; /** - * 限流拦截器 - 支持按 IP 和按用户双重限流 + * 限流拦截器 - 支持按 IP 和按用户双重限流，带容量限制 */ +@Slf4j public class RateLimitInterceptor implements HandlerInterceptor { // 普通接口限流配置 @@ -26,40 +29,70 @@ public class RateLimitInterceptor implements HandlerInterceptor { private static final int MAX_LOGIN_REQUESTS_PER_MINUTE_USER = 10; // 时间窗口（毫秒） private static final long WINDOW_SIZE_MS = 60_000; + // 最大存储记录数（防止内存溢出） + private static final int MAX_RECORDS = 50000; // IP 级别限流 - private final Map ipCounters = new ConcurrentHashMap<>(); - private final Map ipLoginCounters = new ConcurrentHashMap<>(); + private static final LRUCache ipCounters = new LRUCache<>(MAX_RECORDS / 2); + private static final LRUCache ipLoginCounters = new LRUCache<>(MAX_RECORDS / 4); // 用户级别限流 - private final Map userCounters = new ConcurrentHashMap<>(); - private final Map userLoginCounters = new ConcurrentHashMap<>(); + private static final LRUCache userCounters = new LRUCache<>(MAX_RECORDS / 4); + private static final LRUCache userLoginCounters = new LRUCache<>(MAX_RECORDS / 4); + + private static final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); private static class RequestCounter { - AtomicInteger count; + int count; long windowStart; RequestCounter() { - this.count = new AtomicInteger(1); + this.count = 1; this.windowStart = System.currentTimeMillis(); } boolean incrementAndCheck(int maxRequests) { long now = System.currentTimeMillis(); if (now - windowStart > WINDOW_SIZE_MS) { - synchronized (this) { - if (now - windowStart > WINDOW_SIZE_MS) { - count.set(1); - windowStart = now; - return true; - } - } + // 新窗口 + count = 1; + windowStart = now; + return true; } - return count.incrementAndGet() <= maxRequests; + count++; + return count <= maxRequests; + } + } + + /** + * 简单的 LRU 缓存实现 + */ + private static class LRUCache extends LinkedHashMap { + private final int maxSize; + + LRUCache(int maxSize) { + super(maxSize, 0.75f, true); + this.maxSize = maxSize; + } + + @Override + protected boolean removeEldestEntry(Map.Entry eldest) { + boolean shouldRemove = size() > maxSize; + if (shouldRemove) { + log.debug("限流记录达到上限，移除最旧的记录"); + } + return shouldRemove; } } @Override public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception { + // 检查内存状态 + if (MemoryProtector.isMemoryInsufficient()) { + log.warn("内存不足，拒绝请求: {}", request.getRequestURI()); + MemoryProtector.writeMemoryInsufficientResponse(response); + return false; + } + String clientIp = getClientIp(request); String requestUri = request.getRequestURI(); String username = getCurrentUsername(); @@ -84,14 +117,19 @@ public class RateLimitInterceptor implements HandlerInterceptor { */ private boolean checkIpLimit(String clientIp, boolean isLoginRequest, HttpServletResponse response) throws Exception { - Map counters = isLoginRequest ? ipLoginCounters : ipCounters; + LRUCache counters = isLoginRequest ? ipLoginCounters : ipCounters; int maxRequests = isLoginRequest ? MAX_LOGIN_REQUESTS_PER_MINUTE : MAX_REQUESTS_PER_MINUTE; - RequestCounter counter = counters.computeIfAbsent(clientIp, k -> new RequestCounter()); + lock.writeLock().lock(); + try { + RequestCounter counter = counters.computeIfAbsent(clientIp, k -> new RequestCounter()); - if (!counter.incrementAndCheck(maxRequests)) { - writeRateLimitResponse(response, "请求过于频繁，请稍后再试"); - return false; + if (!counter.incrementAndCheck(maxRequests)) { + writeRateLimitResponse(response, "请求过于频繁，请稍后再试"); + return false; + } + } finally { + lock.writeLock().unlock(); } return true; } @@ -101,14 +139,19 @@ public class RateLimitInterceptor implements HandlerInterceptor { */ private boolean checkUserLimit(String username, boolean isLoginRequest, HttpServletResponse response) throws Exception { - Map counters = isLoginRequest ? userLoginCounters : userCounters; + LRUCache counters = isLoginRequest ? userLoginCounters : userCounters; int maxRequests = isLoginRequest ? MAX_LOGIN_REQUESTS_PER_MINUTE_USER : MAX_REQUESTS_PER_MINUTE_USER; - RequestCounter counter = counters.computeIfAbsent(username, k -> new RequestCounter()); + lock.writeLock().lock(); + try { + RequestCounter counter = counters.computeIfAbsent(username, k -> new RequestCounter()); - if (!counter.incrementAndCheck(maxRequests)) { - writeRateLimitResponse(response, "您的操作过于频繁，请稍后再试"); - return false; + if (!counter.incrementAndCheck(maxRequests)) { + writeRateLimitResponse(response, "您的操作过于频繁，请稍后再试"); + return false; + } + } finally { + lock.writeLock().unlock(); } return true; } diff --git a/biji-houdaun/src/main/java/com/test/bijihoudaun/interceptor/ReplayAttackInterceptor.java b/biji-houdaun/src/main/java/com/test/bijihoudaun/interceptor/ReplayAttackInterceptor.java index 5364367..c1ea5fc 100644 --- a/biji-houdaun/src/main/java/com/test/bijihoudaun/interceptor/ReplayAttackInterceptor.java +++ b/biji-houdaun/src/main/java/com/test/bijihoudaun/interceptor/ReplayAttackInterceptor.java @@ -3,26 +3,33 @@ package com.test.bijihoudaun.interceptor; import com.fasterxml.jackson.databind.ObjectMapper; import com.test.bijihoudaun.common.response.R; import com.test.bijihoudaun.common.response.ResultCode; +import com.test.bijihoudaun.util.MemoryProtector; import jakarta.servlet.http.HttpServletRequest; import jakarta.servlet.http.HttpServletResponse; +import lombok.extern.slf4j.Slf4j; import org.springframework.web.servlet.HandlerInterceptor; import java.util.Arrays; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.Map; import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.locks.ReentrantReadWriteLock; /** * 防重放攻击拦截器 - * 通过时间戳和 nonce 机制防止请求被截获重放 + * 通过时间戳和 nonce 机制防止请求被截获重放，带容量限制 */ +@Slf4j public class ReplayAttackInterceptor implements HandlerInterceptor { // 请求时间戳有效期（毫秒）：5分钟 private static final long TIMESTAMP_VALIDITY = 5 * 60 * 1000; // nonce 有效期（毫秒）：10分钟 private static final long NONCE_EXPIRE_TIME = 10 * 60 * 1000; + // 最大存储 nonce 数（防止内存溢出） + private static final int MAX_NONCES = 100000; + // 用于不需要验证的路径 private static final Set EXCLUDE_PATHS = new HashSet<>(Arrays.asList( "/api/user/login", @@ -31,11 +38,40 @@ public class ReplayAttackInterceptor implements HandlerInterceptor { )); // 存储已使用的 nonce：key=nonce，value=使用时间戳 - private static final Map usedNonces = new ConcurrentHashMap<>(); + private static final LRUCache usedNonces = new LRUCache<>(MAX_NONCES); + private static final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); + + /** + * 简单的 LRU 缓存实现 + */ + private static class LRUCache extends LinkedHashMap { + private final int maxSize; + + LRUCache(int maxSize) { + super(maxSize, 0.75f, true); + this.maxSize = maxSize; + } + + @Override + protected boolean removeEldestEntry(Map.Entry eldest) { + boolean shouldRemove = size() > maxSize; + if (shouldRemove) { + log.debug("nonce 存储达到上限，移除最旧的记录"); + } + return shouldRemove; + } + } @Override public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception { + // 检查内存状态 + if (MemoryProtector.isMemoryInsufficient()) { + log.warn("内存不足，拒绝请求: {}", request.getRequestURI()); + MemoryProtector.writeMemoryInsufficientResponse(response); + return false; + } + String requestUri = request.getRequestURI(); // 排除不需要验证的路径 @@ -77,13 +113,23 @@ public class ReplayAttackInterceptor implements HandlerInterceptor { } // 验证 nonce 是否已被使用 - if (usedNonces.containsKey(nonce)) { - writeErrorResponse(response, "检测到重放攻击，请求已被拒绝"); - return false; + lock.readLock().lock(); + try { + if (usedNonces.containsKey(nonce)) { + writeErrorResponse(response, "检测到重放攻击，请求已被拒绝"); + return false; + } + } finally { + lock.readLock().unlock(); } // 记录 nonce - usedNonces.put(nonce, now); + lock.writeLock().lock(); + try { + usedNonces.put(nonce, now); + } finally { + lock.writeLock().unlock(); + } return true; } @@ -92,10 +138,20 @@ public class ReplayAttackInterceptor implements HandlerInterceptor { * 清理过期的 nonce */ private void cleanupExpiredNonces() { - long now = System.currentTimeMillis(); - usedNonces.entrySet().removeIf(entry -> - (now - entry.getValue()) > NONCE_EXPIRE_TIME - ); + // 每80%容量时清理一次，减少开销 + if (usedNonces.size() < MAX_NONCES * 0.8) { + return; + } + + lock.writeLock().lock(); + try { + long now = System.currentTimeMillis(); + usedNonces.entrySet().removeIf(entry -> + (now - entry.getValue()) > NONCE_EXPIRE_TIME + ); + } finally { + lock.writeLock().unlock(); + } } /** diff --git a/biji-houdaun/src/main/java/com/test/bijihoudaun/util/CaptchaUtil.java b/biji-houdaun/src/main/java/com/test/bijihoudaun/util/CaptchaUtil.java index 3cd8775..c643416 100644 --- a/biji-houdaun/src/main/java/com/test/bijihoudaun/util/CaptchaUtil.java +++ b/biji-houdaun/src/main/java/com/test/bijihoudaun/util/CaptchaUtil.java @@ -1,5 +1,7 @@ package com.test.bijihoudaun.util; +import lombok.extern.slf4j.Slf4j; + import javax.imageio.ImageIO; import java.awt.*; import java.awt.image.BufferedImage; @@ -8,14 +10,15 @@ import java.io.IOException; import java.security.SecureRandom; import java.time.LocalDateTime; import java.time.temporal.ChronoUnit; -import java.util.Base64; +import java.util.LinkedHashMap; import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.locks.ReentrantReadWriteLock; /** * 图形验证码工具类 - * 使用本地内存存储验证码 + * 使用本地内存存储验证码，带容量限制 */ +@Slf4j public class CaptchaUtil { // 验证码有效期（分钟） @@ -26,9 +29,12 @@ public class CaptchaUtil { private static final int IMAGE_WIDTH = 120; // 图片高度 private static final int IMAGE_HEIGHT = 40; + // 最大存储验证码数（防止内存溢出） + private static final int MAX_CAPTCHAS = 5000; // 存储验证码：key=验证码ID，value=验证码记录 - private static final Map captchaStore = new ConcurrentHashMap<>(); + private static final LRUCache captchaStore = new LRUCache<>(MAX_CAPTCHAS); + private static final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); // 安全随机数生成器 private static final SecureRandom random = new SecureRandom(); @@ -47,6 +53,27 @@ public class CaptchaUtil { } } + /** + * 简单的 LRU 缓存实现 + */ + private static class LRUCache extends LinkedHashMap { + private final int maxSize; + + LRUCache(int maxSize) { + super(maxSize, 0.75f, true); + this.maxSize = maxSize; + } + + @Override + protected boolean removeEldestEntry(Map.Entry eldest) { + boolean shouldRemove = size() > maxSize; + if (shouldRemove) { + log.warn("验证码存储达到上限，移除最旧的验证码: {}", eldest.getKey()); + } + return shouldRemove; + } + } + /** * 验证码结果 */ @@ -73,6 +100,12 @@ public class CaptchaUtil { * @return 包含验证码ID和Base64图片的结果 */ public static CaptchaResult generateCaptcha() { + // 检查内存状态 + if (MemoryProtector.isMemoryInsufficient()) { + log.error("内存不足，拒绝生成验证码"); + throw new RuntimeException("服务器繁忙，请稍后再试"); + } + // 清理过期验证码 cleanupExpiredCaptchas(); @@ -85,8 +118,13 @@ public class CaptchaUtil { // 生成图片 String base64Image = generateImage(code); - // 存储验证码 - captchaStore.put(captchaId, new CaptchaRecord(code)); + lock.writeLock().lock(); + try { + // 存储验证码 + captchaStore.put(captchaId, new CaptchaRecord(code)); + } finally { + lock.writeLock().unlock(); + } return new CaptchaResult(captchaId, base64Image); } @@ -102,24 +140,43 @@ public class CaptchaUtil { return false; } - CaptchaRecord record = captchaStore.get(captchaId); - if (record == null || record.isExpired()) { - // 验证码不存在或已过期，移除 - if (record != null) { - captchaStore.remove(captchaId); + lock.readLock().lock(); + try { + CaptchaRecord record = captchaStore.get(captchaId); + if (record == null || record.isExpired()) { + // 验证码不存在或已过期，移除 + if (record != null) { + lock.readLock().unlock(); + lock.writeLock().lock(); + try { + captchaStore.remove(captchaId); + } finally { + lock.writeLock().unlock(); + lock.readLock().lock(); + } + } + return false; } - return false; + + // 验证码比对（不区分大小写） + boolean success = record.code.equalsIgnoreCase(code); + + // 验证成功后立即删除（一次性使用） + if (success) { + lock.readLock().unlock(); + lock.writeLock().lock(); + try { + captchaStore.remove(captchaId); + } finally { + lock.writeLock().unlock(); + lock.readLock().lock(); + } + } + + return success; + } finally { + lock.readLock().unlock(); } - - // 验证码比对（不区分大小写） - boolean success = record.code.equalsIgnoreCase(code); - - // 验证成功后立即删除（一次性使用） - if (success) { - captchaStore.remove(captchaId); - } - - return success; } /** @@ -186,7 +243,7 @@ public class CaptchaUtil { try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) { ImageIO.write(image, "png", baos); byte[] imageBytes = baos.toByteArray(); - return "data:image/png;base64," + Base64.getEncoder().encodeToString(imageBytes); + return "data:image/png;base64," + java.util.Base64.getEncoder().encodeToString(imageBytes); } catch (IOException e) { throw new RuntimeException("生成验证码图片失败", e); } @@ -196,6 +253,28 @@ public class CaptchaUtil { * 清理过期验证码 */ private static void cleanupExpiredCaptchas() { - captchaStore.entrySet().removeIf(entry -> entry.getValue().isExpired()); + // 每80%容量时清理一次，减少开销 + if (captchaStore.size() < MAX_CAPTCHAS * 0.8) { + return; + } + + lock.writeLock().lock(); + try { + captchaStore.entrySet().removeIf(entry -> entry.getValue().isExpired()); + } finally { + lock.writeLock().unlock(); + } + } + + /** + * 获取当前验证码数量（用于监控） + */ + public static int getCaptchaCount() { + lock.readLock().lock(); + try { + return captchaStore.size(); + } finally { + lock.readLock().unlock(); + } } } diff --git a/biji-houdaun/src/main/java/com/test/bijihoudaun/util/LoginLockUtil.java b/biji-houdaun/src/main/java/com/test/bijihoudaun/util/LoginLockUtil.java index 51cd7dd..95777dd 100644 --- a/biji-houdaun/src/main/java/com/test/bijihoudaun/util/LoginLockUtil.java +++ b/biji-houdaun/src/main/java/com/test/bijihoudaun/util/LoginLockUtil.java @@ -1,13 +1,18 @@ package com.test.bijihoudaun.util; +import lombok.extern.slf4j.Slf4j; + import java.time.LocalDateTime; import java.time.temporal.ChronoUnit; -import java.util.concurrent.ConcurrentHashMap; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.concurrent.locks.ReentrantReadWriteLock; /** * 登录锁定工具类 - * 使用本地内存存储登录失败记录 + * 使用本地内存存储登录失败记录，带容量限制 */ +@Slf4j public class LoginLockUtil { // 最大失败次数 @@ -16,9 +21,12 @@ public class LoginLockUtil { private static final int LOCK_TIME_MINUTES = 30; // 失败记录过期时间（分钟） private static final int RECORD_EXPIRE_MINUTES = 60; + // 最大存储记录数（防止内存溢出） + private static final int MAX_RECORDS = 10000; // 登录失败记录：key=用户名，value=失败记录 - private static final ConcurrentHashMap attempts = new ConcurrentHashMap<>(); + private static final LRUCache attempts = new LRUCache<>(MAX_RECORDS); + private static final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); private static class LoginAttempt { int failedCount; @@ -32,6 +40,27 @@ public class LoginLockUtil { } } + /** + * 简单的 LRU 缓存实现 + */ + private static class LRUCache extends LinkedHashMap { + private final int maxSize; + + LRUCache(int maxSize) { + super(maxSize, 0.75f, true); + this.maxSize = maxSize; + } + + @Override + protected boolean removeEldestEntry(Map.Entry eldest) { + boolean shouldRemove = size() > maxSize; + if (shouldRemove) { + log.warn("登录锁定记录达到上限，移除最旧的记录: {}", eldest.getKey()); + } + return shouldRemove; + } + } + /** * 记录登录失败 * @param username 用户名 @@ -41,13 +70,19 @@ public class LoginLockUtil { cleanupExpiredRecords(); - LoginAttempt attempt = attempts.computeIfAbsent(username, k -> new LoginAttempt()); - attempt.failedCount++; - attempt.lastAttemptTime = LocalDateTime.now(); + lock.writeLock().lock(); + try { + LoginAttempt attempt = attempts.computeIfAbsent(username, k -> new LoginAttempt()); + attempt.failedCount++; + attempt.lastAttemptTime = LocalDateTime.now(); - // 达到最大失败次数，锁定账号 - if (attempt.failedCount >= MAX_FAILED_ATTEMPTS) { - attempt.lockUntil = LocalDateTime.now().plusMinutes(LOCK_TIME_MINUTES); + // 达到最大失败次数，锁定账号 + if (attempt.failedCount >= MAX_FAILED_ATTEMPTS) { + attempt.lockUntil = LocalDateTime.now().plusMinutes(LOCK_TIME_MINUTES); + log.warn("账号 [{}] 已被锁定 {} 分钟", username, LOCK_TIME_MINUTES); + } + } finally { + lock.writeLock().unlock(); } } @@ -57,7 +92,13 @@ public class LoginLockUtil { */ public static void recordSuccess(String username) { if (username == null || username.isEmpty()) return; - attempts.remove(username); + + lock.writeLock().lock(); + try { + attempts.remove(username); + } finally { + lock.writeLock().unlock(); + } } /** @@ -68,20 +109,32 @@ public class LoginLockUtil { public static boolean isLocked(String username) { if (username == null || username.isEmpty()) return false; - LoginAttempt attempt = attempts.get(username); - if (attempt == null) return false; + lock.readLock().lock(); + try { + LoginAttempt attempt = attempts.get(username); + if (attempt == null) return false; - // 检查是否仍在锁定时间内 - if (attempt.lockUntil != null) { - if (LocalDateTime.now().isBefore(attempt.lockUntil)) { - return true; - } else { - // 锁定时间已过，清除记录 - attempts.remove(username); - return false; + // 检查是否仍在锁定时间内 + if (attempt.lockUntil != null) { + if (LocalDateTime.now().isBefore(attempt.lockUntil)) { + return true; + } else { + // 锁定时间已过，清除记录 + lock.readLock().unlock(); + lock.writeLock().lock(); + try { + attempts.remove(username); + } finally { + lock.writeLock().unlock(); + lock.readLock().lock(); + } + return false; + } } + return false; + } finally { + lock.readLock().unlock(); } - return false; } /** @@ -92,11 +145,16 @@ public class LoginLockUtil { public static long getRemainingLockTime(String username) { if (username == null || username.isEmpty()) return 0; - LoginAttempt attempt = attempts.get(username); - if (attempt == null || attempt.lockUntil == null) return 0; + lock.readLock().lock(); + try { + LoginAttempt attempt = attempts.get(username); + if (attempt == null || attempt.lockUntil == null) return 0; - long remaining = ChronoUnit.SECONDS.between(LocalDateTime.now(), attempt.lockUntil); - return Math.max(0, remaining); + long remaining = ChronoUnit.SECONDS.between(LocalDateTime.now(), attempt.lockUntil); + return Math.max(0, remaining); + } finally { + lock.readLock().unlock(); + } } /** @@ -107,25 +165,52 @@ public class LoginLockUtil { public static int getRemainingAttempts(String username) { if (username == null || username.isEmpty()) return MAX_FAILED_ATTEMPTS; - LoginAttempt attempt = attempts.get(username); - if (attempt == null) return MAX_FAILED_ATTEMPTS; + lock.readLock().lock(); + try { + LoginAttempt attempt = attempts.get(username); + if (attempt == null) return MAX_FAILED_ATTEMPTS; - return Math.max(0, MAX_FAILED_ATTEMPTS - attempt.failedCount); + return Math.max(0, MAX_FAILED_ATTEMPTS - attempt.failedCount); + } finally { + lock.readLock().unlock(); + } } /** * 清理过期记录 */ private static void cleanupExpiredRecords() { - LocalDateTime now = LocalDateTime.now(); - attempts.entrySet().removeIf(entry -> { - LoginAttempt attempt = entry.getValue(); - // 未锁定且长时间没有登录的记录 - if (attempt.lockUntil == null) { - return ChronoUnit.MINUTES.between(attempt.lastAttemptTime, now) > RECORD_EXPIRE_MINUTES; - } - // 锁定已过期的记录 - return now.isAfter(attempt.lockUntil); - }); + // 每100次操作清理一次，减少开销 + if (attempts.size() < MAX_RECORDS * 0.8) { + return; + } + + lock.writeLock().lock(); + try { + LocalDateTime now = LocalDateTime.now(); + attempts.entrySet().removeIf(entry -> { + LoginAttempt attempt = entry.getValue(); + // 未锁定且长时间没有登录的记录 + if (attempt.lockUntil == null) { + return ChronoUnit.MINUTES.between(attempt.lastAttemptTime, now) > RECORD_EXPIRE_MINUTES; + } + // 锁定已过期的记录 + return now.isAfter(attempt.lockUntil); + }); + } finally { + lock.writeLock().unlock(); + } + } + + /** + * 获取当前记录数量（用于监控） + */ + public static int getRecordCount() { + lock.readLock().lock(); + try { + return attempts.size(); + } finally { + lock.readLock().unlock(); + } } } diff --git a/biji-houdaun/src/main/java/com/test/bijihoudaun/util/MemoryProtector.java b/biji-houdaun/src/main/java/com/test/bijihoudaun/util/MemoryProtector.java new file mode 100644 index 0000000..27bcfc6 --- /dev/null +++ b/biji-houdaun/src/main/java/com/test/bijihoudaun/util/MemoryProtector.java @@ -0,0 +1,124 @@ +package com.test.bijihoudaun.util; + +import com.test.bijihoudaun.common.response.R; +import com.test.bijihoudaun.common.response.ResultCode; +import jakarta.servlet.http.HttpServletResponse; +import lombok.extern.slf4j.Slf4j; + +import java.io.IOException; +import java.lang.management.ManagementFactory; +import java.lang.management.MemoryMXBean; +import java.lang.management.MemoryUsage; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * 内存保护工具类 + * 监控 JVM 内存使用情况，当内存不足时拒绝请求 + */ +@Slf4j +public class MemoryProtector { + + // 内存使用率阈值（百分比），超过此值进入保护模式 + private static final double MEMORY_THRESHOLD_PERCENT = 85.0; + // 堆内存使用阈值（MB） + private static final long HEAP_THRESHOLD_MB = 1500; + + // 是否处于保护模式 + private static final AtomicBoolean PROTECTION_MODE = new AtomicBoolean(false); + + // 上次检查时间 + private static volatile long lastCheckTime = 0; + // 检查间隔（毫秒） + private static final long CHECK_INTERVAL_MS = 5000; + + /** + * 检查内存状态，如果内存不足返回 true + */ + public static boolean isMemoryInsufficient() { + long now = System.currentTimeMillis(); + // 每 5 秒检查一次，减少性能开销 + if (now - lastCheckTime < CHECK_INTERVAL_MS) { + return PROTECTION_MODE.get(); + } + + synchronized (MemoryProtector.class) { + if (now - lastCheckTime < CHECK_INTERVAL_MS) { + return PROTECTION_MODE.get(); + } + lastCheckTime = now; + + MemoryMXBean memoryMXBean = ManagementFactory.getMemoryMXBean(); + MemoryUsage heapUsage = memoryMXBean.getHeapMemoryUsage(); + + long usedMB = heapUsage.getUsed() / 1024 / 1024; + long maxMB = heapUsage.getMax() / 1024 / 1024; + double usagePercent = maxMB > 0 ? (usedMB * 100.0 / maxMB) : 0; + + boolean shouldProtect = usagePercent > MEMORY_THRESHOLD_PERCENT || usedMB > HEAP_THRESHOLD_MB; + + if (shouldProtect && !PROTECTION_MODE.get()) { + log.warn("内存不足，进入保护模式 - 使用率: {}%, 已使用: {}MB, 最大: {}MB", + String.format("%.2f", usagePercent), usedMB, maxMB); + PROTECTION_MODE.set(true); + } else if (!shouldProtect && PROTECTION_MODE.get()) { + log.info("内存恢复正常，退出保护模式 - 使用率: {}%, 已使用: {}MB", + String.format("%.2f", usagePercent), usedMB); + PROTECTION_MODE.set(false); + } + + return PROTECTION_MODE.get(); + } + } + + /** + * 写入内存不足响应 + */ + public static void writeMemoryInsufficientResponse(HttpServletResponse response) throws IOException { + response.setContentType("application/json;charset=UTF-8"); + response.setStatus(503); + response.getWriter().write( + "{\"code\":" + ResultCode.FAILED.getCode() + + ",\"msg\":\"服务器繁忙，请稍后再试\",\"data\":null}" + ); + } + + /** + * 获取当前内存状态（用于监控） + */ + public static MemoryStatus getMemoryStatus() { + MemoryMXBean memoryMXBean = ManagementFactory.getMemoryMXBean(); + MemoryUsage heapUsage = memoryMXBean.getHeapMemoryUsage(); + + MemoryStatus status = new MemoryStatus(); + status.setUsedMB(heapUsage.getUsed() / 1024 / 1024); + status.setCommittedMB(heapUsage.getCommitted() / 1024 / 1024); + status.setMaxMB(heapUsage.getMax() / 1024 / 1024); + status.setUsagePercent(status.getMaxMB() > 0 ? + (status.getUsedMB() * 100.0 / status.getMaxMB()) : 0); + status.setProtectionMode(PROTECTION_MODE.get()); + + return status; + } + + /** + * 内存状态信息 + */ + public static class MemoryStatus { + private long usedMB; + private long committedMB; + private long maxMB; + private double usagePercent; + private boolean protectionMode; + + public long getUsedMB() { return usedMB; } + public void setUsedMB(long usedMB) { this.usedMB = usedMB; } + public long getCommittedMB() { return committedMB; } + public void setCommittedMB(long committedMB) { this.committedMB = committedMB; } + public long getMaxMB() { return maxMB; } + public void setMaxMB(long maxMB) { this.maxMB = maxMB; } + public double getUsagePercent() { return usagePercent; } + public void setUsagePercent(double usagePercent) { this.usagePercent = usagePercent; } + public boolean isProtectionMode() { return protectionMode; } + public void setProtectionMode(boolean protectionMode) { this.protectionMode = protectionMode; } + } +}