Code review with various improvements

.pre-commit-config.yaml

@@ -1,19 +1,17 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+    rev: v6.0.0
     hooks:
       - id: check-merge-conflict
       - id: debug-statements
       - id: trailing-whitespace
       - id: end-of-file-fixer
-      - id: check-json
-      - id: check-toml
       - id: check-yaml
+      - id: check-toml
+      - id: check-json
+      - id: check-ast
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.9.10
+    rev: v0.16.9
     hooks:
-      # Run the linter.
-      - id: ruff
-        args: [ --fix ]
-      # Run the formatter.
+      - id: ruff-check
       - id: ruff-format

dnstapir/dns/mozpsl.py

@@ -1,6 +1,6 @@
 import io
 
-import httpx
+import httpx2
 
 
 class TrieNode:
@@ -74,7 +74,7 @@ def __init__(self) -> None:
 
     def load_psl_url(self, url: str) -> None:
         """Load PSL from URL"""
-        response = httpx.get(
+        response = httpx2.get(
             url,
             headers={
                 "Accept-Encoding": "gzip",
@@ -126,7 +126,7 @@ def coredomain(self, domain: str) -> tuple[str, str]:
         if not domain:
             raise ValueError
         try:
-            domain = domain.rstrip(".")
+            domain = domain.lower().rstrip(".")
         except AttributeError as exc:
             raise ValueError from exc
         lbls = domain.split(".")
@@ -151,6 +151,6 @@ def coredomain(self, domain: str) -> tuple[str, str]:
 
     def rdomain(self, rdomain: str) -> tuple[str, str]:
         """Find ICANN and private name cut-off for domain, reverse order process"""
-        lbls = rdomain.split(".")
+        lbls = rdomain.lower().split(".")
         c, p = self.trie.search(lbls)
         return (".".join(lbls[0:c]), ".".join(lbls[0:p]))

dnstapir/jws.py

@@ -16,33 +16,44 @@ class ResolverJWKSet(JWKSet):
     def __init__(self, key_resolver: KeyResolver):
         super().__init__()
         self.key_resolver = key_resolver
-        self._cache: dict[str, JWK] = {}
 
     def get_key(self, kid: str) -> JWK:
-        if kid in self._cache:
-            return self._cache[kid]
-        key = JWK.from_pyca(self.key_resolver.resolve_public_key(kid))  # type: ignore
-        self._cache[kid] = key
-        return key  # type: ignore
+        return JWK.from_pyca(self.key_resolver.resolve_public_key(kid))
 
     def get_keys(self, kid: str) -> list[JWK]:
         return [self.get_key(kid)]
 
     def verify_jws(self, jws: JWS) -> JWK:
         """Verify JWS and return verified key (or raise JWKeyNotFound)"""
-        protected_header: dict[str, str] = json.loads(jws.objects["protected"])
-        if kid := protected_header.get("kid"):
-            logger.debug("Signature by kid=%s", kid)
-            for key in self.get_keys(kid):
-                try:
-                    jws.verify(key=key)
-                    if not hasattr(key, "kid"):
-                        key.kid = kid
-                    return key
-                except InvalidJWSSignature:
-                    pass
-        else:
-            logger.debug("Signature without kid")
+
+        for signature in jws.objects.get("signatures", [jws.objects]):
+            unprotected_header = signature.get("header", {})
+            if protected_header_str := signature.get("protected"):
+                protected_header: dict[str, str] = json.loads(protected_header_str)
+                if (kid := protected_header.get("kid")) and (
+                    alg := protected_header.get("alg", unprotected_header.get("alg"))
+                ):
+                    logger.debug("Signature by kid=%s alg=%s", kid, alg)
+                    try:
+                        keys = self.get_keys(kid)
+                    except KeyError:
+                        logger.debug("Key with kid=%s not found", kid)
+                        continue
+                    except Exception as exc:
+                        logger.debug("Error looking up key with kid=%s: %s", kid, exc, exc_info=exc)
+                        continue
+                    for key in keys:
+                        try:
+                            jws.verify(key=key, alg=alg)
+                            if not hasattr(key, "kid"):
+                                key.kid = kid
+                            return key
+                        except InvalidJWSSignature:
+                            pass
+                else:
+                    logger.debug("Skipping signature without kid")
+            else:
+                logger.debug("No protected header found in signature")
         raise JWKeyNotFound
 
 
@@ -51,7 +62,7 @@ def main() -> None:
 
     parser = argparse.ArgumentParser(description="JWS Verifier")
 
-    parser.add_argument("--nodeman", help="Nodeman API")
+    parser.add_argument("--nodeman", help="Nodeman API", required=True)
     parser.add_argument("--debug", action="store_true", help="Enable debugging")
     parser.add_argument("message", help="JWS message")
 

dnstapir/key_cache.py

@@ -1,6 +1,5 @@
 import logging
-import time
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from datetime import timedelta
 
 import redis
@@ -23,18 +22,17 @@ class KeyCacheSettings(BaseModel):
 
 
 def key_cache_from_settings(settings: KeyCacheSettings):
-    memory_key_cache = MemoryKeyCache(size=settings.size, ttl=settings.ttl)
     if settings.redis:
+        memory_key_cache = MemoryKeyCache(size=settings.size, ttl=settings.ttl) if settings.size else None
         redis_client = redis.StrictRedis(host=settings.redis.host, port=settings.redis.port)
-        redis_key_cache = RedisKeyCache(redis_client=redis_client, ttl=settings.ttl)
-        return CombinedKeyCache([memory_key_cache, redis_key_cache]) if settings.size else redis_key_cache
+        return RedisKeyCache(redis_client=redis_client, ttl=settings.ttl, memory_cache=memory_key_cache)
     elif settings.size:
-        return memory_key_cache
+        return MemoryKeyCache(size=settings.size, ttl=settings.ttl)
     else:
         return DummyKeyCache()
 
 
-class KeyCache:
+class KeyCache(ABC):
     def __init__(self):
         self.logger = logging.getLogger(__name__).getChild(self.__class__.__name__)
 
@@ -74,35 +72,24 @@ def set(self, key: str, value: bytes) -> None:
 
 
 class RedisKeyCache(KeyCache):
-    def __init__(self, redis_client: redis.Redis, ttl: int):
+    def __init__(self, redis_client: redis.Redis, ttl: int, memory_cache: MemoryKeyCache | None = None):
         super().__init__()
         self.redis_client = redis_client
         self.ttl = ttl
+        self.memory_cache = memory_cache
         self.logger.info("Configured Redis key cache ttl=%d", ttl)
 
     def get(self, key: str) -> bytes | None:
+        if self.memory_cache and (res := self.memory_cache.get(key)):
+            return res
         with tracer.start_as_current_span("redis_key_cache_get"):
             res = self.redis_client.get(name=key)
         self.logger.debug("Cache GET %s (%s)", key, "hit" if res else "miss")
         return res
 
     def set(self, key: str, value: bytes) -> None:
         self.logger.debug("Cache SET %s", key)
-        expires_at = int(time.time()) + self.ttl
         with tracer.start_as_current_span("redis_key_cache_set"):
-            self.redis_client.set(name=key, value=value, exat=expires_at)
-
-
-class CombinedKeyCache(KeyCache):
-    def __init__(self, caches: list[KeyCache]):
-        self.caches = caches
-
-    def get(self, key: str) -> bytes | None:
-        for cache in self.caches:
-            if res := cache.get(key):
-                return res
-        return None
-
-    def set(self, key: str, value: bytes) -> None:
-        for cache in self.caches:
-            cache.set(key, value)
+            self.redis_client.set(name=key, value=value, ex=self.ttl)
+        if self.memory_cache:
+            self.memory_cache.set(key, value)

dnstapir/key_resolver.py

@@ -1,30 +1,32 @@
 import logging
 import re
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from pathlib import Path
 from urllib.parse import urljoin, urlparse
 
-import httpx
-from cryptography.hazmat.primitives.asymmetric.ec import EllipticCurvePublicKey
-from cryptography.hazmat.primitives.asymmetric.ed448 import Ed448PublicKey
-from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
-from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
+import httpx2
+from cryptography.hazmat.primitives.asymmetric.types import PublicKeyTypes
 from cryptography.hazmat.primitives.serialization import load_pem_public_key
 from opentelemetry import metrics, trace
 
 from .key_cache import KeyCache
 
-PublicKey = Ed25519PublicKey | Ed448PublicKey | EllipticCurvePublicKey | RSAPublicKey
+type PublicKey = PublicKeyTypes
 
 tracer = trace.get_tracer("dnstapir.tracer")
 meter = metrics.get_meter("dnstapir.meter")
 
 public_key_get_counter = meter.create_counter(
-    "aggregates.public_key_get_counter",
+    "dnstapir.public_key_get_counter",
     description="The number of public key lookups",
 )
 
-KEY_ID_VALIDATOR = re.compile(r"^[a-zA-Z0-9_\-.]+$")
+public_key_resolve_counter = meter.create_counter(
+    "dnstapir.public_key_resolve_counter",
+    description="The number of public key resolutions",
+)
+
+KEY_ID_VALIDATOR = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9_\-.]*$")
 
 
 def key_resolver_from_client_database(client_database: str, key_cache: KeyCache | None = None):
@@ -34,17 +36,17 @@ def key_resolver_from_client_database(client_database: str, key_cache: KeyCache
         return FileKeyResolver(client_database_directory=client_database, key_cache=key_cache)
 
 
-class KeyResolver:
+class KeyResolver(ABC):
     def __init__(self):
         self.logger = logging.getLogger(__name__).getChild(self.__class__.__name__)
         self.key_id_validator = KEY_ID_VALIDATOR
 
     @abstractmethod
-    def resolve_public_key(self, key_id: str) -> PublicKey:
+    def resolve_public_key(self, key_id: str) -> PublicKeyTypes:
         pass
 
     def validate_key_id(self, key_id: str) -> None:
-        if not self.key_id_validator.match(key_id):
+        if not self.key_id_validator.fullmatch(key_id):
             raise ValueError(f"Invalid key_id format: {key_id}")
 
 
@@ -54,19 +56,24 @@ def __init__(self, key_cache: KeyCache | None):
         self.key_cache = key_cache
 
     @abstractmethod
-    def get_public_key_pem(self, key_id: str) -> bytes:
+    def _get_public_key_pem(self, key_id: str) -> bytes:
         pass
 
-    def resolve_public_key(self, key_id: str):
+    def resolve_public_key(self, key_id: str) -> PublicKeyTypes:
+        self.validate_key_id(key_id)
+        public_key_resolve_counter.add(1)
         with tracer.start_as_current_span("resolve_public_key"):
             if self.key_cache:
                 public_key_pem = self.key_cache.get(key_id)
                 if not public_key_pem:
-                    public_key_pem = self.get_public_key_pem(key_id)
+                    public_key_pem = self._get_public_key_pem(key_id)
+                    # Load the public key from PEM format before caching and returning it
+                    res = load_pem_public_key(public_key_pem)
                     self.key_cache.set(key_id, public_key_pem)
                     public_key_get_counter.add(1)
+                    return res
             else:
-                public_key_pem = self.get_public_key_pem(key_id)
+                public_key_pem = self._get_public_key_pem(key_id)
         return load_pem_public_key(public_key_pem)
 
 
@@ -75,9 +82,8 @@ def __init__(self, client_database_directory: str, key_cache: KeyCache | None =
         super().__init__(key_cache=key_cache)
         self.client_database_directory = client_database_directory
 
-    def get_public_key_pem(self, key_id: str) -> bytes:
+    def _get_public_key_pem(self, key_id: str) -> bytes:
         with tracer.start_as_current_span("get_public_key_pem_from_file"):
-            self.validate_key_id(key_id)
             filename = Path(self.client_database_directory) / f"{key_id}.pem"
             self.logger.debug("Fetching public key for %s from %s", key_id, filename)
             try:
@@ -92,7 +98,7 @@ def __init__(self, client_database_base_url: str, key_cache: KeyCache | None = N
         super().__init__(key_cache=key_cache)
 
         self.client_database_base_url = client_database_base_url
-        self._httpx_client: httpx.Client | None = None
+        self._httpx_client: httpx2.Client | None = None
         self.key_id_pattern = "{key_id}"
 
         if urlparse(self.client_database_base_url).scheme not in ("http", "https"):
@@ -103,10 +109,8 @@ def __init__(self, client_database_base_url: str, key_cache: KeyCache | None = N
             if urlparse(test_url).scheme not in ("http", "https"):
                 raise ValueError(f"Invalid URL pattern: {self.client_database_base_url}")
 
-    def get_public_key_pem(self, key_id: str) -> bytes:
+    def _get_public_key_pem(self, key_id: str) -> bytes:
         with tracer.start_as_current_span("get_public_key_pem_from_url"):
-            self.validate_key_id(key_id)
-
             if self.key_id_pattern in self.client_database_base_url:
                 public_key_url = self.client_database_base_url.replace(self.key_id_pattern, key_id)
             else:
@@ -120,13 +124,13 @@ def get_public_key_pem(self, key_id: str) -> bytes:
                 response = self.httpx_client.get(public_key_url)
                 response.raise_for_status()
                 return response.content
-            except httpx.HTTPError as exc:
+            except httpx2.HTTPError as exc:
                 raise KeyError(key_id) from exc
 
     @property
-    def httpx_client(self) -> httpx.Client:
+    def httpx_client(self) -> httpx2.Client:
         if self._httpx_client is None:
-            self._httpx_client = httpx.Client(headers={"Accept": "application/x-pem-file"})
+            self._httpx_client = httpx2.Client(http2=True, headers={"Accept": "application/x-pem-file"})
         return self._httpx_client
 
     def __enter__(self):

dnstapir/logging.py

@@ -9,6 +9,8 @@
 
 VALID_LOG_LEVELS = set(["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG"])
 
+HANDLER_NAME = "dnstapir"
+
 
 def drop_color_message_key(_, __, event_dict: EventDict) -> EventDict:
     """
@@ -67,9 +69,13 @@ def setup_logging(json_logs: bool = False, log_level: str = "INFO") -> None:
     )
 
     # Reconfigure the root logger to use our structlog formatter, effectively emitting the logs via structlog
+    # Replace any handler from a previous call to avoid duplicate log lines
     handler = logging.StreamHandler()
+    handler.set_name(HANDLER_NAME)
     handler.setFormatter(formatter)
     root_logger = logging.getLogger()
+    for existing in [h for h in root_logger.handlers if h.get_name() == HANDLER_NAME]:
+        root_logger.removeHandler(existing)
     root_logger.addHandler(handler)
     root_logger.setLevel(log_level.upper())