luoyb
2021-02-02 a6c5984b2347a368e4e9ea7dbfdb95d199d6d4f4
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
package cc.mrbird.febs.gateway.enhance.service.impl;
 
import cc.mrbird.febs.common.redis.service.RedisService;
import cc.mrbird.febs.gateway.enhance.entity.BlackList;
import cc.mrbird.febs.gateway.enhance.entity.RateLimitRule;
import cc.mrbird.febs.gateway.enhance.service.RouteEnhanceCacheService;
import cc.mrbird.febs.gateway.enhance.utils.RouteEnhanceCacheUtil;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
 
import java.util.Set;
 
/**
 * @author MrBird
 */
@Slf4j
@Service
public class RouteEnhanceCacheServiceImpl implements RouteEnhanceCacheService {
 
    private RedisService redisService;
 
    @Autowired(required = false)
    public void setRedisService(RedisService redisService) {
        this.redisService = redisService;
    }
 
    @Override
    public void saveAllBlackList(Flux<BlackList> blackList) {
        blackList.subscribe(b -> {
            String key = StringUtils.isNotBlank(b.getIp()) ?
                    RouteEnhanceCacheUtil.getBlackListCacheKey(b.getIp()) :
                    RouteEnhanceCacheUtil.getBlackListCacheKey();
            String value = JSONObject.toJSONString(b);
            redisService.sSet(key, value);
        });
        log.info("Cache blacklist into redis >>>");
    }
 
    @Override
    public void saveBlackList(BlackList blackList) {
        String key = StringUtils.isNotBlank(blackList.getIp()) ?
                RouteEnhanceCacheUtil.getBlackListCacheKey(blackList.getIp()) :
                RouteEnhanceCacheUtil.getBlackListCacheKey();
        redisService.sSet(key, JSONObject.toJSONString(blackList));
    }
 
    @Override
    public Set<Object> getBlackList(String ip) {
        String key = RouteEnhanceCacheUtil.getBlackListCacheKey(ip);
        return redisService.sGet(key);
    }
 
    @Override
    public Set<Object> getBlackList() {
        String key = RouteEnhanceCacheUtil.getBlackListCacheKey();
        return redisService.sGet(key);
    }
 
    @Override
    public void removeBlackList(BlackList blackList) {
        String key = StringUtils.isNotBlank(blackList.getIp()) ?
                RouteEnhanceCacheUtil.getBlackListCacheKey(blackList.getIp()) :
                RouteEnhanceCacheUtil.getBlackListCacheKey();
        redisService.setRemove(key, JSONObject.toJSONString(blackList));
    }
 
    @Override
    public void saveAllRateLimitRules(Flux<RateLimitRule> rateLimitRules) {
        rateLimitRules.subscribe(r -> {
            String key = RouteEnhanceCacheUtil.getRateLimitCacheKey(r.getRequestUri(), r.getRequestMethod());
            String value = JSONObject.toJSONString(r);
            redisService.set(key, value);
        });
        log.info("Cache rate limit rules into redis >>>");
    }
 
    @Override
    public void saveRateLimitRule(RateLimitRule rateLimitRule) {
        String key = RouteEnhanceCacheUtil.getRateLimitCacheKey(rateLimitRule.getRequestUri(), rateLimitRule.getRequestMethod());
        redisService.set(key, JSONObject.toJSONString(rateLimitRule));
    }
 
 
    @Override
    public Object getRateLimitRule(String uri, String method) {
        String key = RouteEnhanceCacheUtil.getRateLimitCacheKey(uri, method);
        return redisService.get(key);
    }
 
    @Override
    public int getCurrentRequestCount(String uri, String ip) {
        String key = RouteEnhanceCacheUtil.getRateLimitCountKey(uri, ip);
        return redisService.hasKey(key) ? (int) redisService.get(key) : 0;
    }
 
    @Override
    public void removeRateLimitRule(RateLimitRule rateLimitRule) {
        String key = RouteEnhanceCacheUtil.getRateLimitCacheKey(rateLimitRule.getRequestUri(), rateLimitRule.getRequestMethod());
        redisService.del(key);
    }
 
    @Override
    public void setCurrentRequestCount(String uri, String ip, Long time) {
        String key = RouteEnhanceCacheUtil.getRateLimitCountKey(uri, ip);
        redisService.set(key, 1, time);
    }
 
    @Override
    public void incrCurrentRequestCount(String uri, String ip) {
        String key = RouteEnhanceCacheUtil.getRateLimitCountKey(uri, ip);
        redisService.incr(key, 1L);
    }
}