diff --git a/README.rst b/README.rst index 3e0671c..6049c3e 100755 --- a/README.rst +++ b/README.rst @@ -45,14 +45,14 @@ Shortcut -------- Manual usage is too board. So I make shortcut to use easily. -You just replace ``aiohttp.ClientSession`` to ``aiohttp_doh.ClientSession``. +You just replace ``aiohttp.ClientSession`` to ``aiohttp_doh.DNSOverHTTPSClientSession``. .. code-block:: python3 - from aiohttp_doh import ClientSession + from aiohttp_doh import DNSOverHTTPSClientSession async def main(): - async with ClientSession() as session: + async with DNSOverHTTPSClientSession() as session: async with session.get('http://example.com') as resp: data = await resp.text() @@ -68,7 +68,7 @@ endpoints List of str. DNS over HTTPS endpoints. Shortcut use `'https://dns.google.com/resolve'` and `'https://cloudflare-dns.com/dns-query'` both in default. - You can also use others instead. + You can also use others instead, make sure DNS endpoints support JSON response format (application/dns-json) json_loads Function for loads json. default is Python builtin json module's one. diff --git a/aiohttp_doh.py b/aiohttp_doh.py index 165646d..bbf69a3 100644 --- a/aiohttp_doh.py +++ b/aiohttp_doh.py @@ -9,7 +9,7 @@ from aiohttp.connector import TCPConnector from aiohttp.resolver import DefaultResolver -__all__ = 'ClientSession', 'DNSOverHTTPSResolver', 'RecordType' +__all__ = 'DNSOverHTTPSClientSession', 'DNSOverHTTPSResolver', 'RecordType' class RecordType(enum.Enum): @@ -22,13 +22,7 @@ class RecordType(enum.Enum): class DNSOverHTTPSResolver(AbstractResolver): """DNS over HTTPS Resolver""" - def __init__( - self, - *, - endpoints: List[str], - json_loads=json.loads, - resolver_class=None, - ) -> None: + def __init__(self, *, endpoints: List[str], json_loads=json.loads, resolver_class=None) -> None: self.endpoints = endpoints self.json_loads = json_loads if resolver_class is None: @@ -42,15 +36,14 @@ async def _resolve(self, endpoint: str, host, port, family): record_type = RecordType.A params = { - 'ct': 'application/dns-json', 'name': host, 'type': record_type.name, } resolver = self.resolveer_class() connector = TCPConnector(resolver=resolver) - - async with CS(connector=connector) as session: + + async with CS(connector=connector, headers={"accept": "application/dns-json"}) as session: async with session.get(endpoint, params=params) as resp: data = self.json_loads(await resp.text()) @@ -79,19 +72,15 @@ async def resolve(self, host, port=0, family=socket.AF_INET): self._resolve(endpoint, host, port, family) for endpoint in self.endpoints ] - done, pending = await asyncio.wait( - tasks, - return_when=asyncio.FIRST_COMPLETED, - ) - for p in pending: - p.cancel() - return list(done)[0].result() + done = await asyncio.gather(*tasks) + return done[0] async def close(self): pass -def ClientSession(*args, **kwargs) -> CS: # noqa +# noinspection PyTypeChecker +def DNSOverHTTPSClientSession(*args, **kwargs) -> CS: # noqa """Shortcut of aiohttp.ClientSession and DNSOverHTTPSResolver""" endpoints = kwargs.pop( diff --git a/tests/test_aiohttp_doh.py b/tests/test_aiohttp_doh.py new file mode 100644 index 0000000..a0d1fec --- /dev/null +++ b/tests/test_aiohttp_doh.py @@ -0,0 +1,27 @@ +import pytest +from aiohttp import ClientSession, TCPConnector + +import aiohttp_doh +from aiohttp_doh import DNSOverHTTPSClientSession + +pytest_plugins = ('pytest_asyncio',) + + +@pytest.mark.asyncio +async def test_DNSOverHTTPSResolver(): + endpoints = [ + 'https://dns.google.com/resolve', + 'https://cloudflare-dns.com/dns-query', + ] + resolver = aiohttp_doh.DNSOverHTTPSResolver(endpoints=endpoints) + connector = TCPConnector(resolver=resolver) + async with ClientSession(connector=connector) as session: + async with session.head("https://example.com") as response: + assert response.status == 200 + + +@pytest.mark.asyncio +async def test_DNSOverHTTPSClientSession(): + async with DNSOverHTTPSClientSession() as session: + async with session.head("https://example.com") as response: + assert response.status == 200