@@ -2,9 +2,12 @@ use std::collections::HashSet;
22use std:: net:: IpAddr ;
33use std:: sync:: Arc ;
44
5+ use ipnet:: IpNet ;
6+
57#[ derive( Clone , Default , Debug ) ]
68pub struct IpBlacklist {
79 ips : Arc < HashSet < IpAddr > > ,
10+ networks : Arc < Vec < IpNet > > ,
811 is_loading : bool ,
912}
1013
@@ -13,24 +16,32 @@ impl IpBlacklist {
1316 let Ok ( buf) = tokio:: fs:: read_to_string ( & path) . await else {
1417 return IpBlacklist :: default ( ) ;
1518 } ;
16- let mut set = HashSet :: new ( ) ;
19+ let mut ips = HashSet :: new ( ) ;
20+ let mut networks = Vec :: new ( ) ;
1721 for line in buf. lines ( ) {
18- if let Ok ( ip) = line. trim ( ) . parse :: < IpAddr > ( ) {
19- set. insert ( ip) ;
22+ let Some ( line) = clean_blacklist_line ( line) else {
23+ continue ;
24+ } ;
25+
26+ if let Ok ( ip) = line. parse :: < IpAddr > ( ) {
27+ ips. insert ( ip) ;
28+ } else if let Ok ( network) = line. parse :: < IpNet > ( ) {
29+ networks. push ( network) ;
2030 }
2131 }
2232 IpBlacklist {
23- ips : Arc :: new ( set) ,
33+ ips : Arc :: new ( ips) ,
34+ networks : Arc :: new ( networks) ,
2435 is_loading : false ,
2536 }
2637 }
2738
2839 pub fn contains ( & self , ip : & IpAddr ) -> bool {
29- self . ips . contains ( ip)
40+ self . ips . contains ( ip) || self . networks . iter ( ) . any ( |network| network . contains ( ip ) )
3041 }
3142
3243 pub fn is_invalid ( & self ) -> bool {
33- self . ips . is_empty ( ) && !self . is_loading
44+ self . ips . is_empty ( ) && self . networks . is_empty ( ) && !self . is_loading
3445 }
3546
3647 pub fn is_loading ( & self ) -> bool {
@@ -42,6 +53,20 @@ impl IpBlacklist {
4253 }
4354}
4455
56+ fn clean_blacklist_line ( line : & str ) -> Option < & str > {
57+ let line = line. trim ( ) ;
58+
59+ if line. is_empty ( ) || line. starts_with ( '#' ) || line. starts_with ( ';' ) {
60+ return None ;
61+ }
62+
63+ let line = line. split ( ';' ) . next ( ) . unwrap_or ( line) ;
64+ let line = line. split ( '#' ) . next ( ) . unwrap_or ( line) ;
65+ let line = line. trim ( ) ;
66+
67+ if line. is_empty ( ) { None } else { Some ( line) }
68+ }
69+
4570#[ cfg( test) ]
4671mod tests {
4772 use super :: * ;
@@ -55,6 +80,7 @@ mod tests {
5580 assert ! ( !blacklist. is_invalid( ) ) ;
5681 assert ! ( !blacklist. is_loading( ) ) ;
5782 assert_eq ! ( blacklist. ips. len( ) , 4 ) ;
83+ assert_eq ! ( blacklist. networks. len( ) , 0 ) ;
5884
5985 assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 8 , 8 , 8 , 8 ) ) ) ) ;
6086 assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 1 , 2 , 3 , 255 ) ) ) ) ;
@@ -75,10 +101,86 @@ mod tests {
75101 assert ! ( blacklist. is_invalid( ) ) ;
76102 assert ! ( !blacklist. is_loading( ) ) ;
77103 assert_eq ! ( blacklist. ips. len( ) , 0 ) ;
104+ assert_eq ! ( blacklist. networks. len( ) , 0 ) ;
78105
79106 assert ! ( !blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 8 , 8 , 8 , 8 ) ) ) ) ;
80107 assert ! ( !blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 0 , 0 , 0 , 0 ) ) ) ) ;
81108 assert ! ( !blacklist. contains( & "::123" . parse:: <IpAddr >( ) . unwrap( ) ) ) ;
82109 assert ! ( !blacklist. contains( & "::" . parse:: <IpAddr >( ) . unwrap( ) ) ) ;
83110 }
111+
112+ #[ tokio:: test]
113+ async fn test_ip_blacklist_valid_with_cidr ( ) {
114+ let blacklist =
115+ IpBlacklist :: from_file ( "resources/test/ip_blacklist_valid_with_cidr.txt" . to_string ( ) )
116+ . await ;
117+
118+ assert ! ( !blacklist. is_invalid( ) ) ;
119+ assert ! ( !blacklist. is_loading( ) ) ;
120+ assert_eq ! ( blacklist. ips. len( ) , 2 ) ;
121+ assert_eq ! ( blacklist. networks. len( ) , 4 ) ;
122+
123+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 8 , 8 , 8 , 8 ) ) ) ) ;
124+ assert ! ( blacklist. contains( & "2001:db8::1" . parse:: <IpAddr >( ) . unwrap( ) ) ) ;
125+ assert ! ( !blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 8 , 8 , 8 , 9 ) ) ) ) ;
126+ assert ! ( !blacklist. contains( & "2001:db8::2" . parse:: <IpAddr >( ) . unwrap( ) ) ) ;
127+
128+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 1 , 2 , 3 , 1 ) ) ) ) ;
129+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 1 , 2 , 3 , 255 ) ) ) ) ;
130+ assert ! ( !blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 1 , 2 , 5 , 1 ) ) ) ) ;
131+
132+ assert ! ( blacklist. contains( & "2001:db9::1" . parse:: <IpAddr >( ) . unwrap( ) ) ) ;
133+ assert ! ( blacklist. contains( & "2001:db9:ffff::1" . parse:: <IpAddr >( ) . unwrap( ) ) ) ;
134+ assert ! ( !blacklist. contains( & "2001:dba::1" . parse:: <IpAddr >( ) . unwrap( ) ) ) ;
135+
136+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 5 , 6 , 7 , 10 ) ) ) ) ;
137+ assert ! ( !blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 5 , 6 , 8 , 10 ) ) ) ) ;
138+
139+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 9 , 9 , 9 , 9 ) ) ) ) ;
140+ assert ! ( !blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 9 , 9 , 10 , 9 ) ) ) ) ;
141+ }
142+
143+ #[ tokio:: test]
144+ async fn test_ip_blacklist_valid_with_cidr_only ( ) {
145+ let blacklist =
146+ IpBlacklist :: from_file ( "resources/test/ip_blacklist_valid_cidr_only.txt" . to_string ( ) )
147+ . await ;
148+
149+ assert ! ( !blacklist. is_invalid( ) ) ;
150+ assert ! ( !blacklist. is_loading( ) ) ;
151+ assert_eq ! ( blacklist. ips. len( ) , 0 ) ;
152+ assert_eq ! ( blacklist. networks. len( ) , 1 ) ;
153+
154+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 1 , 2 , 3 , 1 ) ) ) ) ;
155+ assert ! ( !blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 1 , 2 , 4 , 1 ) ) ) ) ;
156+ }
157+
158+ #[ tokio:: test]
159+ async fn test_ip_blacklist_real_cidr_ranges ( ) {
160+ let blacklist =
161+ IpBlacklist :: from_file ( "resources/test/ip_blacklist_real_cidr_ranges.txt" . to_string ( ) )
162+ . await ;
163+
164+ assert ! ( !blacklist. is_invalid( ) ) ;
165+ assert_eq ! ( blacklist. ips. len( ) , 0 ) ;
166+ assert_eq ! ( blacklist. networks. len( ) , 6 ) ;
167+
168+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 186 , 20 , 0 ) ) ) ) ;
169+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 186 , 23 , 255 ) ) ) ) ;
170+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 186 , 24 , 0 ) ) ) ) ;
171+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 186 , 31 , 255 ) ) ) ) ;
172+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 186 , 224 , 0 ) ) ) ) ;
173+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 186 , 231 , 255 ) ) ) ) ;
174+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 186 , 232 , 0 ) ) ) ) ;
175+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 186 , 235 , 255 ) ) ) ) ;
176+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 186 , 236 , 0 ) ) ) ) ;
177+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 186 , 236 , 255 ) ) ) ) ;
178+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 233 , 156 , 0 ) ) ) ) ;
179+ assert ! ( blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 233 , 159 , 255 ) ) ) ) ;
180+
181+ assert ! ( !blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 186 , 19 , 255 ) ) ) ) ;
182+ assert ! ( !blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 186 , 32 , 0 ) ) ) ) ;
183+ assert ! ( !blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 186 , 237 , 0 ) ) ) ) ;
184+ assert ! ( !blacklist. contains( & IpAddr :: V4 ( Ipv4Addr :: new( 209 , 233 , 160 , 0 ) ) ) ) ;
185+ }
84186}
0 commit comments