and even more fixes

This commit is contained in:
Sarah 2021-10-17 03:09:59 +02:00
parent d112cf8796
commit 566e8a070a
No known key found for this signature in database
GPG key ID: 708F7ACE058F0186
2 changed files with 16 additions and 5 deletions

View file

@ -16,7 +16,8 @@ async def _setup_stores(local_port: int):
global l_access, r_access global l_access, r_access
async with local() as l: async with local() as l:
l_access = PrefixStore("local/nar", l) l_access = PrefixStore("local/nar", l)
async with remote(l_access, local_port, "0.0.0.0") as r: lp = PrefixStore("local", l)
async with remote(lp, local_port, "0.0.0.0", lp.prefix) as r:
r_access = PrefixStore("remote", r) r_access = PrefixStore("remote", r)
yield yield
setup_store = _setup_stores(12304) setup_store = _setup_stores(12304)

View file

@ -41,13 +41,15 @@ class DiscoveryProtocol(asyncio.DatagramProtocol, Store):
store: Store store: Store
session: aiohttp.ClientSession session: aiohttp.ClientSession
local_port: int local_port: int
prefix: str
def __init__(self, store: Store, session: aiohttp.ClientSession, local_port: int): def __init__(self, store: Store, session: aiohttp.ClientSession, local_port: int, prefix: str):
self.idx = 0 self.idx = 0
self.waiters = {} self.waiters = {}
self.store = store self.store = store
self.session = session self.session = session
self.local_port = local_port self.local_port = local_port
self.prefix = prefix
def connection_made(self, transport): def connection_made(self, transport):
self.transport = transport self.transport = transport
@ -82,7 +84,15 @@ class DiscoveryProtocol(asyncio.DatagramProtocol, Store):
if narinfo is None: if narinfo is None:
return return
self.transport.sendto(b"\x01" + data[1:5] + self.local_port.to_bytes(4, "big") + narinfo.url.encode("utf-8"), addr) self.transport.sendto(b"".join([
b"\x01",
data[1:5],
self.local_port.to_bytes(4, "big"),
self.prefix.encode("utf-8"),
b"/",
hsh.encode("utf-8"),
b".narinfo"
]), addr)
async def narinfo(self, hsh: str) -> t.Optional[NarInfo]: async def narinfo(self, hsh: str) -> t.Optional[NarInfo]:
fut = asyncio.get_running_loop().create_future() fut = asyncio.get_running_loop().create_future()
@ -120,11 +130,11 @@ class DiscoveryProtocol(asyncio.DatagramProtocol, Store):
@contextlib.asynccontextmanager @contextlib.asynccontextmanager
async def remote(store: Store, local_port: int, local_addr: str="0.0.0.0"): async def remote(store: Store, local_port: int, local_addr: str="0.0.0.0", prefix: str="local"):
protocol: DiscoveryProtocol protocol: DiscoveryProtocol
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
_, protocol = await asyncio.get_running_loop().create_datagram_endpoint( _, protocol = await asyncio.get_running_loop().create_datagram_endpoint(
lambda: DiscoveryProtocol(store, session, local_port), lambda: DiscoveryProtocol(store, session, local_port, prefix),
local_addr=(local_addr, local_port), local_addr=(local_addr, local_port),
family=socket.AF_INET, family=socket.AF_INET,
allow_broadcast=True allow_broadcast=True