From 43c93dc54a822bfb3231ba82a8b4da05346393ca Mon Sep 17 00:00:00 2001
From: Muaz Ahmad <mahmad2000@protonmail.com>
Date: Sat, 11 May 2024 15:14:32 +0500
Subject: [PATCH] Add UDP send functionality

---
 build.zig    |  1 +
 src/mdns.zig | 74 ++++++++++++++++++++++++++++++++++++++--------------
 src/util.zig | 19 ++++++++++++--
 3 files changed, 73 insertions(+), 21 deletions(-)

diff --git a/build.zig b/build.zig
index 4c75803..a401b77 100644
--- a/build.zig
+++ b/build.zig
@@ -21,6 +21,7 @@ pub fn build(b: *std.Build) void {
         .target = target,
         .optimize = optimize,
     });
+    exe.linkLibC();
 
     // This declares intent for the executable to be installed into the
     // standard location when the user invokes the "install" step (the default
diff --git a/src/mdns.zig b/src/mdns.zig
index 27ef84a..107a8ea 100644
--- a/src/mdns.zig
+++ b/src/mdns.zig
@@ -1,32 +1,68 @@
 const std = @import("std");
 const util = @import("util.zig");
 
-pub fn get_mdns(domain: util.Domain, ip_ver: util.IP_VER_ENUM) !util.IPs {
-    var listener = try init(ip_ver);
-    try send_query(&listener, domain, ip_ver);
-    return receive_query(&listener);
+const MDNSError = error{
+    Unimplemented,
+    SocketInitFail,
+    UDPConnectFail,
+    UDPSendFail,
+};
+
+pub fn get_mdns(domain: util.Domain, ip_info: util.IPInfo) !util.IPs {
+    const sock = try send_query(domain, ip_info);
+    return receive_query(sock);
 }
 
-fn init(ip_ver: util.IP_VER_ENUM) !std.net.Server {
-    return try std.net.Address.listen(try std.net.Address.parseIp(switch (ip_ver) {
-        util.IP_VER_ENUM.IPv4 => "0.0.0.0",
-        util.IP_VER_ENUM.IPv6 => "::",
-    }, 5353), .{});
+const socket = c_int;
+
+fn get_mdns_socket(ip_info: util.IPInfo) !socket {
+    const target_addr: []const u8 = switch (ip_info.version) {
+        util.IP_VER_ENUM.IPv4 => "224.0.0.251",
+        util.IP_VER_ENUM.IPv6 => blk: {
+            var buf = [_]u8{0x00} ** 100;
+            var byte_buf = std.io.fixedBufferStream(&buf);
+            const writer = byte_buf.writer();
+            try std.fmt.format(writer, "ff02::fb%{s}", .{ip_info.interface.?});
+            break :blk buf[0..writer.context.pos];
+        },
+    };
+    const addr = try std.net.Address.resolveIp(target_addr, 5353);
+    const sock = std.c.socket(switch (ip_info.version) {
+        util.IP_VER_ENUM.IPv4 => std.posix.AF.INET,
+        util.IP_VER_ENUM.IPv6 => std.posix.AF.INET6,
+    }, std.c.SOCK.DGRAM, std.c.IPPROTO.UDP);
+    if (sock == -1) {
+        return MDNSError.SocketInitFail;
+    }
+    if (std.c.connect(sock, &addr.any, addr.getOsSockLen()) == -1) {
+        return MDNSError.UDPConnectFail;
+    }
+    return sock;
 }
 
-fn send_query(listener: *std.net.Server, domain: util.Domain, ip_ver: util.IP_VER_ENUM) !void {
-    const header =
-        [_]u8{0x00} ** 4 ++ // transaction id + is question
-        [_]u8{ 0x00, 0x01 } ++ // 1 question
-        [_]u8{0x00} ** 6; // misc flags
-    _ = listener;
+fn construct_mdns_query(domain: util.Domain, ip_info: util.IPInfo, buff: []u8) !usize {
     _ = domain;
-    _ = header;
-    _ = ip_ver;
+    _ = ip_info;
+    var n: usize = 0;
+    var buff_writer = std.io.fixedBufferStream(buff);
+    n += try buff_writer.write("test");
+    return n;
 }
 
-fn receive_query(listener: *std.net.Server) !util.IPs {
-    _ = listener;
+fn send_query(domain: util.Domain, ip_info: util.IPInfo) !socket {
+    const sock = try get_mdns_socket(ip_info);
+    errdefer _ = std.c.close(sock);
+    var buff = [_]u8{0x00} ** 100;
+
+    const n = try construct_mdns_query(domain, ip_info, &buff);
+    if (std.c.send(sock, &buff, n, 0) == -1) {
+        return MDNSError.UDPSendFail;
+    }
+    return sock;
+}
+
+fn receive_query(sock: socket) !util.IPs {
+    _ = sock;
     return util.IPs{
         .v4 = "",
         .v6 = "",
diff --git a/src/util.zig b/src/util.zig
index fc6c8c3..a63b190 100644
--- a/src/util.zig
+++ b/src/util.zig
@@ -4,6 +4,8 @@ const ArgError = error{
     NotEnoughArgs,
     BadDomain,
     InvalidAddressVer,
+    InterfaceRequired,
+    InvalidInterface,
 };
 
 pub const IP_VER_ENUM = enum(u3) {
@@ -16,6 +18,11 @@ pub const IPs = struct {
     v6: []const u8,
 };
 
+pub const IPInfo = struct {
+    version: IP_VER_ENUM,
+    interface: ?[:0]const u8,
+};
+
 pub const Domain = struct {
     name: [:0]const u8,
     labels: [5][]const u8,
@@ -28,7 +35,7 @@ pub fn check_perms() !void {
 
 pub fn get_input() !struct {
     Domain,
-    IP_VER_ENUM,
+    IPInfo,
 } {
     var args = std.process.args();
     _ = args.next(); // consume calling binary arg
@@ -39,7 +46,15 @@ pub fn get_input() !struct {
     const ip_ver_str = args.next() orelse return ArgError.NotEnoughArgs;
     const ip_ver = std.meta.intToEnum(IP_VER_ENUM, std.fmt.parseInt(u3, ip_ver_str, 10) catch return ArgError.InvalidAddressVer) catch return ArgError.InvalidAddressVer;
 
-    return .{ domain, ip_ver };
+    var iface: ?[:0]const u8 = null;
+    if (ip_ver == IP_VER_ENUM.IPv6) {
+        iface = args.next() orelse return ArgError.InterfaceRequired;
+        if (std.c.if_nametoindex(iface.?) == 0) {
+            return ArgError.InvalidInterface;
+        }
+    }
+
+    return .{ domain, IPInfo{ .version = ip_ver, .interface = iface } };
 }
 
 fn check_domain(domain_str: [:0]const u8) !Domain {