Skip to content

Commit d05edaf

Browse files
timokoesslertimokoessler
committed
Prepare ServiceConfig and stats
1 parent 48776d1 commit d05edaf

File tree

9 files changed

+258
-62
lines changed

9 files changed

+258
-62
lines changed

‎library/agent/Agent.test.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,10 @@ wrap(fetch, "fetch", function mock() {
5858
pattern: "Bytespider",
5959
},
6060
],
61+
domains: [
62+
{ hostname: "example.com", mode: "block" },
63+
{ hostname: "aikido.dev", mode: "allow" },
64+
],
6165
} satisfies Response),
6266
};
6367
};
@@ -1110,6 +1114,9 @@ t.test("it fetches blocked lists", async () => {
11101114
t.same(agent.getConfig().isUserAgentBlocked("Mozilla/5.0 (compatible)"), {
11111115
blocked: false,
11121116
});
1117+
1118+
t.same(agent.getConfig().shouldBlockOutgoingRequest("example.com"), true);
1119+
t.same(agent.getConfig().shouldBlockOutgoingRequest("aikido.dev"), false);
11131120
});
11141121

11151122
t.test("it does not fetch blocked IPs if serverless", async () => {

‎library/agent/Agent.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -294,6 +294,12 @@ export class Agent {
294294
) {
295295
this.sendHeartbeatEveryMS = response.heartbeatIntervalInMS;
296296
}
297+
298+
if (typeof response.blockNewOutgoingRequests === "boolean") {
299+
this.serviceConfig.setBlockNewOutgoingRequests(
300+
response.blockNewOutgoingRequests
301+
);
302+
}
297303
}
298304
}
299305

@@ -402,13 +408,15 @@ export class Agent {
402408
monitoredIPAddresses,
403409
monitoredUserAgents,
404410
userAgentDetails,
411+
domains,
405412
} = await fetchBlockedLists(this.token);
406413
this.serviceConfig.updateBlockedIPAddresses(blockedIPAddresses);
407414
this.serviceConfig.updateBlockedUserAgents(blockedUserAgents);
408415
this.serviceConfig.updateAllowedIPAddresses(allowedIPAddresses);
409416
this.serviceConfig.updateMonitoredIPAddresses(monitoredIPAddresses);
410417
this.serviceConfig.updateMonitoredUserAgents(monitoredUserAgents);
411418
this.serviceConfig.updateUserAgentDetails(userAgentDetails);
419+
this.serviceConfig.updateDomains(domains);
412420
} catch (error: any) {
413421
console.error(`Aikido: Failed to update blocked lists: ${error.message}`);
414422
}

‎library/agent/Hostnames.test.ts‎

