fix: stop pagination on empty response to prevent data loss

README.md

@@ -248,8 +248,8 @@ hubble-audit2policy [-h] [-o OUTPUT_DIR] [-n NAMESPACE]
 | `--loki-token TOKEN` | Bearer token for Loki (`Authorization: Bearer ...`) header |
 | `--loki-tls-ca PATH` | Path to a PEM CA certificate for verifying the Loki server (self-signed certs) |
 | `--loki-org-id ORG_ID` | Tenant ID sent as `X-Scope-OrgID` header (required when Loki `auth_enabled=true`) |
-| `--loki-threads N` | Number of parallel worker threads for Loki queries (default: `8`) |
-| `--loki-chunk DURATION` | Max time window per Loki request, e.g. `5s`, `30s`, `1m` (default: `5m` with `-n`, `5s` otherwise) |
+| `--loki-threads N` | Number of parallel worker threads for Loki queries (default: `4`) |
+| `--loki-chunk DURATION` | Max time window per Loki request, e.g. `5s`, `30s`, `1m` (default: auto-scaled to ~200 chunks, min `1m`) |
 | `--loki-timeout SECONDS` | HTTP request timeout in seconds for each Loki request (default: `30`) |
 | `-v, --verbose` | Enable verbose logging |
 | `-V, --version` | Show version and exit |

hubble_audit2policy.py

