同一個IP 訪問網站， 不能同時超過配置的最大值， 可以用來保護REST IP 或者DOS攻擊

1. /* 
2.  * Pprun's Public Domain. 
3.  */  
4. package org.pprun.common.security;  
5. import java.io.IOException;  
6. import java.util.Collections;  
7. import java.util.HashMap;  
8. import java.util.HashSet;  
9. import java.util.Map;  
10. import java.util.Set;  
11. import javax.servlet.Filter;  
12. import javax.servlet.FilterChain;  
13. import javax.servlet.FilterConfig;  
14. import javax.servlet.ServletException;  
15. import javax.servlet.ServletRequest;  
16. import javax.servlet.ServletResponse;  
17. import javax.servlet.http.HttpServletResponse;  
18. import org.slf4j.Logger;  
19. import org.slf4j.LoggerFactory;  
20. import org.springframework.beans.factory.annotation.Required;  
21. import org.springframework.jmx.export.annotation.ManagedAttribute;  
22. import org.springframework.scheduling.annotation.Scheduled;  
23. /** 
24.  * A filter to throttle only a limited number of requests from the same ip per second.
     
25.  * <p> 
26.  * Two parameters need to inject: 
27.  * <ol> 
28.  * <li>common.throttle.maxConcurrentRequests=10</li> 
29.  * <li>common.throttle.period=1000</li> 
30.  * </ol> 
31.  * <br /> 
32.  * If you use {@literal Spring} for {@lit DI}, it can be done as below in application {@literal web.xml}: 
33.  * <br /> 
34.  * {@code  
35.  *      <filter>   
36.  *          <description>A filter to throttle only a limited number of requests from the same ip per second.</description>   
37.  *          <filter-name>throttleFilter</filter-name>   
38.  *          <filter-class>org.springframework.web.filter.DelegatingFilterProxy</filter-class>   
39.  *      </filter> 
40.  * } 
41.  * <br /> 
42.  * And a bean with {@code id="throttleFilter"} should be defined in {@literal application.xml}. 
43.  * </p> 
44.  * @author <a href="mailto:[email protected]">pizhigang</a> 
45.  */  
46. //@Component  
47. public class ThrottleFilter implements Filter, ThrottleFilterMXBean {  
48.     private static Logger log = LoggerFactory.getLogger(ThrottleFilter.class);  
49.     // map(ip, requestCount)  
50.     private Map<String, Integer> ip2countCache = new HashMap<String, Integer>();  
51.     private Set<String> blackList = new HashSet<String>();  
52.     private int maxConcurrentRequests;  
53.     private static final long PERIOD = 1L; // second  
54.     private boolean enable = false;  
55.     @Override  
56.     public void init(FilterConfig config) throws ServletException {  
57.         // nothing  
58.     }  
59.     @Override  
60.     public void doFilter(ServletRequest request, ServletResponse response, FilterChain nextFilter) throws ServletException, IOException {  
61.         if (enable) {  
62.             final String ip = request.getRemoteAddr();  
63.             boolean isOverflow;  
64.             synchronized (this) {  
65.                 Integer count = ip2countCache.get(ip);  
66.                 if (count == null || count.intValue() == 0) {  
67.                     count = 0;  
68.                 }  
69.                 if (count < maxConcurrentRequests) {  
70.                     isOverflow = false;  
71.                     ip2countCache.put(ip, count + 1);  
72.                 } else {  
73.                     isOverflow = true;  
74.                     blackList.add(ip);  
75.                 }  
76.             }  
77.             if (isOverflow) {  
78.                 // block it  
79.                 log.info(" ip {} has reached the threshold {} in {} second, block it!", new Object[]{ip, maxConcurrentRequests, PERIOD});  
80.                 if (response instanceof HttpServletResponse) {  
81.                     ((HttpServletResponse) response).sendError(503, ip + " has too many concurrent requests per " + PERIOD + " second");  
82.                 }  
83.                 return;  
84.             }  
85.         } // else go ahead below  
86.         // every thing is okay, go ahead  
87.         nextFilter.doFilter(request, response);  
88.     }  
89.     // every 1 second  
90.     @Scheduled(fixedRate = PERIOD * 1000)  
91.     public void throttlingJob() {  
92.         if (enable) {  
93.             log.debug("Throttle Filter clean up job is running");  
94.             synchronized (ThrottleFilter.this) {  
95.                 for (Map.Entry<String, Integer> ip2count : ip2countCache.entrySet()) {  
96.                     Integer count = ip2count.getValue();  
97.                     String ip = ip2count.getKey();  
98.                     if (count == null || count <= 1) {  
99.                         ip2countCache.remove(ip);  
100.                     } else {  
101.                         if (count == maxConcurrentRequests) {  
102.                             // remove from blacklist  
103.                             log.info("Throttle Filter: removing {} from black list", ip);  
104.                             blackList.remove(ip);  
105.                         }  
106.                         ip2countCache.put(ip, count - 1);  
107.                     }  
108.                 }  
109.             }  
110.             log.debug("Throttle Filter clean up job is done");  
111.         }  
112.     }  
113.     /** 
114.      * Any cleanup for the filter. 
115.      */  
116.     @Override  
117.     public void destroy() {  
118.         log.warn("destorying Throttle Filter");  
119.     }  
120.     /** 
121.      * Sets the maximum number of concurrent requests for a single IP. 
122.      */  
123.     @Required  
124.     public void setMaxConcurrentRequests(int max) {  
125.         maxConcurrentRequests = max;  
126.     }  
127.     @ManagedAttribute(description = "is the throttle filter enable or not")  
128.     @Override  
129.     public boolean isEnable() {  
130.         return enable;  
131.     }  
132.     @Required  
133.     @Override  
134.     public void setEnable(boolean enable) {  
135.         log.info("set enable to {}", enable);  
136.         if (enable == false) {  
137.             log.info("Throttle filter will be disabled");  
138.         }  
139.         this.enable = enable;  
140.     }  
141.     public Set<String> getBlackList() {  
142.         // for exactly, it might need synchronized, but no hurt for snapshoot in one or severl seconds  
143.         return Collections.unmodifiableSet(blackList);  
144.     }  
145. }