1
- use super :: PortStrategy ;
1
+ use super :: { warning , PortStrategy } ;
2
2
3
3
mod socket_iterator;
4
4
use socket_iterator:: SocketIterator ;
@@ -11,6 +11,7 @@ use futures::stream::FuturesUnordered;
11
11
use std:: {
12
12
io:: ErrorKind ,
13
13
net:: { IpAddr , Shutdown , SocketAddr } ,
14
+ num:: NonZeroU8 ,
14
15
time:: Duration ,
15
16
} ;
16
17
@@ -26,6 +27,7 @@ pub struct Scanner {
26
27
ips : Vec < IpAddr > ,
27
28
batch_size : u16 ,
28
29
timeout : Duration ,
30
+ tries : NonZeroU8 ,
29
31
greppable : bool ,
30
32
port_strategy : PortStrategy ,
31
33
accessible : bool ,
@@ -36,13 +38,20 @@ impl Scanner {
36
38
ips : & [ IpAddr ] ,
37
39
batch_size : u16 ,
38
40
timeout : Duration ,
41
+ tries : u8 ,
39
42
greppable : bool ,
40
43
port_strategy : PortStrategy ,
41
44
accessible : bool ,
42
45
) -> Self {
43
46
Self {
44
47
batch_size,
45
48
timeout,
49
+ tries : if tries == 0 {
50
+ warning ! ( "tries is 0, corrected to 1!" ) ;
51
+ NonZeroU8 :: new ( 1 ) . unwrap ( )
52
+ } else {
53
+ NonZeroU8 :: new ( tries) . expect ( "tries cannot be 0!" )
54
+ } ,
46
55
greppable,
47
56
port_strategy,
48
57
ips : ips. iter ( ) . map ( |ip| ip. to_owned ( ) ) . collect ( ) ,
@@ -87,54 +96,67 @@ impl Scanner {
87
96
open_sockets
88
97
}
89
98
90
- /// Given a port , scan it.
99
+ /// Given a socket , scan it self.tries times .
91
100
/// Turns the address into a SocketAddr
92
101
/// Deals with the <result> type
93
102
/// If it experiences error ErrorKind::Other then too many files are open and it Panics!
94
- /// ese any other error, it returns the error in Result as a string
95
- /// If no errors occur, it returns the port number in Result to signify the port is open.
103
+ /// Else any other error, it returns the error in Result as a string
104
+ /// If no errors occur, it returns the port number in Result to signify the port is open.
96
105
/// This function mainly deals with the logic of Results handling.
97
106
/// # Example
98
107
///
99
- /// self.scan_port(10:u16 )
108
+ /// self.scan_socket(socket )
100
109
///
101
110
/// Note: `self` must contain `self.ip`.
102
111
async fn scan_socket ( & self , socket : SocketAddr ) -> io:: Result < SocketAddr > {
103
- match self . connect ( socket) . await {
104
- Ok ( x) => {
105
- debug ! (
106
- "Connection was successful, shutting down stream {}" ,
107
- & socket
108
- ) ;
109
- match x. shutdown ( Shutdown :: Both ) {
110
- Err ( e) => debug ! ( "Shutdown stream error {}" , & e) ,
111
- _ => { }
112
- }
113
- if !self . greppable {
114
- if self . accessible {
115
- println ! ( "Open {}" , socket. to_string( ) ) ;
116
- } else {
117
- println ! ( "Open {}" , socket. to_string( ) . purple( ) ) ;
112
+ let tries = self . tries . get ( ) ;
113
+
114
+ debug ! ( "self.tries: {}" , tries) ;
115
+
116
+ for nr_try in 1 ..=tries {
117
+ debug ! ( "Try number: {}" , nr_try) ;
118
+
119
+ match self . connect ( socket) . await {
120
+ Ok ( x) => {
121
+ debug ! (
122
+ "Connection was successful, shutting down stream {}" ,
123
+ & socket
124
+ ) ;
125
+ if let Err ( e) = x. shutdown ( Shutdown :: Both ) {
126
+ debug ! ( "Shutdown stream error {}" , & e) ;
127
+ }
128
+ if !self . greppable {
129
+ if self . accessible {
130
+ println ! ( "Open {}" , socket. to_string( ) ) ;
131
+ } else {
132
+ println ! ( "Open {}" , socket. to_string( ) . purple( ) ) ;
133
+ }
118
134
}
135
+
136
+ debug ! ( "Return Ok after {} tries" , nr_try) ;
137
+ return Ok ( socket) ;
119
138
}
139
+ Err ( e) => {
140
+ let error_string = e. to_string ( ) ;
120
141
121
- Ok ( socket)
122
- }
123
- Err ( e) => match e. kind ( ) {
124
- ErrorKind :: Other => {
125
- if e. to_string ( ) . contains ( "No route to host" )
126
- || e. to_string ( ) . contains ( "Network is unreachable" )
127
- {
128
- debug ! ( "Socket connect error: {} {}" , & e. to_string( ) , & socket) ;
129
- Err ( io:: Error :: new ( io:: ErrorKind :: Other , e. to_string ( ) ) )
130
- } else {
131
- debug ! ( "Socket connect error: {} {}" , & e. to_string( ) , & socket) ;
132
- panic ! ( "Too many open files. Please reduce batch size. The default is 5000. Try -b 2500." ) ;
142
+ if e. kind ( ) == ErrorKind :: Other {
143
+ debug ! ( "Socket connect error: {} {}" , & error_string, & socket) ;
144
+
145
+ if !error_string. contains ( "No route to host" )
146
+ && !error_string. contains ( "Network is unreachable" )
147
+ {
148
+ panic ! ( "Too many open files. Please reduce batch size. The default is 5000. Try -b 2500." ) ;
149
+ }
150
+ }
151
+
152
+ if nr_try == tries {
153
+ debug ! ( "Return Err after {} tries" , tries) ;
154
+ return Err ( io:: Error :: new ( io:: ErrorKind :: Other , error_string) ) ;
133
155
}
134
156
}
135
- _ => Err ( io:: Error :: new ( io:: ErrorKind :: Other , e. to_string ( ) ) ) ,
136
- } ,
157
+ } ;
137
158
}
159
+ unreachable ! ( ) ;
138
160
}
139
161
140
162
/// Performs the connection to the socket with timeout
@@ -174,7 +196,15 @@ mod tests {
174
196
end : 1_000 ,
175
197
} ;
176
198
let strategy = PortStrategy :: pick ( Some ( range) , None , ScanOrder :: Random ) ;
177
- let scanner = Scanner :: new ( & addrs, 10 , Duration :: from_millis ( 100 ) , true , strategy, true ) ;
199
+ let scanner = Scanner :: new (
200
+ & addrs,
201
+ 10 ,
202
+ Duration :: from_millis ( 100 ) ,
203
+ 1 ,
204
+ true ,
205
+ strategy,
206
+ true ,
207
+ ) ;
178
208
block_on ( scanner. run ( ) ) ;
179
209
// if the scan fails, it wouldn't be able to assert_eq! as it panicked!
180
210
assert_eq ! ( 1 , 1 ) ;
@@ -188,7 +218,15 @@ mod tests {
188
218
end : 1_000 ,
189
219
} ;
190
220
let strategy = PortStrategy :: pick ( Some ( range) , None , ScanOrder :: Random ) ;
191
- let scanner = Scanner :: new ( & addrs, 10 , Duration :: from_millis ( 100 ) , true , strategy, true ) ;
221
+ let scanner = Scanner :: new (
222
+ & addrs,
223
+ 10 ,
224
+ Duration :: from_millis ( 100 ) ,
225
+ 1 ,
226
+ true ,
227
+ strategy,
228
+ true ,
229
+ ) ;
192
230
block_on ( scanner. run ( ) ) ;
193
231
// if the scan fails, it wouldn't be able to assert_eq! as it panicked!
194
232
assert_eq ! ( 1 , 1 ) ;
@@ -201,7 +239,15 @@ mod tests {
201
239
end : 1_000 ,
202
240
} ;
203
241
let strategy = PortStrategy :: pick ( Some ( range) , None , ScanOrder :: Random ) ;
204
- let scanner = Scanner :: new ( & addrs, 10 , Duration :: from_millis ( 100 ) , true , strategy, true ) ;
242
+ let scanner = Scanner :: new (
243
+ & addrs,
244
+ 10 ,
245
+ Duration :: from_millis ( 100 ) ,
246
+ 1 ,
247
+ true ,
248
+ strategy,
249
+ true ,
250
+ ) ;
205
251
block_on ( scanner. run ( ) ) ;
206
252
assert_eq ! ( 1 , 1 ) ;
207
253
}
@@ -213,7 +259,15 @@ mod tests {
213
259
end : 445 ,
214
260
} ;
215
261
let strategy = PortStrategy :: pick ( Some ( range) , None , ScanOrder :: Random ) ;
216
- let scanner = Scanner :: new ( & addrs, 10 , Duration :: from_millis ( 100 ) , true , strategy, true ) ;
262
+ let scanner = Scanner :: new (
263
+ & addrs,
264
+ 10 ,
265
+ Duration :: from_millis ( 100 ) ,
266
+ 1 ,
267
+ true ,
268
+ strategy,
269
+ true ,
270
+ ) ;
217
271
block_on ( scanner. run ( ) ) ;
218
272
assert_eq ! ( 1 , 1 ) ;
219
273
}
@@ -228,7 +282,15 @@ mod tests {
228
282
end : 600 ,
229
283
} ;
230
284
let strategy = PortStrategy :: pick ( Some ( range) , None , ScanOrder :: Random ) ;
231
- let scanner = Scanner :: new ( & addrs, 10 , Duration :: from_millis ( 100 ) , true , strategy, true ) ;
285
+ let scanner = Scanner :: new (
286
+ & addrs,
287
+ 10 ,
288
+ Duration :: from_millis ( 100 ) ,
289
+ 1 ,
290
+ true ,
291
+ strategy,
292
+ true ,
293
+ ) ;
232
294
block_on ( scanner. run ( ) ) ;
233
295
assert_eq ! ( 1 , 1 ) ;
234
296
}
0 commit comments