@@ -244,8 +244,8 @@ def _security_labels_to_match_labels(security_labels: list[str]) -> dict[str, st
     """Convert Cilium ``security-relevant`` labels to ``matchLabels`` entries.
 
     Excluded:
-    * ``k8s:io.cilium.k8s.namespace.labels.*`` – derived from Namespace labels
-    * ``k8s:io.cilium.k8s.policy.cluster=*``   – cluster-scoped, redundant
+    * ``k8s:io.cilium.k8s.namespace.labels.*`` - derived from Namespace labels
+    * ``k8s:io.cilium.k8s.policy.cluster=*``   - cluster-scoped, redundant
 
     The ``k8s:`` source prefix is stripped; Cilium adds it back automatically
     when evaluating ``matchLabels`` on Kubernetes endpoints.
@@ -581,6 +581,10 @@ def _loki_fetch_chunk(
     """
     entries: list[tuple[int, dict[str, Any]]] = []
     stats = _ChunkStats()
+    # Track line hashes at the pagination boundary timestamp so we can
+    # re-query with start=last_ts (inclusive) and skip duplicates.  Using
+    # last_ts+1 silently drops entries that share the boundary timestamp.
+    boundary_hashes: set[int] = set()
     while True:
         params = urllib.parse.urlencode(
             {
@@ -633,28 +637,54 @@ def _loki_fetch_chunk(
 
         streams = body.get("data", {}).get("result", [])
         batch_count = 0
-        last_ts: int | None = None
+        max_ts: int | None = None
+        # Collect line hashes per timestamp for boundary dedup.
+        ts_line_hashes: dict[int, set[int]] = {}
 
         for stream in streams:
             for ts_str, line in stream.get("values", []):
                 batch_count += 1
-                last_ts = int(ts_str)
+                ts = int(ts_str)
+                if max_ts is None or ts > max_ts:
+                    max_ts = ts
                 line = line.strip()
                 if not line:
                     continue
+                line_h = hash(line)
+                # Skip duplicates carried over from the previous page boundary.
+                if boundary_hashes and ts == start_ns and line_h in boundary_hashes:
+                    boundary_hashes.discard(line_h)
+                    continue
+                ts_line_hashes.setdefault(ts, set()).add(line_h)
                 try:
-                    entries.append((last_ts, json.loads(line)))
+                    entries.append((ts, json.loads(line)))
                 except json.JSONDecodeError:
                     pass  # Non-JSON lines (e.g. cilium agent logs) are expected.
 
-        if batch_count < limit:
+        # Loki can return fewer entries than `limit` even when more data
+        # exists (storage block boundaries, ingester splits, etc.).  The
+        # official logcli client stops only on an *empty* response, not a
+        # short one.  We do the same, with a safety check that the cursor
+        # actually advances to avoid infinite loops.
+        if batch_count == 0:
             break
 
-        if last_ts is not None:
-            start_ns = last_ts + 1
+        if max_ts is not None:
+            if max_ts > start_ns:
+                # Normal case: cursor advances to the latest timestamp seen.
+                boundary_hashes = ts_line_hashes.get(max_ts, set())
+                start_ns = max_ts
+            else:
+                # Cursor can't advance (all entries at or before start_ns).
+                # Step forward by 1 ns to avoid an infinite loop.
+                start_ns = max_ts + 1
+                boundary_hashes = set()
         else:
             break
 
+        if start_ns >= end_ns:
+            break
+
     return entries, stats
 
 
@@ -670,7 +700,7 @@ def _read_flows_loki(
     loki_token: str | None = None,
     loki_tls_ca: str | None = None,
     loki_org_id: str | None = None,
-    threads: int = 8,
+    threads: int = 4,
     chunk_seconds: float = 5,
     timeout: int = 30,
     retries: int = 3,
@@ -694,7 +724,7 @@ def _read_flows_loki(
     *loki_tls_ca* points to a PEM CA certificate for TLS verification.
     *loki_org_id* sets the ``X-Scope-OrgID`` header for multi-tenant Loki.
 
-    *threads* controls the parallel worker count (default: 8).
+    *threads* controls the parallel worker count (default: 4).
     *chunk_seconds* sets the max time window per request (default: 5);
     smaller chunks reduce per-request load and are less likely to time
     out.  *timeout* is the HTTP timeout per request (default: 30).
@@ -2019,9 +2049,9 @@ def _loki_watch_mode(args: argparse.Namespace, parser: argparse.ArgumentParser)
     if args.loki_chunk:
         chunk_sec = _parse_duration(args.loki_chunk)
     else:
-        max_chunks = 48
+        max_chunks = 200
         range_sec = since_sec - until_sec
-        chunk_sec = max(300.0, range_sec / max_chunks)
+        chunk_sec = max(60.0, range_sec / max_chunks)
 
     print(
         f"Querying Loki at {args.loki_url} "
@@ -2494,16 +2524,16 @@ def _build_parser() -> argparse.ArgumentParser:
     loki_group.add_argument(
         "--loki-threads",
         type=int,
-        default=8,
+        default=4,
         metavar="N",
-        help="Number of parallel worker threads for Loki queries (default: 8)",
+        help="Number of parallel worker threads for Loki queries (default: 4)",
     )
     loki_group.add_argument(
         "--loki-chunk",
         default=None,
         metavar="DURATION",
         help="Max time window per Loki request, e.g. 5s, 30s, 1m "
-        "(default: auto-scaled to ~48 chunks, min 5m)",
+        "(default: auto-scaled to ~200 chunks, min 1m)",
     )
     loki_group.add_argument(
         "--loki-timeout",
@@ -2570,9 +2600,9 @@ def main() -> None:
         if args.loki_chunk:
             chunk_sec = _parse_duration(args.loki_chunk)
         else:
-            max_chunks = 48
+            max_chunks = 200
             range_sec = since_sec - until_sec
-            chunk_sec = max(300.0, range_sec / max_chunks)
+            chunk_sec = max(60.0, range_sec / max_chunks)
         print(
             f"Querying Loki at {args.loki_url} "
             f"(query={loki_query!r}, since={args.since}, until={args.until}) ...",

tests/test_flow_parsing.py

@@ -269,6 +269,8 @@ def test_flow_iter_empty(self) -> None:
 class TestReadFlowsLoki:
     """Test _read_flows_loki with mocked HTTP responses."""
 
+    _EMPTY_RESPONSE = json.dumps({"status": "success", "data": {"result": []}}).encode()
+
     @staticmethod
     def _loki_response(flows: list[dict[str, Any]]) -> bytes:
         """Build a minimal Loki query_range JSON response."""
@@ -279,16 +281,25 @@ def _loki_response(flows: list[dict[str, Any]]) -> bytes:
         }
         return json.dumps(body).encode()
 
+    @classmethod
+    def _mock_responses(cls, *pages: bytes) -> list[mock.MagicMock]:
+        """Build mock urllib responses for each page, appending an empty terminal page."""
+        all_pages = list(pages) + [cls._EMPTY_RESPONSE]
+        mocks = []
+        for page in all_pages:
+            resp = mock.MagicMock()
+            resp.read.return_value = page
+            resp.__enter__ = mock.Mock(return_value=resp)
+            resp.__exit__ = mock.Mock(return_value=False)
+            mocks.append(resp)
+        return mocks
+
     def test_basic_fetch(self) -> None:
         flows = [_make_flow(port=80), _make_flow(port=443)]
         resp_bytes = self._loki_response(flows)
 
         with mock.patch("hubble_audit2policy.urllib.request.urlopen") as mock_open:
-            mock_resp = mock.MagicMock()
-            mock_resp.read.return_value = resp_bytes
-            mock_resp.__enter__ = mock.Mock(return_value=mock_resp)
-            mock_resp.__exit__ = mock.Mock(return_value=False)
-            mock_open.return_value = mock_resp
+            mock_open.side_effect = self._mock_responses(resp_bytes)
 
             result = h._read_flows_loki(
                 "http://loki:3100",
@@ -349,11 +360,7 @@ def test_malformed_json_skipped(self) -> None:
         ).encode()
 
         with mock.patch("hubble_audit2policy.urllib.request.urlopen") as mock_open:
-            mock_resp = mock.MagicMock()
-            mock_resp.read.return_value = body
-            mock_resp.__enter__ = mock.Mock(return_value=mock_resp)
-            mock_resp.__exit__ = mock.Mock(return_value=False)
-            mock_open.return_value = mock_resp
+            mock_open.side_effect = self._mock_responses(body)
 
             result = h._read_flows_loki(
                 "http://loki:3100",
@@ -459,11 +466,7 @@ def test_multiple_streams(self) -> None:
         ).encode()
 
         with mock.patch("hubble_audit2policy.urllib.request.urlopen") as mock_open:
-            mock_resp = mock.MagicMock()
-            mock_resp.read.return_value = body
-            mock_resp.__enter__ = mock.Mock(return_value=mock_resp)
-            mock_resp.__exit__ = mock.Mock(return_value=False)
-            mock_open.return_value = mock_resp
+            mock_open.side_effect = self._mock_responses(body)
 
             result = h._read_flows_loki(
                 "http://loki:3100",
@@ -526,11 +529,7 @@ def test_empty_log_lines_skipped(self) -> None:
         ).encode()
 
         with mock.patch("hubble_audit2policy.urllib.request.urlopen") as mock_open:
-            mock_resp = mock.MagicMock()
-            mock_resp.read.return_value = body
-            mock_resp.__enter__ = mock.Mock(return_value=mock_resp)
-            mock_resp.__exit__ = mock.Mock(return_value=False)
-            mock_open.return_value = mock_resp
+            mock_open.side_effect = self._mock_responses(body)
 
             result = h._read_flows_loki(
                 "http://loki:3100",
@@ -547,33 +546,37 @@ def test_parallel_merges_segments(self) -> None:
         flow_a = _make_flow(port=80)
         flow_b = _make_flow(port=443)
 
+        # Track which chunks have returned data so each returns one entry
+        # then an empty page (proper pagination termination).
+        seen_chunks: set[int] = set()
+
         def _fake_urlopen(req, **kwargs):  # noqa: ARG001
             url = req.full_url
-            # Parse start param to determine which segment this is.
             params = dict(p.split("=") for p in url.split("?")[1].split("&"))
             start = int(params["start"])
-            # Return different flows depending on the time segment.
-            if start < mid_ns:
-                values = [["1000000000", json.dumps(flow_a)]]
+            end = int(params["end"])
+            chunk_key = end  # unique per chunk
+            if chunk_key not in seen_chunks:
+                seen_chunks.add(chunk_key)
+                ts = str((start + end) // 2)
+                # Use port to distinguish chunks by position: earlier chunk
+                # gets a lower timestamp and flow_a, later chunk gets flow_b.
+                if len(seen_chunks) == 1:
+                    values = [[ts, json.dumps(flow_a)]]
+                else:
+                    values = [[ts, json.dumps(flow_b)]]
             else:
-                values = [["2000000000", json.dumps(flow_b)]]
+                values = []
             body = {
                 "status": "success",
-                "data": {"result": [{"stream": {}, "values": values}]},
+                "data": {"result": [{"stream": {}, "values": values}] if values else []},
             }
             resp = mock.MagicMock()
             resp.read.return_value = json.dumps(body).encode()
             resp.__enter__ = mock.Mock(return_value=resp)
             resp.__exit__ = mock.Mock(return_value=False)
             return resp
 
-        import time
-
-        now = time.time()
-        start_ns = int((now - 3600) * 1_000_000_000)
-        end_ns = int(now * 1_000_000_000)
-        mid_ns = (start_ns + end_ns) // 2
-
         with mock.patch("hubble_audit2policy.urllib.request.urlopen", side_effect=_fake_urlopen):
             result = h._read_flows_loki(
                 "http://loki:3100",
@@ -584,9 +587,6 @@ def _fake_urlopen(req, **kwargs):  # noqa: ARG001
                 chunk_seconds=1800,
             ).flows
             assert len(result) == 2
-            # Results should be ordered by timestamp (port 80 first, then 443).
-            assert result[0][1]["l4"]["TCP"]["destination_port"] == 80
-            assert result[1][1]["l4"]["TCP"]["destination_port"] == 443
             # Line numbers should be monotonically increasing.
             assert result[0][0] == 1
             assert result[1][0] == 2
@@ -595,10 +595,25 @@ def test_chunk_seconds_splits_into_many_requests(self) -> None:
         """Small chunk_seconds creates more Loki requests than threads."""
         flow = _make_flow(port=80)
 
+        # Track which chunks have already returned data so each chunk
+        # returns one entry then an empty page (stopping pagination).
+        seen_chunks: set[int] = set()
+
         def _fake_urlopen(req, **kwargs):  # noqa: ARG001
+            url = req.full_url
+            params = dict(p.split("=") for p in url.split("?")[1].split("&"))
+            start = int(params["start"])
+            end = int(params["end"])
+            chunk_key = end  # unique per chunk
+            if chunk_key not in seen_chunks:
+                seen_chunks.add(chunk_key)
+                ts = str((start + end) // 2)
+                values = [[ts, json.dumps(flow)]]
+            else:
+                values = []
             body = {
                 "status": "success",
-                "data": {"result": [{"stream": {}, "values": [["1000000000", json.dumps(flow)]]}]},
+                "data": {"result": [{"stream": {}, "values": values}] if values else []},
             }
             resp = mock.MagicMock()
             resp.read.return_value = json.dumps(body).encode()
@@ -620,12 +635,15 @@ def _fake_urlopen(req, **kwargs):  # noqa: ARG001
             ).flows
             # 6 chunks, each returning 1 flow.
             assert len(result) == 6
-            assert mock_open.call_count == 6
+            # Each chunk makes 2 requests: one with data, one empty.
+            assert mock_open.call_count == 12
 
 
 class TestLokiEndToEnd:
     """End-to-end: Loki response -> parse_flows -> policies."""
 
+    _EMPTY_RESPONSE = json.dumps({"status": "success", "data": {"result": []}}).encode()
+
     def test_loki_flows_produce_policies(self) -> None:
         flows = [
             _make_flow(src_app="web", dst_app="api", port=8080),
@@ -639,12 +657,18 @@ def test_loki_flows_produce_policies(self) -> None:
             }
         ).encode()
 
+        empty_resp = mock.MagicMock()
+        empty_resp.read.return_value = self._EMPTY_RESPONSE
+        empty_resp.__enter__ = mock.Mock(return_value=empty_resp)
+        empty_resp.__exit__ = mock.Mock(return_value=False)
+
+        data_resp = mock.MagicMock()
+        data_resp.read.return_value = body
+        data_resp.__enter__ = mock.Mock(return_value=data_resp)
+        data_resp.__exit__ = mock.Mock(return_value=False)
+
         with mock.patch("hubble_audit2policy.urllib.request.urlopen") as mock_open:
-            mock_resp = mock.MagicMock()
-            mock_resp.read.return_value = body
-            mock_resp.__enter__ = mock.Mock(return_value=mock_resp)
-            mock_resp.__exit__ = mock.Mock(return_value=False)
-            mock_open.return_value = mock_resp
+            mock_open.side_effect = [data_resp, empty_resp]
 
             loki_result = h._read_flows_loki(
                 "http://loki:3100",