001/** 002 * Copyright (c) 2015-2022, Michael Yang 杨福海 (fuhai999@gmail.com). 003 * <p> 004 * Licensed under the Apache License, Version 2.0 (the "License"); 005 * you may not use this file except in compliance with the License. 006 * You may obtain a copy of the License at 007 * <p> 008 * http://www.apache.org/licenses/LICENSE-2.0 009 * <p> 010 * Unless required by applicable law or agreed to in writing, software 011 * distributed under the License is distributed on an "AS IS" BASIS, 012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 013 * See the License for the specific language governing permissions and 014 * limitations under the License. 015 */ 016package io.jboot.components.limiter; 017 018 019import com.github.benmanes.caffeine.cache.Cache; 020import com.github.benmanes.caffeine.cache.Caffeine; 021import com.google.common.util.concurrent.RateLimiter; 022import com.jfinal.aop.Invocation; 023import io.jboot.Jboot; 024import io.jboot.utils.ClassUtil; 025import io.jboot.utils.StrUtil; 026 027import java.util.*; 028import java.util.concurrent.Semaphore; 029import java.util.concurrent.TimeUnit; 030import java.util.regex.Pattern; 031 032public class LimiterManager { 033 034 private HashSet<LimitConfigBean> limitConfigBeans = new HashSet<>(); 035 private HashSet<String> ipWhitelist = new HashSet<>(); 036 037 private LimitConfig limitConfig = Jboot.config(LimitConfig.class); 038 039 040 private Cache<String, Semaphore> semaphoreCache = Caffeine.newBuilder() 041 .expireAfterAccess(10, TimeUnit.MINUTES) 042 .expireAfterWrite(10, TimeUnit.MINUTES) 043 .build(); 044 045 046 private Cache<String, RateLimiter> rateLimiterCache = Caffeine.newBuilder() 047 .expireAfterAccess(10, TimeUnit.MINUTES) 048 .expireAfterWrite(10, TimeUnit.MINUTES) 049 .build(); 050 051 052 private LimitFallbackProcesser fallbackProcesser; 053 054 private static LimiterManager me = new LimiterManager(); 055 056 private LimiterManager() { 057 } 058 059 public static LimiterManager me() { 060 return me; 061 } 062 063 public void init() { 064 doInitFallbackProcesser(); 065 doParseConfig(); 066 067 Set<String> ips = StrUtil.splitToSetByComma(limitConfig.getIpWhitelist()); 068 if (ips != null) { 069 ipWhitelist.addAll(ips); 070 } 071 } 072 073 private void doInitFallbackProcesser() { 074 LimitConfig config = Jboot.config(LimitConfig.class); 075 if (StrUtil.isBlank(config.getFallbackProcesser())) { 076 this.fallbackProcesser = new LimitFallbackProcesserDefault(); 077 } else { 078 this.fallbackProcesser = Objects.requireNonNull(ClassUtil.newInstance(config.getFallbackProcesser()), 079 "can not newInstance class for " + config.getFallbackProcesser()); 080 } 081 } 082 083 /** 084 * 解析用户配置 085 */ 086 private void doParseConfig() { 087 088 LimitConfig config = Jboot.config(LimitConfig.class); 089 090 if (!isEnable()) { 091 return; 092 } 093 094 String rule = config.getRule(); 095 if (StrUtil.isBlank(rule)) { 096 return; 097 } 098 099 String[] rules = rule.split(","); 100 for (String r : rules) { 101 String[] confs = r.split(":"); 102 if (confs == null || confs.length != 3) { 103 continue; 104 } 105 106 String packageOrTarget = confs[0]; 107 String type = confs[1]; 108 String rate = confs[2]; 109 110 if (!ensureLegal(packageOrTarget, type, rate.trim())) { 111 continue; 112 } 113 114 packageOrTarget = packageOrTarget.replace(".", "\\.") 115 .replace("(", "\\(") 116 .replace(")", "\\)") 117 .replace("*", ".*"); 118 119 limitConfigBeans.add(new LimitConfigBean(packageOrTarget.trim(),type.trim(), Integer.valueOf(rate.trim()))); 120 } 121 } 122 123 124 /** 125 * 匹配用户配置 126 * 127 * @param packageOrTarget 128 * @return 129 */ 130 public LimitConfigBean matchConfig(String packageOrTarget) { 131 132 if (!isEnable() || limitConfigBeans.isEmpty()) { 133 return null; 134 } 135 136 for (LimitConfigBean value : limitConfigBeans) { 137 if (value.isMatched(packageOrTarget)){ 138 return value; 139 } 140 } 141 142 return null; 143 } 144 145 146 public boolean isInIpWhitelist(String ip){ 147 return !ipWhitelist.isEmpty() && ipWhitelist.contains(ip); 148 } 149 150 public RateLimiter getOrCreateRateLimiter(String resKey, int rate) { 151 return rateLimiterCache.get(resKey, s -> RateLimiter.create(rate)); 152 } 153 154 155 public Semaphore getOrCreateSemaphore(String resKey, int rate) { 156 return semaphoreCache.get(resKey, s -> new Semaphore(rate)); 157 } 158 159 public HashSet<LimitConfigBean> getLimitConfigBeans() { 160 return limitConfigBeans; 161 } 162 163 /** 164 * 确保配置合法 165 * 166 * @param packageOrTarget 167 * @param type 168 * @param rate 169 * @return 170 */ 171 private boolean ensureLegal(String packageOrTarget, String type, String rate) { 172 if (StrUtil.isBlank(packageOrTarget)) { 173 return false; 174 } 175 176 if (!LimitType.types.contains(type)) { 177 return false; 178 } 179 180 if (!StrUtil.isNumeric(rate)) { 181 return false; 182 } 183 184 return true; 185 } 186 187 public boolean isEnable() { 188 return limitConfig.isEnable(); 189 } 190 191 public void processFallback(String resource, String fallback, Invocation inv) { 192 fallbackProcesser.process(resource, fallback, inv); 193 } 194 195 public static class LimitConfigBean { 196 197 private String packageOrTarget; 198 private String type; 199 private int rate; 200 private Pattern pattern; 201 202 203 public LimitConfigBean(String packageOrTarget, String type, int rate) { 204 this.packageOrTarget = packageOrTarget; 205 this.type = type; 206 this.rate = rate; 207 this.pattern = Pattern.compile(packageOrTarget); 208 } 209 210 public String getPackageOrTarget() { 211 return packageOrTarget; 212 } 213 214 public void setPackageOrTarget(String packageOrTarget) { 215 this.packageOrTarget = packageOrTarget; 216 } 217 218 public String getType() { 219 return type; 220 } 221 222 public void setType(String type) { 223 this.type = type; 224 } 225 226 public int getRate() { 227 return rate; 228 } 229 230 public void setRate(int rate) { 231 this.rate = rate; 232 } 233 234 public boolean isMatched(String packageOrTarget){ 235 return pattern.matcher(packageOrTarget).matches(); 236 } 237 } 238}