Lines changed: 73 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -6,41 +6,43 @@ t.test("it works", async () => {
66
t.same(hostnames.asArray(), []);
77

88
hostnames.add("aikido.dev", 443);
9-
t.same(hostnames.asArray(), [{ hostname: "aikido.dev", port: 443, hits: 1 }]);
9+
t.same(hostnames.asArray(), [
10+
{ hostname: "aikido.dev", port: 443, hits: 1, blockedHits: 0 },
11+
]);
1012

1113
hostnames.add("aikido.dev", 80);
1214
t.same(hostnames.asArray(), [
13-
{ hostname: "aikido.dev", port: 443, hits: 1 },
14-
{ hostname: "aikido.dev", port: 80, hits: 1 },
15+
{ hostname: "aikido.dev", port: 443, hits: 1, blockedHits: 0 },
16+
{ hostname: "aikido.dev", port: 80, hits: 1, blockedHits: 0 },
1517
]);
1618

1719
hostnames.add("google.com", 80);
1820
t.same(hostnames.asArray(), [
19-
{ hostname: "aikido.dev", port: 443, hits: 1 },
20-
{ hostname: "aikido.dev", port: 80, hits: 1 },
21-
{ hostname: "google.com", port: 80, hits: 1 },
21+
{ hostname: "aikido.dev", port: 443, hits: 1, blockedHits: 0 },
22+
{ hostname: "aikido.dev", port: 80, hits: 1, blockedHits: 0 },
23+
{ hostname: "google.com", port: 80, hits: 1, blockedHits: 0 },
2224
]);
2325

2426
hostnames.add("google.com", 0);
2527
hostnames.add("google.com", -1);
2628
t.same(hostnames.asArray(), [
27-
{ hostname: "aikido.dev", port: 443, hits: 1 },
28-
{ hostname: "aikido.dev", port: 80, hits: 1 },
29-
{ hostname: "google.com", port: 80, hits: 1 },
29+
{ hostname: "aikido.dev", port: 443, hits: 1, blockedHits: 0 },
30+
{ hostname: "aikido.dev", port: 80, hits: 1, blockedHits: 0 },
31+
{ hostname: "google.com", port: 80, hits: 1, blockedHits: 0 },
3032
]);
3133

3234
hostnames.add("github.com", 80);
3335
t.same(hostnames.asArray(), [
34-
{ hostname: "aikido.dev", port: 80, hits: 1 },
35-
{ hostname: "google.com", port: 80, hits: 1 },
36-
{ hostname: "github.com", port: 80, hits: 1 },
36+
{ hostname: "aikido.dev", port: 80, hits: 1, blockedHits: 0 },
37+
{ hostname: "google.com", port: 80, hits: 1, blockedHits: 0 },
38+
{ hostname: "github.com", port: 80, hits: 1, blockedHits: 0 },
3739
]);
3840

3941
hostnames.add("jetbrains.com", 80);
4042
t.same(hostnames.asArray(), [
41-
{ hostname: "google.com", port: 80, hits: 1 },
42-
{ hostname: "github.com", port: 80, hits: 1 },
43-
{ hostname: "jetbrains.com", port: 80, hits: 1 },
43+
{ hostname: "google.com", port: 80, hits: 1, blockedHits: 0 },
44+
{ hostname: "github.com", port: 80, hits: 1, blockedHits: 0 },
45+
{ hostname: "jetbrains.com", port: 80, hits: 1, blockedHits: 0 },
4446
]);
4547

4648
hostnames.clear();
@@ -53,30 +55,30 @@ t.test("it respects max size", async () => {
5355
hostnames.add("aikido.dev", 2);
5456

5557
t.same(hostnames.asArray(), [
56-
{ hostname: "aikido.dev", port: 1, hits: 1 },
57-
{ hostname: "aikido.dev", port: 2, hits: 1 },
58+
{ hostname: "aikido.dev", port: 1, hits: 1, blockedHits: 0 },
59+
{ hostname: "aikido.dev", port: 2, hits: 1, blockedHits: 0 },
5860
]);
5961

6062
hostnames.add("aikido.dev", 3);
6163
hostnames.add("aikido.dev", 4);
6264

6365
t.same(hostnames.asArray(), [
64-
{ hostname: "aikido.dev", port: 3, hits: 1 },
65-
{ hostname: "aikido.dev", port: 4, hits: 1 },
66+
{ hostname: "aikido.dev", port: 3, hits: 1, blockedHits: 0 },
67+
{ hostname: "aikido.dev", port: 4, hits: 1, blockedHits: 0 },
6668
]);
6769

6870
hostnames.add("google.com", 1);
6971

7072
t.same(hostnames.asArray(), [
71-
{ hostname: "aikido.dev", port: 4, hits: 1 },
72-
{ hostname: "google.com", port: 1, hits: 1 },
73+
{ hostname: "aikido.dev", port: 4, hits: 1, blockedHits: 0 },
74+
{ hostname: "google.com", port: 1, hits: 1, blockedHits: 0 },
7375
]);
7476

7577
hostnames.add("google.com", 2);
7678

7779
t.same(hostnames.asArray(), [
78-
{ hostname: "google.com", port: 1, hits: 1 },
79-
{ hostname: "google.com", port: 2, hits: 1 },
80+
{ hostname: "google.com", port: 1, hits: 1, blockedHits: 0 },
81+
{ hostname: "google.com", port: 2, hits: 1, blockedHits: 0 },
8082
]);
8183
});
8284

@@ -85,25 +87,64 @@ t.test("it tracks hits", async () => {
8587

8688
hostnames.add("aikido.dev", 443);
8789
hostnames.add("aikido.dev", 443);
88-
t.same(hostnames.asArray(), [{ hostname: "aikido.dev", port: 443, hits: 2 }]);
90+
t.same(hostnames.asArray(), [
91+
{ hostname: "aikido.dev", port: 443, hits: 2, blockedHits: 0 },
92+
]);
8993

9094
hostnames.add("aikido.dev", 80);
9195
t.same(hostnames.asArray(), [
92-
{ hostname: "aikido.dev", port: 443, hits: 2 },
93-
{ hostname: "aikido.dev", port: 80, hits: 1 },
96+
{ hostname: "aikido.dev", port: 443, hits: 2, blockedHits: 0 },
97+
{ hostname: "aikido.dev", port: 80, hits: 1, blockedHits: 0 },
9498
]);
9599

96100
hostnames.add("google.com", 80);
97101
t.same(hostnames.asArray(), [
98-
{ hostname: "aikido.dev", port: 443, hits: 2 },
99-
{ hostname: "aikido.dev", port: 80, hits: 1 },
100-
{ hostname: "google.com", port: 80, hits: 1 },
102+
{ hostname: "aikido.dev", port: 443, hits: 2, blockedHits: 0 },
103+
{ hostname: "aikido.dev", port: 80, hits: 1, blockedHits: 0 },
104+
{ hostname: "google.com", port: 80, hits: 1, blockedHits: 0 },
101105
]);
102106

103107
hostnames.add("aikido.dev", 443);
104108
t.same(hostnames.asArray(), [
105-
{ hostname: "aikido.dev", port: 443, hits: 3 },
106-
{ hostname: "aikido.dev", port: 80, hits: 1 },
107-
{ hostname: "google.com", port: 80, hits: 1 },
109+
{ hostname: "aikido.dev", port: 443, hits: 3, blockedHits: 0 },
110+
{ hostname: "aikido.dev", port: 80, hits: 1, blockedHits: 0 },
111+
{ hostname: "google.com", port: 80, hits: 1, blockedHits: 0 },
112+
]);
113+
});
114+
115+
t.test("it tracks blocked hits", async () => {
116+
const hostnames = new Hostnames(3);
117+
118+
hostnames.add("aikido.dev", 443, true);
119+
hostnames.add("aikido.dev", 443, true);
120+
t.same(hostnames.asArray(), [
121+
{ hostname: "aikido.dev", port: 443, hits: 2, blockedHits: 2 },
122+
]);
123+
124+
hostnames.add("aikido.dev", 80, true);
125+
t.same(hostnames.asArray(), [
126+
{ hostname: "aikido.dev", port: 443, hits: 2, blockedHits: 2 },
127+
{ hostname: "aikido.dev", port: 80, hits: 1, blockedHits: 1 },
128+
]);
129+
130+
hostnames.add("google.com", 80, true);
131+
t.same(hostnames.asArray(), [
132+
{ hostname: "aikido.dev", port: 443, hits: 2, blockedHits: 2 },
133+
{ hostname: "aikido.dev", port: 80, hits: 1, blockedHits: 1 },
134+
{ hostname: "google.com", port: 80, hits: 1, blockedHits: 1 },
135+
]);
136+
137+
hostnames.add("aikido.dev", 443, false);
138+
t.same(hostnames.asArray(), [
139+
{ hostname: "aikido.dev", port: 443, hits: 3, blockedHits: 2 },
140+
{ hostname: "aikido.dev", port: 80, hits: 1, blockedHits: 1 },
141+
{ hostname: "google.com", port: 80, hits: 1, blockedHits: 1 },
142+
]);
143+
144+
hostnames.add("aikido.dev", 443, false);
145+
t.same(hostnames.asArray(), [
146+
{ hostname: "aikido.dev", port: 443, hits: 4, blockedHits: 2 },
147+
{ hostname: "aikido.dev", port: 80, hits: 1, blockedHits: 1 },
148+
{ hostname: "google.com", port: 80, hits: 1, blockedHits: 1 },
108149
]);
109150
});

‎library/agent/Hostnames.ts‎

Lines changed: 33 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,48 @@
1-
type Ports = Map<number, number>;
1+
type Ports = Map<
2+
number,
3+
{
4+
hits: number;
5+
blockedHits: number;
6+
}
7+
>;
28

39
export class Hostnames {
410
private map: Map<string, Ports> = new Map();
511

612
constructor(private readonly maxEntries: number = 200) {}
713

8-
add(hostname: string, port: number) {
14+
add(hostname: string, port: number, blocked = false) {
915
if (port <= 0) {
1016
return;
1117
}
1218

1319
if (!this.map.has(hostname)) {
14-
this.map.set(hostname, new Map([[port, 1]]));
20+
this.map.set(
21+
hostname,
22+
new Map([
23+
[
24+
port,
25+
{
26+
hits: 1,
27+
blockedHits: blocked ? 1 : 0,
28+
},
29+
],
30+
])
31+
);
1532
} else {
1633
const ports = this.map.get(hostname) as Ports;
1734
if (!ports.has(port)) {
18-
ports.set(port, 1);
35+
ports.set(port, {
36+
hits: 1,
37+
blockedHits: blocked ? 1 : 0,
38+
});
1939
} else {
20-
ports.set(port, ports.get(port)! + 1);
40+
ports.set(port, {
41+
hits: ports.get(port)!.hits + 1,
42+
blockedHits: blocked
43+
? ports.get(port)!.blockedHits + 1
44+
: ports.get(port)!.blockedHits,
45+
});
2146
}
2247
}
2348

@@ -47,11 +72,12 @@ export class Hostnames {
4772

4873
asArray() {
4974
return Array.from(this.map.entries()).flatMap(([hostname, ports]) =>
50-
Array.from(ports.entries()).map(([port, hits]) => {
75+
Array.from(ports.entries()).map(([port, stats]) => {
5176
return {
5277
hostname,
5378
port,
54-
hits,
79+
hits: stats.hits,
80+
blockedHits: stats.blockedHits,
5581
};
5682
})
5783
);

‎library/agent/ServiceConfig.test.ts‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -406,3 +406,34 @@ t.test(
406406
t.same(config.getMatchingUserAgentKeys("googlebot"), []);
407407
}
408408
);
409+
410+
t.test("outbound request blocking", async (t) => {
411+
const config = new ServiceConfig([], 0, [], [], false, [], []);
412+
413+
t.same(config.shouldBlockOutgoingRequest("example.com"), false);
414+
415+
config.setBlockNewOutgoingRequests(true);
416+
t.same(config.shouldBlockOutgoingRequest("example.com"), true);
417+
418+
config.updateDomains([
419+
{ hostname: "example.com", mode: "allow" },
420+
{ hostname: "aikido.dev", mode: "block" },
421+
]);
422+
t.same(config.shouldBlockOutgoingRequest("example.com"), false);
423+
t.same(config.shouldBlockOutgoingRequest("aikido.dev"), true);
424+
t.same(config.shouldBlockOutgoingRequest("unknown.com"), true);
425+
426+
config.updateDomains([
427+
{ hostname: "example.com", mode: "block" },
428+
{ hostname: "aikido.dev", mode: "allow" },
429+
]);
430+
t.same(config.shouldBlockOutgoingRequest("example.com"), true);
431+
t.same(config.shouldBlockOutgoingRequest("aikido.dev"), false);
432+
t.same(config.shouldBlockOutgoingRequest("unknown.com"), true);
433+
434+
config.setBlockNewOutgoingRequests(false);
435+
436+
t.same(config.shouldBlockOutgoingRequest("example.com"), true);
437+
t.same(config.shouldBlockOutgoingRequest("aikido.dev"), false);
438+
t.same(config.shouldBlockOutgoingRequest("unknown.com"), false);
439+
});

‎library/agent/ServiceConfig.ts‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { IPMatcher } from "../helpers/ip-matcher/IPMatcher";
22
import { LimitedContext, matchEndpoints } from "../helpers/matchEndpoints";
33
import { isPrivateIP } from "../vulnerabilities/ssrf/isPrivateIP";
44
import type { Endpoint, EndpointConfig } from "./Config";
5-
import { IPList, UserAgentDetails } from "./api/fetchBlockedLists";
5+
import type { Domain, IPList, UserAgentDetails } from "./api/fetchBlockedLists";
66
import { safeCreateRegExp } from "./safeCreateRegExp";
77

88
export class ServiceConfig {
@@ -27,6 +27,9 @@ export class ServiceConfig {
2727
private monitoredUserAgentRegex: RegExp | undefined;
2828
private userAgentDetails: { pattern: RegExp; key: string }[] = [];
2929

30+
private blockNewOutgoingRequests = false;
31+
private domains = new Map<string, Domain["mode"]>();
32+
3033
constructor(
3134
endpoints: EndpointConfig[],
3235
private lastUpdatedAt: number,
@@ -278,4 +281,24 @@ export class ServiceConfig {
278281
hasReceivedAnyStats() {
279282
return this.receivedAnyStats;
280283
}
284+
285+
setBlockNewOutgoingRequests(block: boolean) {
286+
this.blockNewOutgoingRequests = block;
287+
}
288+
289+
updateDomains(domains: Domain[]) {
290+
this.domains = new Map(domains.map((i) => [i.hostname, i.mode]));
291+
}
292+
293+
shouldBlockOutgoingRequest(hostname: string): boolean {
294+
const mode = this.domains.get(hostname);
295+
296+
if (this.blockNewOutgoingRequests) {
297+
// Only allow outgoing requests if the mode is "allow"
298+
return mode !== "allow";
299+
}
300+
301+
// Only block outgoing requests if the mode is "block"
302+
return mode === "block";
303+
}
281304
}

‎library/sinks/Fetch.test.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -90,14 +90,14 @@ t.test(
9090
await fetch("http://app.aikido.dev");
9191

9292
t.same(agent.getHostnames().asArray(), [
93-
{ hostname: "app.aikido.dev", port: 80, hits: 1 },
93+
{ hostname: "app.aikido.dev", port: 80, hits: 1, blockedHits: 0 },
9494
]);
9595
agent.getHostnames().clear();
9696

9797
await fetch(new URL("https://app.aikido.dev"));
9898

9999
t.same(agent.getHostnames().asArray(), [
100-
{ hostname: "app.aikido.dev", port: 443, hits: 1 },
100+
{ hostname: "app.aikido.dev", port: 443, hits: 1, blockedHits: 0 },
101101
]);
102102
agent.getHostnames().clear();
103103

@@ -110,7 +110,7 @@ t.test(
110110
await fetch(new Request("https://app.aikido.dev"));
111111

112112
t.same(agent.getHostnames().asArray(), [
113-
{ hostname: "app.aikido.dev", port: 443, hits: 1 },
113+
{ hostname: "app.aikido.dev", port: 443, hits: 1, blockedHits: 0 },
114114
]);
115115

116116
agent.getHostnames().clear();

0 commit comments

Comments
 (0)