zcatp2p/src/stun.zig

//! Módulo STUN - Session Traversal Utilities for NAT (RFC 5389)
//!
//! Cliente STUN para descubrir dirección IP externa y tipo de NAT.

const std = @import("std");
const crypto = @import("crypto.zig");

/// Puerto STUN estándar
pub const STUN_PORT: u16 = 3478;

/// Magic cookie STUN
const MAGIC_COOKIE: u32 = 0x2112A442;

/// Tipos de mensaje STUN
pub const MessageType = enum(u16) {
    binding_request = 0x0001,
    binding_response = 0x0101,
    binding_error = 0x0111,
    _,
};

/// Tipos de atributo STUN
pub const AttributeType = enum(u16) {
    mapped_address = 0x0001,
    response_address = 0x0002,
    change_request = 0x0003,
    source_address = 0x0004,
    changed_address = 0x0005,
    username = 0x0006,
    password = 0x0007,
    message_integrity = 0x0008,
    error_code = 0x0009,
    unknown_attributes = 0x000A,
    reflected_from = 0x000B,
    realm = 0x0014,
    nonce = 0x0015,
    xor_mapped_address = 0x0020,
    software = 0x8022,
    alternate_server = 0x8023,
    fingerprint = 0x8028,
    other_address = 0x802C,
    _,
};

/// Familia de direcciones
pub const AddressFamily = enum(u8) {
    ipv4 = 0x01,
    ipv6 = 0x02,
    _,
};

/// Dirección mapeada
pub const MappedAddress = struct {
    family: AddressFamily,
    port: u16,
    address: union {
        ipv4: [4]u8,
        ipv6: [16]u8,
    },

    pub fn format(self: MappedAddress, buf: []u8) []const u8 {
        if (self.family == .ipv4) {
            return std.fmt.bufPrint(buf, "{}.{}.{}.{}:{}", .{
                self.address.ipv4[0],
                self.address.ipv4[1],
                self.address.ipv4[2],
                self.address.ipv4[3],
                self.port,
            }) catch "";
        }
        return "";
    }
};

