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
|
package com.it_jaros.jscanner.scan;
import com.it_jaros.jscanner.scan.domain.ScanResult;
import com.it_jaros.jscanner.scan.engine.*;
import java.time.Duration;
import java.util.concurrent.*;
import java.util.function.Consumer;
public class Scanner implements AutoCloseable {
private final CancelledToken cancelledToken = new CancelledToken();
private final ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor();
private final ScanOptions scanOptions;
private final Semaphore socketLimit;
public Scanner(ScanOptions options) {
this.scanOptions = options;
this.socketLimit = new Semaphore(scanOptions.socketLimit());
}
/**
* Starts a given scan.
*
* @param scan
*/
public void runScan(final Scan scan, final Consumer<ScanResult> consumer) {
if (scan == null) {
throw new IllegalArgumentException("Scan argument cannot be null");
}
scan.start();
long delayInNanos = TimeUnit.MILLISECONDS.toNanos(Math.max(0, scanOptions.delayInMillis()));
ScanExecutionContext context = new ScanExecutionContext(
scanOptions,
executor,
socketLimit,
cancelledToken,
scan,
new PortScanRateLimiter(cancelledToken, delayInNanos)
);
// start producer thread
scan.producerStart();
final ProducerState<ScanResult> state = new ProducerThread(context).startProducer(
scan.getHosts().iterator(),
scanOptions.maxHostsLimit(),
host -> new ScanHostTask(scan, host, context)
);
// the main thread is the consumer
// Let the consumer run as long as the producer runs
// or if still tasks are pending in pipeline
// we do not listen to canceled here because we want
// all results (also partial) collected for the consumer
// with whatever is there already
while (state.running().get() || state.inPipeline().get() > 0) {
try {
PollState<ScanResult> poll = getHostResult(state);
if (poll instanceof PollState.Success<ScanResult>(ScanResult value)) {
consumer.accept(value);
} else if (poll instanceof PollState.Failure<ScanResult>(Throwable error)) {
System.err.printf("runScan(): ScanHostTask() failed with error %s -> %s%n", error.getClass().getSimpleName(), error.getMessage());
}
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
}
}
scan.producerStop();
scan.stop();
}
private PollState<ScanResult> getHostResult(ProducerState<ScanResult> state) throws InterruptedException {
Future<ScanResult> finishedHost = state.completionService().poll(ProducerThread.pollInterval.toMillis(), TimeUnit.MILLISECONDS);
if (finishedHost == null) {
return new PollState.Unavailable<>();
}
try {
ScanResult result = finishedHost.get();
return new PollState.Success<>(result);
} catch (ExecutionException e) {
return new PollState.Failure<>(e.getCause());
} finally {
state.activeWorkers().release();
state.inPipeline().decrementAndGet();
}
}
public boolean awaitTermination(Duration duration) throws InterruptedException {
return executor.awaitTermination(duration.toMillis(), TimeUnit.MILLISECONDS);
}
public void cancel() {
this.cancelledToken.signal();
executor.shutdown();
}
public void cancelNow() {
this.cancelledToken.signal();
executor.shutdownNow();
}
@Override
public void close() throws Exception {
cancel();
}
}
|
