Side-Channel Risks for Geospatial Enclave Queries

from __future__ import annotations

import math
import random
from dataclasses import dataclass
from typing import Iterable, Mapping, Sequence


@dataclass(frozen=True)
class AccessTrace:
    """What a malicious host can observe: pages touched, in order, with timings."""

    pages: tuple[int, ...]
    total_units: int

    def distinct_pages(self) -> int:
        return len(set(self.pages))


def naive_lookup(
    query_cells: Sequence[str], occupancy: Mapping[str, int], *, cells_per_page: int
) -> AccessTrace:
    """The default implementation: touch only what is needed, cost what it costs.

    Both observables are data-dependent — the page sequence names the region, and
    the unit count reveals its density.
    """
    pages: list[int] = []
    units = 0
    for cell in query_cells:
        page = abs(hash(cell)) % 4096 // max(1, cells_per_page)
        pages.append(page)
        units += max(1, occupancy.get(cell, 0))
    return AccessTrace(tuple(pages), units)


def padded_lookup(
    query_cells: Sequence[str],
    occupancy: Mapping[str, int],
    *,
    cells_per_page: int,
    padding_target: int,
) -> AccessTrace:
    """Oblivious in cost: every cell is processed for exactly `padding_target` units.

    Cells above the target are truncated, which is a correctness trade — set the
    target from a high occupancy percentile so truncation is rare, and account for
    the truncation as a bias in the released aggregate.
    """
    pages: list[int] = []
    for cell in query_cells:
        pages.append(abs(hash(cell)) % 4096 // max(1, cells_per_page))
    return AccessTrace(tuple(pages), padding_target * len(query_cells))


def dummy_padded_lookup(
    query_cells: Sequence[str],
    occupancy: Mapping[str, int],
    all_cells: Sequence[str],
    *,
    cells_per_page: int,
    dummy_ratio: int,
    seed: int = 0,
) -> AccessTrace:
    """Interleave `dummy_ratio` decoy cells per real cell, in shuffled order.

    Weaker than padding: the observer still sees a cost that correlates with real
    occupancy, just diluted. Shuffling matters — appending dummies after the real
    accesses makes them trivially separable by position.
    """
    rng = random.Random(seed)
    plan = list(query_cells)
    for _ in range(dummy_ratio * len(query_cells)):
        plan.append(rng.choice(all_cells))
    rng.shuffle(plan)
    pages = [abs(hash(c)) % 4096 // max(1, cells_per_page) for c in plan]
    units = sum(max(1, occupancy.get(c, 0)) for c in plan)
    return AccessTrace(tuple(pages), units)


def leakage_score(traces: Sequence[AccessTrace]) -> float:
    """How distinguishable a set of traces is: 0 means identical, 1 means unique.

    A crude but honest proxy — if two queries over different regions produce traces
    an observer can tell apart, the access pattern is leaking the region regardless
    of what the formal enclave guarantee says.
    """
    if len(traces) < 2:
        return 0.0
    signatures = {(t.distinct_pages(), t.total_units) for t in traces}
    return (len(signatures) - 1) / (len(traces) - 1)

def _validate() -> None:
    rng = random.Random(11)
    all_cells = [f"c{i}" for i in range(4_000)]
    occupancy = {c: max(1, int(rng.lognormvariate(3.0, 1.4))) for c in all_cells}

    dense = [c for c in all_cells if occupancy[c] > 120][:24]
    sparse = [c for c in all_cells if occupancy[c] < 8][:24]

    # 1. The naive implementation leaks: dense and sparse queries look different.
    naive_traces = [naive_lookup(dense, occupancy, cells_per_page=32),
                    naive_lookup(sparse, occupancy, cells_per_page=32)]
    assert naive_traces[0].total_units > 4 * naive_traces[1].total_units
    assert leakage_score(naive_traces) == 1.0

    # 2. Padding removes the cost signal entirely.
    padded = [padded_lookup(dense, occupancy, cells_per_page=32, padding_target=200),
              padded_lookup(sparse, occupancy, cells_per_page=32, padding_target=200)]
    assert padded[0].total_units == padded[1].total_units

    # 3. Padding does NOT hide which cells were queried — be honest about that.
    assert padded[0].pages != padded[1].pages

    # 4. Padding's cost is the target times the cell count, which is the price.
    assert padded[0].total_units == 200 * len(dense)
    assert padded[0].total_units > naive_traces[1].total_units * 10

    # 5. Dummy accesses reduce, but do not remove, the correlation.
    d_dense = dummy_padded_lookup(dense, occupancy, all_cells,
                                  cells_per_page=32, dummy_ratio=3, seed=1)
    d_sparse = dummy_padded_lookup(sparse, occupancy, all_cells,
                                   cells_per_page=32, dummy_ratio=3, seed=1)
    ratio_naive = naive_traces[0].total_units / naive_traces[1].total_units
    ratio_dummy = d_dense.total_units / d_sparse.total_units
    assert 1.0 < ratio_dummy < ratio_naive, (ratio_dummy, ratio_naive)

    # 6. Coarser page packing raises the observer's uncertainty at no runtime cost.
    fine = naive_lookup(dense, occupancy, cells_per_page=8)
    coarse = naive_lookup(dense, occupancy, cells_per_page=128)
    assert coarse.distinct_pages() <= fine.distinct_pages()

    print("side-channel analysis: all assertions passed")


_validate()