/// Mensaje STUN
pub const StunMessage = struct {
    message_type: MessageType,
    transaction_id: [12]u8,
    attributes: std.ArrayListUnmanaged(Attribute),
    allocator: std.mem.Allocator,

    pub const Attribute = struct {
        attr_type: AttributeType,
        data: []const u8,
    };

        pub fn init(io: std.Io, allocator: std.mem.Allocator, msg_type: MessageType) StunMessage {

            var transaction_id: [12]u8 = undefined;

            io.random(&transaction_id);

            return .{

                .message_type = msg_type,

                .transaction_id = transaction_id,

                .attributes = .{},

                .allocator = allocator,

            };

        }

    pub fn deinit(self: *StunMessage) void {
        for (self.attributes.items) |attr| {
            self.allocator.free(attr.data);
        }
        self.attributes.deinit(self.allocator);
    }

    /// Codifica el mensaje STUN
    pub fn encode(self: *StunMessage) ![]u8 {
        // Calcular longitud de atributos
        var attrs_len: usize = 0;
        for (self.attributes.items) |attr| {
            attrs_len += 4 + attr.data.len;
            // Padding a 4 bytes
            if (attr.data.len % 4 != 0) {
                attrs_len += 4 - (attr.data.len % 4);
            }
        }

        const total_len = 20 + attrs_len;
        const buf = try self.allocator.alloc(u8, total_len);
        errdefer self.allocator.free(buf);

        var pos: usize = 0;

        // Header
        std.mem.writeInt(u16, buf[0..2], @intFromEnum(self.message_type), .big);
        std.mem.writeInt(u16, buf[2..4], @intCast(attrs_len), .big);
        std.mem.writeInt(u32, buf[4..8], MAGIC_COOKIE, .big);
        @memcpy(buf[8..20], &self.transaction_id);
        pos = 20;

        // Atributos
        for (self.attributes.items) |attr| {
            std.mem.writeInt(u16, buf[pos..][0..2], @intFromEnum(attr.attr_type), .big);
            std.mem.writeInt(u16, buf[pos + 2 ..][0..2], @intCast(attr.data.len), .big);
            @memcpy(buf[pos + 4 .. pos + 4 + attr.data.len], attr.data);
            pos += 4 + attr.data.len;

            // Padding
            const pad = (4 - (attr.data.len % 4)) % 4;
            if (pad > 0) {
                @memset(buf[pos .. pos + pad], 0);
                pos += pad;
            }
        }

        return buf;
    }

    /// Decodifica un mensaje STUN
    pub fn decode(allocator: std.mem.Allocator, data: []const u8) !StunMessage {
        if (data.len < 20) return error.MessageTooShort;

        const msg_type: MessageType = @enumFromInt(std.mem.readInt(u16, data[0..2], .big));
        const msg_len = std.mem.readInt(u16, data[2..4], .big);
        const magic = std.mem.readInt(u32, data[4..8], .big);

        if (magic != MAGIC_COOKIE) return error.InvalidMagicCookie;
        if (data.len < 20 + msg_len) return error.MessageTooShort;

        var msg = StunMessage{
            .message_type = msg_type,
            .transaction_id = data[8..20].*,
            .attributes = .{},
            .allocator = allocator,
        };
        errdefer msg.deinit();

        // Parsear atributos
        var pos: usize = 20;
        const end = 20 + msg_len;

        while (pos + 4 <= end) {
            const attr_type: AttributeType = @enumFromInt(std.mem.readInt(u16, data[pos..][0..2], .big));
            const attr_len = std.mem.readInt(u16, data[pos + 2 ..][0..2], .big);
            pos += 4;

            if (pos + attr_len > end) break;

            const attr_data = try allocator.dupe(u8, data[pos .. pos + attr_len]);
            try msg.attributes.append(allocator, .{
                .attr_type = attr_type,
                .data = attr_data,
            });

            pos += attr_len;
            // Skip padding
            pos += (4 - (attr_len % 4)) % 4;
        }

        return msg;
    }

    /// Obtiene la dirección XOR-MAPPED-ADDRESS
    pub fn getXorMappedAddress(self: *StunMessage) ?MappedAddress {
        for (self.attributes.items) |attr| {
            if (attr.attr_type == .xor_mapped_address) {
                return parseXorMappedAddress(attr.data, self.transaction_id);
            }
        }
        // Fallback a MAPPED-ADDRESS
        for (self.attributes.items) |attr| {
            if (attr.attr_type == .mapped_address) {
                return parseMappedAddress(attr.data);
            }
        }
        return null;
    }

    /// Obtiene OTHER-ADDRESS (para detección de NAT)
    pub fn getOtherAddress(self: *StunMessage) ?MappedAddress {
        for (self.attributes.items) |attr| {
            if (attr.attr_type == .other_address or attr.attr_type == .changed_address) {
                return parseMappedAddress(attr.data);
            }
        }
        return null;
    }
};

fn parseMappedAddress(data: []const u8) ?MappedAddress {
    if (data.len < 8) return null;

    const family: AddressFamily = @enumFromInt(data[1]);

    if (family == .ipv4 and data.len >= 8) {
        return .{
            .family = .ipv4,
            .port = std.mem.readInt(u16, data[2..4], .big),
            .address = .{ .ipv4 = data[4..8].* },
        };
    } else if (family == .ipv6 and data.len >= 20) {
        return .{
            .family = .ipv6,
            .port = std.mem.readInt(u16, data[2..4], .big),
            .address = .{ .ipv6 = data[4..20].* },
        };
    }
    return null;
}

fn parseXorMappedAddress(data: []const u8, transaction_id: [12]u8) ?MappedAddress {
    if (data.len < 8) return null;

    const family: AddressFamily = @enumFromInt(data[1]);

    // XOR port with magic cookie high bytes
    const port = std.mem.readInt(u16, data[2..4], .big) ^ @as(u16, @truncate(MAGIC_COOKIE >> 16));

    if (family == .ipv4 and data.len >= 8) {
        // XOR address with magic cookie
        var addr: [4]u8 = data[4..8].*;
        const magic_bytes = std.mem.toBytes(std.mem.nativeToBig(u32, MAGIC_COOKIE));
        for (0..4) |i| {
            addr[i] ^= magic_bytes[i];
        }
        return .{
            .family = .ipv4,
            .port = port,
            .address = .{ .ipv4 = addr },
        };
    } else if (family == .ipv6 and data.len >= 20) {
        // XOR address with magic cookie + transaction_id
        var addr: [16]u8 = data[4..20].*;
        const magic_bytes = std.mem.toBytes(std.mem.nativeToBig(u32, MAGIC_COOKIE));
        for (0..4) |i| {
            addr[i] ^= magic_bytes[i];
        }
        for (0..12) |i| {
            addr[4 + i] ^= transaction_id[i];
        }
        return .{
            .family = .ipv6,
            .port = port,
            .address = .{ .ipv6 = addr },
        };
    }
    return null;
}

