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}