Repository: samuelcolvin/dnserver

Issue:
Let upstream be `None`
Hi. Thank you again for your amazing work!

I can now use dnserver in unit tests :rocket:

```python
@pytest.fixture(scope="session")
def dns_server():
    from dnserver import DNSServer

    server = DNSServer(
        zones_file="tests/fixtures/zones.toml", port=DNS_PORT, upstream=None
    )
    server.start()
    yield server
    server.stop()
```
However I would like to be able to disable any upstream DNS server, so my unit tests won't make any network I/O. At the moment this is prevented by this line:

https://github.com/samuelcolvin/dnserver/blob/ef85e97e5b80b0f031fa4f1bedf90440339eb800/dnserver/main.py#L125

What do you think?

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/samuelcolvin/dnserver/pull/9?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Samuel+Colvin) Report
> Merging [#9](https://codecov.io/gh/samuelcolvin/dnserver/pull/9?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Samuel+Colvin) (3026ee8) into [main](https://codecov.io/gh/samuelcolvin/dnserver/commit/4b90d04e53060c5f08dbdb4981156e46452633c1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Samuel+Colvin) (4b90d04) will **decrease** coverage by `1.06%`.
> The diff coverage is `100.00%`.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main       #9      +/-   ##
==========================================
- Coverage   98.30%   97.23%   -1.07%     
==========================================
  Files           5        5              
  Lines         177      181       +4     
  Branches       31       33       +2     
==========================================
+ Hits          174      176       +2     
- Misses          2        4       +2     
  Partials        1        1              
```


| [Impacted Files](https://codecov.io/gh/samuelcolvin/dnserver/pull/9?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Samuel+Colvin) | Coverage Δ | |
|---|---|---|
| [dnserver/main.py](https://codecov.io/gh/samuelcolvin/dnserver/pull/9/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Samuel+Colvin#diff-ZG5zZXJ2ZXIvbWFpbi5weQ==) | `98.90% <100.00%> (+0.05%)` | :arrow_up: |
| [dnserver/version.py](https://codecov.io/gh/samuelcolvin/dnserver/pull/9/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Samuel+Colvin#diff-ZG5zZXJ2ZXIvdmVyc2lvbi5weQ==) | `100.00% <100.00%> (ø)` | |
| [dnserver/load\_records.py](https://codecov.io/gh/samuelcolvin/dnserver/pull/9/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Samuel+Colvin#diff-ZG5zZXJ2ZXIvbG9hZF9yZWNvcmRzLnB5) | `92.30% <0.00%> (-3.85%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/samuelcolvin/dnserver/pull/9?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Samuel+Colvin).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Samuel+Colvin)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/samuelcolvin/dnserver/pull/9?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Samuel+Colvin). Last update [ef85e97...3026ee8](https://codecov.io/gh/samuelcolvin/dnserver/pull/9?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Samuel+Colvin). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Samuel+Colvin).

</details>
samuelcolvin: let me know when this is ready to review again.
azmeuk: It is ready!
azmeuk: This should be better now.
azmeuk: We are close :)

Patch:
diff --git a/dnserver/cli.py b/dnserver/cli.py
index c74de55..15d4803 100644
--- a/dnserver/cli.py
+++ b/dnserver/cli.py
@@ -6,7 +6,7 @@ import signal
 import sys
 from time import sleep
 
-from .main import DNSServer, logger
+from .main import DEFAULT_UPSTREAM, DNSServer, logger
 from .version import VERSION
 
 __all__ = ('cli',)
@@ -43,11 +43,25 @@ def cli_logic(args: list[str]) -> int:
             'if omitted will use DNSERVER_UPSTREAM env var, or 1.1.1.1'
         ),
     )
+    parser.add_argument(
+        '--no-upstream',
+        action='store_true',
+        default=False,
+        help=(
+            'Disable upstream DNS server. '
+            'If set and a domain is not in the zone TOML file, the DNS server will '
+            'consider the domain is not valid. If omitted will use DNSERVER_NO_UPSTREAM env var, or False'
+        ),
+    )
     parser.add_argument('--version', action='version', version=f'%(prog)s v{VERSION}')
     parsed_args = parser.parse_args(args)
 
     port = parsed_args.port or os.getenv('DNSERVER_PORT', None)
-    upstream = parsed_args.upstream or os.getenv('DNSERVER_UPSTREAM', None)
+    no_upstream = parsed_args.no_upstream or os.getenv('DNSERVER_NO_UPSTREAM', False)
+    if no_upstream:
+        upstream = None
+    else:
+        upstream = parsed_args.upstream or os.getenv('DNSERVER_UPSTREAM', DEFAULT_UPSTREAM)
     zones_file = parsed_args.zones_file or os.getenv('DNSERVER_ZONE_FILE', None)
     if zones_file is None:
         print('no zones file specified, use --help for more information', file=sys.stderr)
diff --git a/dnserver/main.py b/dnserver/main.py
index b8d4115..8bae461 100755
--- a/dnserver/main.py
+++ b/dnserver/main.py
@@ -7,8 +7,8 @@ from textwrap import wrap
 from typing import Any
 
 from dnslib import QTYPE, RR, DNSLabel, dns
-from dnslib.proxy import ProxyResolver
-from dnslib.server import DNSServer as LibDNSServer
+from dnslib.proxy import ProxyResolver as LibProxyResolver
+from dnslib.server import BaseResolver as LibBaseResolver, DNSServer as LibDNSServer
 
 from .load_records import Zone, load_records
 