/// Tipo de NAT detectado
pub const NatType = enum {
    unknown,
    open_internet, // Sin NAT
    full_cone, // Cualquier host externo puede enviar
    restricted, // Solo hosts a los que hemos enviado
    port_restricted, // Solo host:port a los que hemos enviado
    symmetric, // Diferente mapeo por destino
    blocked, // UDP bloqueado

    pub fn canPunch(self: NatType) bool {
        return switch (self) {
            .open_internet, .full_cone, .restricted, .port_restricted => true,
            .symmetric, .blocked, .unknown => false,
        };
    }

    pub fn needsRelay(self: NatType) bool {
        return self == .symmetric or self == .blocked;
    }
};

/// Cliente STUN
pub const StunClient = struct {
    io: std.Io,
    allocator: std.mem.Allocator,
    servers: std.ArrayListUnmanaged([]const u8),
    socket: ?std.Io.net.Socket,
    external_address: ?MappedAddress,
    nat_type: NatType,

    pub fn init(io: std.Io, allocator: std.mem.Allocator) StunClient {
        return .{
            .io = io,
            .allocator = allocator,
            .servers = .{},
            .socket = null,
            .external_address = null,
            .nat_type = .unknown,
        };
    }

    pub fn deinit(self: *StunClient) void {
        if (self.socket) |*sock| {
            sock.close(self.io);
        }
        for (self.servers.items) |server| {
            self.allocator.free(server);
        }
        self.servers.deinit(self.allocator);
    }

    /// Añade un servidor STUN
    pub fn addServer(self: *StunClient, server: []const u8) !void {
        const owned = try self.allocator.dupe(u8, server);
        try self.servers.append(self.allocator, owned);
    }

    /// Crea el socket UDP
    pub fn createSocket(self: *StunClient) !void {
        if (self.socket != null) return;

        const addr = std.Io.net.IpAddress.unspecified(0);
        self.socket = try std.Io.net.bind(&addr, self.io, .{ .mode = .dgram });
    }

    /// Envía un Binding Request a un servidor
    pub fn sendBindingRequest(self: *StunClient, server_addr: std.net.Address) !StunMessage {
        if (self.socket == null) try self.createSocket();

        var request = StunMessage.init(self.io, self.allocator, .binding_request);
        errdefer request.deinit();

        const encoded = try request.encode();
        defer self.allocator.free(encoded);

        // Convertir std.net.Address a std.Io.net.IpAddress
        const io_addr = switch (server_addr.any.family) {
            std.posix.AF.INET => std.Io.net.IpAddress{ .ip4 = server_addr.in },
            std.posix.AF.INET6 => std.Io.net.IpAddress{ .ip6 = server_addr.in6 },
            else => return error.InvalidAddressFamily,
        };

        try self.socket.?.send(self.io, &io_addr, encoded);

        return request;
    }

    /// Recibe una respuesta STUN
    pub fn receiveResponse(self: *StunClient, timeout_ms: u32) !StunMessage {
        const sock = self.socket orelse return error.SocketNotCreated;

        var buf: [1024]u8 = undefined;
        const timeout = std.Io.Timeout{ .duration = .{
            .raw = .{ .nanoseconds = @as(i96, timeout_ms) * std.time.ns_per_ms },
            .clock = .real,
        } };
        const msg = try sock.receiveTimeout(self.io, &buf, timeout);

        return StunMessage.decode(self.allocator, msg.data);
    }

    /// Descubre la dirección externa
    pub fn discoverExternalAddress(self: *StunClient) !?MappedAddress {
        if (self.servers.items.len == 0) {
            // Añadir servidores por defecto
            try self.addServer("stun.l.google.com:19302");
            try self.addServer("stun.syncthing.net:3478");
        }

        for (self.servers.items) |server| {
            const result = self.queryServer(server) catch continue;
            if (result) |addr| {
                self.external_address = addr;
                return addr;
            }
        }

        return null;
    }

    fn queryServer(self: *StunClient, server: []const u8) !?MappedAddress {
        // Parsear host:port
        var host_end: usize = server.len;
        var port: u16 = STUN_PORT;

        if (std.mem.lastIndexOf(u8, server, ":")) |colon| {
            host_end = colon;
            port = std.fmt.parseInt(u16, server[colon + 1 ..], 10) catch STUN_PORT;
        }
        const host = server[0..host_end];

        // Resolver DNS (simplificado - solo IPv4)
        const addr = try parseIpv4(host, port);

        const request = try self.sendBindingRequest(addr);
        var response = try self.receiveResponse(3000);
        defer response.deinit();

        // Verificar transaction ID
        if (!std.mem.eql(u8, &request.transaction_id, &response.transaction_id)) {
            return error.TransactionIdMismatch;
        }

        return response.getXorMappedAddress();
    }

    /// Detecta el tipo de NAT
    pub fn detectNatType(self: *StunClient) !NatType {
        // Algoritmo simplificado de detección de NAT
        // Para detección completa se necesitan 2 servidores STUN con 2 IPs cada uno

        const external = try self.discoverExternalAddress();
        if (external == null) {
            self.nat_type = .blocked;
            return .blocked;
        }

        // Verificar si la IP externa coincide con la local (sin NAT)
        // Simplificado: asumir que hay NAT
        self.nat_type = .restricted;
        return .restricted;
    }
};

