From db1fbeb92ac89ffc3e7823aa82cce350b5875ede Mon Sep 17 00:00:00 2001
From: Muaz Ahmad <mahmad2000@protonmail.com>
Date: Mon, 13 May 2024 01:34:30 +0500
Subject: [PATCH] Add additional mapping for local domain

---
 src/hosts.zig | 11 ++++++-----
 src/main.zig  |  2 +-
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/hosts.zig b/src/hosts.zig
index 8748e9a..a1a8b31 100644
--- a/src/hosts.zig
+++ b/src/hosts.zig
@@ -5,8 +5,8 @@ const hosts_header = "# local-etc-hosts-updater";
 
 const FILE_LINE_BUFF_SIZE = 100;
 
-pub fn update_hosts(ip: util.IP) !void {
-    try create_tmp_hosts(ip);
+pub fn update_hosts(ip: util.IP, domain: util.Domain) !void {
+    try create_tmp_hosts(ip, domain);
     try move_tmp_hosts();
 }
 
@@ -61,7 +61,7 @@ fn purge_existing(old_hosts: std.fs.File, tmp_hosts: std.fs.File) !void {
     }
 }
 
-fn append_new(tmp_hosts: std.fs.File, addr: std.net.Address) !void {
+fn append_new(tmp_hosts: std.fs.File, addr: std.net.Address, local_domain: util.Domain) !void {
     const writer = tmp_hosts.writer();
     var buff = [_]u8{0x00} ** FILE_LINE_BUFF_SIZE;
 
@@ -85,6 +85,7 @@ fn append_new(tmp_hosts: std.fs.File, addr: std.net.Address) !void {
     while (subdomains.next()) |subdomain| {
         try std.fmt.format(writer, "{s} {s}.{s}\n", .{ addr_str, subdomain, domain });
     }
+    try std.fmt.format(writer, "{s} {s}\n", .{ addr_str, local_domain.name });
 
     // add extra final newline
     try writer.writeByte('\n');
@@ -95,7 +96,7 @@ fn get_subdomains() !std.mem.SplitIterator(u8, .scalar) {
     return std.mem.splitScalar(u8, std.mem.span(subdomain_list), ' ');
 }
 
-fn create_tmp_hosts(ip: util.IP) !void {
+fn create_tmp_hosts(ip: util.IP, domain: util.Domain) !void {
     var old_hosts = try std.fs.openFileAbsoluteZ(try util.getenv("OLD_HOSTS_PATH"), .{ .mode = .read_only });
     defer old_hosts.close();
     var tmp_hosts = try std.fs.createFileAbsoluteZ(try util.getenv("TMP_HOSTS_PATH"), .{ .truncate = true });
@@ -103,7 +104,7 @@ fn create_tmp_hosts(ip: util.IP) !void {
 
     try purge_existing(old_hosts, tmp_hosts);
     if (ip) |addr| {
-        try append_new(tmp_hosts, addr);
+        try append_new(tmp_hosts, addr, domain);
     }
 }
 
diff --git a/src/main.zig b/src/main.zig
index de7ed0d..819039b 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -7,5 +7,5 @@ pub fn main() !void {
     try util.check_perms();
     const domain, const ip_info = try util.get_input();
     const ip = try mdns.get_mdns(domain, ip_info);
-    try hosts.update_hosts(ip);
+    try hosts.update_hosts(ip, domain);
 }