retr0.id
A from-scratch atproto PDS implementation in Python (mirrors https://github.com/DavidBuchanan314/millipds)
add SSRF mitigation
+4
-2
src/millipds/__main__.py
+4
-2
src/millipds/__main__.py
···
67
67
from getpass import getpass
68
68
69
69
from docopt import docopt
70
-
import aiohttp
70
+
from .ssrf import get_ssrf_safe_client
71
+
71
72
72
73
import cbrrr
73
74
···
234
235
elif args["run"]:
235
236
236
237
async def run_service_with_client():
237
-
async with aiohttp.ClientSession() as client:
238
+
# TODO: option to use regular unsafe client for local dev testing
239
+
async with get_ssrf_safe_client() as client:
238
240
await service.run(
239
241
db=db,
240
242
client=client,
+37
src/millipds/ssrf.py
+37
src/millipds/ssrf.py
···
1
+
"""
2
+
This is a bit of a bodge, for now.
3
+
4
+
See https://github.com/aio-libs/aiohttp/discussions/10224 for the discussion
5
+
that led to this, and maybe a better solution in the future.
6
+
"""
7
+
8
+
import ipaddress
9
+
from aiohttp import TCPConnector, ClientSession
10
+
import aiohttp.connector
11
+
from aiohttp.resolver import DefaultResolver, AbstractResolver
12
+
13
+
# XXX: monkeypatch to force all hosts to go through the resolver
14
+
# (without this, bare IPs in the URL will bypass the resolver, where our SSRF check is)
15
+
aiohttp.connector.is_ip_address = lambda _: False
16
+
17
+
class SSRFException(ValueError):
18
+
pass
19
+
20
+
class SSRFSafeResolverWrapper(AbstractResolver):
21
+
def __init__(self, resolver: AbstractResolver):
22
+
self.resolver = resolver
23
+
24
+
async def resolve(self, host: str, port: int, family: int):
25
+
result = await self.resolver.resolve(host, port, family)
26
+
for host in result:
27
+
if ipaddress.ip_address(host["host"]).is_private:
28
+
raise SSRFException("Can't connect to private IP: " + host["host"])
29
+
return result
30
+
31
+
async def close(self) -> None:
32
+
await self.resolver.close()
33
+
34
+
def get_ssrf_safe_client() -> ClientSession:
35
+
resolver = SSRFSafeResolverWrapper(DefaultResolver())
36
+
connector = TCPConnector(resolver=resolver)
37
+
return ClientSession(connector=connector)