hyperswarm-peer / src / holepuncher.h

#ifndef HYPERSWARM_HOLEPUNCHER_H
#define HYPERSWARM_HOLEPUNCHER_H

// NAT hole-punching state machine.
// Drives UDP probe sequences for consistent, random, and birthday-attack
// NAT traversal strategies. Wire-compatible with holepunchto/hyperdht.

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

#include <cstdint>
#include <cstring>
#include <functional>

extern "C" {
#include <udx.h>
}

// ---- Holepunch stages ----
// Matches JS holepuncher.js state machine.

enum HolepunchStage {
	HP_IDLE = 0,
	HP_ANALYZE,          // Waiting for firewall detection to complete.
	HP_PUNCH_CONSISTENT, // Both OPEN/CONSISTENT — direct probing.
	HP_PUNCH_RANDOM,     // Local CONSISTENT, remote RANDOM — random port guessing.
	HP_PUNCH_BIRTHDAY,   // Local RANDOM, remote CONSISTENT — birthday attack (256 sockets).
	HP_CONNECTED,
	HP_FAILED
};

// ---- Constants ----

static constexpr int HP_MAX_ADDRESSES = 8;
static constexpr int HP_TTL_LOW = 5;
static constexpr int HP_TTL_NORMAL = 64;

// Consistent probe: 10 rounds × 1 second = ~10 seconds.
static constexpr int HP_CONSISTENT_PROBE_INTERVAL_MS = 1000;
static constexpr int HP_CONSISTENT_PROBE_MAX_ROUNDS = 10;

// Random/birthday probe: 1750 iterations × 20ms = ~35 seconds.
static constexpr int HP_RANDOM_PROBE_INTERVAL_MS = 20;
static constexpr int HP_RANDOM_PROBE_MAX_ITERATIONS = 1750;

// Birthday attack: 256 sockets (matching JS BIRTHDAY_SOCKETS constant).
static constexpr int HP_BIRTHDAY_SOCKET_COUNT = 256;

// Analyze timeout — fallback to consistent if firewalls unknown.
static constexpr uint64_t HP_ANALYZE_TIMEOUT_MS = 5000;

// ---- Remote peer info for holepunching ----

struct RemotePeerInfo {
	uint32_t firewall;
	IPv4Address addresses[HP_MAX_ADDRESSES];
	int address_count;
	IPv4Address verified_address; // Address confirmed by token echo.
	bool has_verified;
	bool punching; // Remote is actively punching.
};

// ---- Holepuncher ----

class Holepuncher {
public:
	Holepuncher();
	~Holepuncher();

	// Initialize with the UDP socket used for probing.
	void init(udx_socket_t *socket, uint32_t local_firewall);

	// Update remote peer info (called when holepunch payload received).
	void update_remote(const RemotePeerInfo &info);

	// Start the punch sequence. Call tick() periodically after this.
	void start();

	// Stop punching.
	void stop();

	// Full reset — clears ALL state for slot reuse.
	// Must be called when reusing a DhtConnection slot.
	void reset();

	// Tick — advance the probe state machine.
	// Returns true if still punching, false if done or aborted.
	bool tick(uint64_t now_ms);

	// Send a single probe to an address with specified TTL.
	void send_probe(const IPv4Address &to, int ttl);

	bool is_punching() const { return _punching; }
	bool is_connected() const { return _connected; }
	HolepunchStage stage() const { return _stage; }

	// Take ownership of the winning birthday socket (if any).
	// Caller is responsible for closing it when done.
	// Returns nullptr if no birthday attack was used or no winner yet.
	udx_socket_t *take_winning_socket();

	// Called when a holepunch probe is received from the remote
	// on the main socket. The caller dispatches raw UDP to us.
	void on_probe_received(const IPv4Address &from);

	// True if `from`'s IP belongs to the remote peer this puncher is punching.
	// Lets the shared-socket dispatcher route a probe to the OWNING puncher
	// (mirrors JS holepuncher per-socket routing) instead of first-active.
	bool matches_remote(const IPv4Address &from) const { return _is_remote_ip(from); }

	// ---- Callbacks ----

	// Connection succeeded — provides the punched address.
	std::function<void(const IPv4Address &addr)> on_connect;

	// Punch failed / aborted.
	std::function<void()> on_abort;

private:
	// Strategy decision based on local/remote firewall.
	void _decide_strategy();

	// Probe strategies.
	void _consistent_probe(uint64_t now_ms);
	void _random_probe_tick(uint64_t now_ms);
	void _birthday_probe_tick(uint64_t now_ms);

	// Birthday attack lifecycle.
	void _start_birthday_attack();
	void _cleanup_birthday_sockets();

	// Send probe from a specific socket (for birthday attack).
	void _send_probe_from(udx_socket_t *sock, const IPv4Address &to, int ttl);

	// Get the best remote address for probing.
	const IPv4Address &_remote_target() const;

	// Check if an address IP matches any known remote peer IP.
	bool _is_remote_ip(const IPv4Address &from) const;

	// Birthday socket recv callback (static, routes to instance).
	static void _on_birthday_recv(udx_socket_t *socket, ssize_t read_len,
								  const uv_buf_t *buf, const struct sockaddr *from);
	static void _on_birthday_close(udx_socket_t *socket);

	// Handle probe on a birthday socket — determines the winner.
	void _on_birthday_probe(udx_socket_t *winner, const IPv4Address &from);

	// Core state.
	udx_socket_t *_socket;
	uint32_t _local_firewall;
	RemotePeerInfo _remote;
	HolepunchStage _stage;

	bool _punching;
	bool _connected;
	bool _started;
	uint64_t _start_time;
	uint64_t _last_probe_time;
	int _probe_round;

	// Birthday attack state.
	udx_socket_t *_birthday_sockets[HP_BIRTHDAY_SOCKET_COUNT];
	int _birthday_count;
	bool _birthday_active;
	int _birthday_probe_index;
	int _birthday_iteration;
	bool _birthday_low_ttl_round;
	udx_socket_t *_winning_socket;

	// Random probe state (local CONSISTENT → remote RANDOM).
	int _random_iteration;
};

#endif // HYPERSWARM_HOLEPUNCHER_H