fn parseIpv4(host: []const u8, port: u16) !std.net.Address {
    // Parsear IP directamente o usar DNS
    var octets: [4]u8 = undefined;
    var octet_idx: usize = 0;
    var current: u16 = 0;

    for (host) |c| {
        if (c == '.') {
            if (octet_idx >= 4) return error.InvalidAddress;
            octets[octet_idx] = @intCast(current);
            octet_idx += 1;
            current = 0;
        } else if (c >= '0' and c <= '9') {
            current = current * 10 + (c - '0');
            if (current > 255) return error.InvalidAddress;
        } else {
            // Es un hostname, no una IP
            // Usar lookup DNS sería necesario aquí
            // Por ahora, usar Google STUN como fallback
            return std.net.Address.initIp4(.{ 142, 250, 187, 127 }, port);
        }
    }

    if (octet_idx == 3) {
        octets[3] = @intCast(current);
        return std.net.Address.initIp4(octets, port);
    }

    return error.InvalidAddress;
}

// =============================================================================
// Tests
// =============================================================================

test "stun message encode/decode" {
    const allocator = std.testing.allocator;
    const io = std.testing.io;

    var msg = StunMessage.init(io, allocator, .binding_request);
    defer msg.deinit();

    const encoded = try msg.encode();
    defer allocator.free(encoded);

    try std.testing.expect(encoded.len == 20); // Header only

    var decoded = try StunMessage.decode(allocator, encoded);
    defer decoded.deinit();

    try std.testing.expect(decoded.message_type == .binding_request);
    try std.testing.expectEqualSlices(u8, &msg.transaction_id, &decoded.transaction_id);
}

test "parse xor mapped address ipv4" {
    const transaction_id = [_]u8{0} ** 12;

    // XOR-MAPPED-ADDRESS para 192.0.2.1:32853
    // Family: 0x01 (IPv4), XOR'd Port: 0x1234, XOR'd IP: XOR con magic cookie
    const port_xored: u16 = 32853 ^ 0x2112; // XOR with high bytes of magic cookie
    const ip = [4]u8{ 192 ^ 0x21, 0 ^ 0x12, 2 ^ 0xA4, 1 ^ 0x42 }; // XOR with magic cookie

    var data: [8]u8 = undefined;
    data[0] = 0; // Reserved
    data[1] = 0x01; // IPv4
    std.mem.writeInt(u16, data[2..4], port_xored, .big);
    @memcpy(data[4..8], &ip);

    const addr = parseXorMappedAddress(&data, transaction_id);
    try std.testing.expect(addr != null);
    try std.testing.expect(addr.?.family == .ipv4);
    try std.testing.expect(addr.?.port == 32853);
    try std.testing.expectEqual([4]u8{ 192, 0, 2, 1 }, addr.?.address.ipv4);
}

test "stun client init" {
    const allocator = std.testing.allocator;
    const io = std.testing.io;

    var client = StunClient.init(io, allocator);
    defer client.deinit();

    try std.testing.expect(client.nat_type == .unknown);
}

test "nat type capabilities" {
    try std.testing.expect(NatType.full_cone.canPunch());
    try std.testing.expect(NatType.restricted.canPunch());
    try std.testing.expect(!NatType.symmetric.canPunch());
    try std.testing.expect(!NatType.blocked.canPunch());
    try std.testing.expect(NatType.symmetric.needsRelay());
}