hyperswarm-peer / src / nat_sampler.h

#ifndef HYPERSWARM_NAT_SAMPLER_H
#define HYPERSWARM_NAT_SAMPLER_H

// NatSampler — Firewall detection via external address sampling.
// Wire-compatible algorithm with holepunchto/hyperdht/lib/nat.js.
// Collects (host, port) samples from DHT response `to` fields,
// then classifies our NAT type as OPEN, CONSISTENT, or RANDOM.

#include "compact_encoding.h"
#include "dht_messages.h"

#include <cstdint>
#include <cstring>

static constexpr int NAT_SAMPLER_MAX_SAMPLES = 32;
static constexpr int NAT_SAMPLER_MIN_SAMPLES = 3;
static constexpr int NAT_SAMPLER_MAX_ADDRESSES = 8;

class NatSampler {
public:
	NatSampler() : _count(0), _address_count(0) {}

	// Add a sample: `reported` is the address the remote sees us at (res.to),
	// `from` is the remote's address (res.from).
	void add(const IPv4Address &reported, const IPv4Address &from) {
		if (_count >= NAT_SAMPLER_MAX_SAMPLES) return;

		// Skip loopback and unspecified.
		if (reported.ip[0] == 127 || reported.ip[0] == 0) return;
		if (reported.port == 0) return;

		// Don't add duplicate (from same source IP).
		for (int i = 0; i < _count; i++) {
			if (memcmp(_samples[i].from_ip, from.ip, 4) == 0) return;
		}

		auto &s = _samples[_count++];
		memcpy(s.reported_ip, reported.ip, 4);
		s.reported_port = reported.port;
		memcpy(s.from_ip, from.ip, 4);

		_rebuild_addresses();
	}

	// Determine firewall type based on collected samples.
	uint32_t firewall() const {
		if (_count < NAT_SAMPLER_MIN_SAMPLES) return FIREWALL_UNKNOWN;

		// Count how many times each (ip, port) pair appears.
		int max_hits = 0;
		bool all_unique = true;

		for (int i = 0; i < _address_count; i++) {
			int hits = 0;
			for (int j = 0; j < _count; j++) {
				if (memcmp(_addresses[i].ip, _samples[j].reported_ip, 4) == 0 &&
					_addresses[i].port == _samples[j].reported_port) {
					hits++;
				}
			}
			if (hits > max_hits) max_hits = hits;
			if (hits > 1) all_unique = false;
		}

		// If max_hits >= 3, our external address is consistent.
		if (max_hits >= 3) return FIREWALL_CONSISTENT;

		// All unique → random NAT.
		if (all_unique && _count >= NAT_SAMPLER_MIN_SAMPLES) return FIREWALL_RANDOM;

		// Edge case: same IP, different ports with 2 hits.
		// Check if all samples have the same IP (port-restricted cone).
		if (_count >= NAT_SAMPLER_MIN_SAMPLES) {
			bool same_ip = true;
			for (int i = 1; i < _count; i++) {
				if (memcmp(_samples[0].reported_ip, _samples[i].reported_ip, 4) != 0) {
					same_ip = false;
					break;
				}
			}
			if (same_ip && _address_count == 1) {
				// Same IP and port for all → OPEN (no NAT at all or port forwarded).
				return FIREWALL_OPEN;
			}
			if (same_ip) {
				// Same IP but different ports → CONSISTENT (port-restricted cone).
				return FIREWALL_CONSISTENT;
			}
		}

		return FIREWALL_UNKNOWN;
	}

	// Get the deduplicated external addresses we've observed.
	const IPv4Address *addresses() const { return _addresses; }
	int address_count() const { return _address_count; }
	int sample_count() const { return _count; }

	void reset() {
		_count = 0;
		_address_count = 0;
	}

private:
	void _rebuild_addresses() {
		_address_count = 0;
		for (int i = 0; i < _count; i++) {
			bool found = false;
			for (int j = 0; j < _address_count; j++) {
				if (memcmp(_addresses[j].ip, _samples[i].reported_ip, 4) == 0 &&
					_addresses[j].port == _samples[i].reported_port) {
					found = true;
					break;
				}
			}
			if (!found && _address_count < NAT_SAMPLER_MAX_ADDRESSES) {
				memcpy(_addresses[_address_count].ip, _samples[i].reported_ip, 4);
				_addresses[_address_count].port = _samples[i].reported_port;
				_address_count++;
			}
		}
	}

	struct Sample {
		uint8_t reported_ip[4];
		uint16_t reported_port;
		uint8_t from_ip[4];
	};

	Sample _samples[NAT_SAMPLER_MAX_SAMPLES];
	int _count;

	IPv4Address _addresses[NAT_SAMPLER_MAX_ADDRESSES];
	int _address_count;
};

#endif // HYPERSWARM_NAT_SAMPLER_H