@@ -85,33 +85,59 @@ class Record:
         return str(self.rr)
 
 
-class Resolver(ProxyResolver):
-    def __init__(self, zones_file: str | Path, upstream: str):
+def resolve(request, handler, records):
+    type_name = QTYPE[request.q.qtype]
+    reply = request.reply()
+    for record in records:
+        if record.match(request.q):
+            reply.add_answer(record.rr)
+
+    if reply.rr:
+        logger.info('found zone for %s[%s], %d replies', request.q.qname, type_name, len(reply.rr))
+        return reply
+
+    # no direct zone so look for an SOA record for a higher level zone
+    for record in records:
+        if record.sub_match(request.q):
+            reply.add_answer(record.rr)
+
+    if reply.rr:
+        logger.info('found higher level SOA resource for %s[%s]', request.q.qname, type_name)
+        return reply
+
+
+class BaseResolver(LibBaseResolver):
+    def __init__(self, zones_file: str | Path):
         records = load_records(zones_file)
         self.records = [Record(zone) for zone in records.zones]
-        logger.info('loaded %d zone record from %s', len(self.records), zones_file)
-        super().__init__(address=upstream, port=53, timeout=5)
+        logger.info('loaded %d zone record from %s, without proxy DNS server', len(self.records), zones_file)
+        super().__init__()
 
     def resolve(self, request, handler):
+        answer = resolve(request, handler, self.records)
+        if answer:
+            return answer
+
         type_name = QTYPE[request.q.qtype]
-        reply = request.reply()
-        for record in self.records:
-            if record.match(request.q):
-                reply.add_answer(record.rr)
+        logger.info('no local zone found, not proxying %s[%s]', request.q.qname, type_name)
+        return request.reply()
 
-        if reply.rr:
-            logger.info('found zone for %s[%s], %d replies', request.q.qname, type_name, len(reply.rr))
-            return reply
 
-        # no direct zone so look for an SOA record for a higher level zone
-        for record in self.records:
-            if record.sub_match(request.q):
-                reply.add_answer(record.rr)
+class ProxyResolver(LibProxyResolver):
+    def __init__(self, zones_file: str | Path, upstream: str):
+        records = load_records(zones_file)
+        self.records = [Record(zone) for zone in records.zones]
+        logger.info(
+            'loaded %d zone record from %s, with %s as a proxy DNS server', len(self.records), zones_file, upstream
+        )
+        super().__init__(address=upstream, port=53, timeout=5)
 
-        if reply.rr:
-            logger.info('found higher level SOA resource for %s[%s]', request.q.qname, type_name)
-            return reply
+    def resolve(self, request, handler):
+        answer = resolve(request, handler, self.records)
+        if answer:
+            return answer
 
+        type_name = QTYPE[request.q.qtype]
         logger.info('no local zone found, proxying %s[%s]', request.q.qname, type_name)
         return super().resolve(request, handler)
 
@@ -122,13 +148,16 @@ class DNSServer:
     ):
         self.zones_file = zones_file
         self.port: int = DEFAULT_PORT if port is None else int(port)
-        self.upstream: str = DEFAULT_UPSTREAM if upstream is None else upstream
+        self.upstream: str | None = upstream
         self.udp_server: LibDNSServer | None = None
         self.tcp_server: LibDNSServer | None = None
 
     def start(self):
         logger.info('starting DNS server on port %d, upstream DNS server "%s"', self.port, self.upstream)
-        resolver = Resolver(self.zones_file, self.upstream)
+        if self.upstream:
+            resolver = ProxyResolver(self.zones_file, self.upstream)
+        else:
+            resolver = BaseResolver(self.zones_file)
 
         self.udp_server = LibDNSServer(resolver, port=self.port)
         self.tcp_server = LibDNSServer(resolver, port=self.port, tcp=True)


Test patch:
diff --git a/tests/test_cli.py b/tests/test_cli.py
index a775113..1cdc24f 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -27,7 +27,7 @@ def test_cli(mocker):
     mock_signal = mocker.patch('dnserver.cli.signal.signal')
     assert cli_logic(['--port', '1234', 'zones.txt']) == 0
     assert calls == [
-        "init ('zones.txt',) {'port': '1234', 'upstream': None}",
+        "init ('zones.txt',) {'port': '1234', 'upstream': '1.1.1.1'}",
         'start',
         'is_running',
         'is_running',
diff --git a/tests/test_dnsserver.py b/tests/test_dnsserver.py
index 6788b0f..caf1774 100644
--- a/tests/test_dnsserver.py
+++ b/tests/test_dnsserver.py
@@ -128,3 +128,30 @@ def test_soa_higher(dns_resolver: Resolver):
     assert str(exc_info.value) == (
         'The DNS response does not contain an answer to the question: subdomain.example.com. IN SOA'
     )
+
+
+def test_dns_server_without_upstream():
+    port = 5054
+
+    server = DNSServer('example_zones.toml', port=port, upstream=None)
+    server.start()
+
+    resolver = RawResolver()
+    resolver.nameservers = ['127.0.0.1']
+    resolver.port = port
+
+    def resolve(name: str, type_: str) -> List[Dict[str, Any]]:
+        answers = resolver.resolve(name, type_)
+        return [convert_answer(answer) for answer in answers]
+
+    try:
+        assert resolve('example.com', 'A') == [
+            {
+                'type': 'A',
+                'value': '1.2.3.4',
+            },
+        ]
+        with pytest.raises(NoAnswer):
+            resolve('python.org', 'A')
+    finally:
+        server.stop()