internal/router: add UniFi Gateway support

UniFi Gateway (USG) uses its own DNS forwarding rule, which is
configured default in /etc/dnsmasq.conf file. Adding ctrld own config in
/etc/dnsmasq.d won't take effects. Instead, we must make changes
directly to /etc/dnsmasq.conf, configuring ctrld as the only upstream.

internal/router/edgeos.go

@@ -1,37 +1,119 @@
 package router
 
 import (
+	"bufio"
+	"bytes"
 	"fmt"
 	"os"
 	"os/exec"
+	"strings"
 )
 
-const edgeOSDNSMasqConfigPath = "/etc/dnsmasq.d/dnsmasq-zzz-ctrld.conf"
+const (
+	edgeOSDNSMasqConfigPath    = "/etc/dnsmasq.d/dnsmasq-zzz-ctrld.conf"
+	UsgDNSMasqConfigPath       = "/etc/dnsmasq.conf"
+	UsgDNSMasqBackupConfigPath = "/etc/dnsmasq.conf.bak"
+)
+
+var (
+	isUSG bool
+)
 
 func setupEdgeOS() error {
+	if isUSG {
+		return setupUSG()
+	}
+	return setupUDM()
+}
+
+func setupUDM() error {
 	// Disable dnsmasq as DNS server.
 	dnsMasqConfigContent, err := dnsMasqConf()
 	if err != nil {
-		return err
+		return fmt.Errorf("setupUDM: generating dnsmasq config: %w", err)
 	}
 	if err := os.WriteFile(edgeOSDNSMasqConfigPath, []byte(dnsMasqConfigContent), 0600); err != nil {
-		return err
+		return fmt.Errorf("setupUDM: generating dnsmasq config: %w", err)
 	}
 	// Restart dnsmasq service.
 	if err := restartDNSMasq(); err != nil {
-		return err
+		return fmt.Errorf("setupUDM: restartDNSMasq: %w", err)
+	}
+	return nil
+}
+
+func setupUSG() error {
+	// On USG, dnsmasq is configured to forward queries to external provider by default.
+	// So instead of generating config in /etc/dnsmasq.d, we need to create a backup of
+	// the config, then modify it to forward queries to ctrld listener.
+
+	// Creating a backup.
+	buf, err := os.ReadFile(UsgDNSMasqConfigPath)
+	if err != nil {
+		return fmt.Errorf("setupUSG: reading current config: %w", err)
+	}
+	if err := os.WriteFile(UsgDNSMasqBackupConfigPath, buf, 0600); err != nil {
+		return fmt.Errorf("setupUSG: backup current config: %w", err)
+	}
+
+	// Removing all configured upstreams.
+	var sb strings.Builder
+	scanner := bufio.NewScanner(bytes.NewReader(buf))
+	for scanner.Scan() {
+		line := scanner.Text()
+		if strings.HasPrefix(line, "server=") {
+			continue
+		}
+		if strings.HasPrefix(line, "all-servers") {
+			continue
+		}
+		sb.WriteString(line)
+	}
+
+	// Adding ctrld listener as the only upstream.
+	dnsMasqConfigContent, err := dnsMasqConf()
+	if err != nil {
+		return fmt.Errorf("setupUSG: generating dnsmasq config: %w", err)
+	}
+	sb.WriteString("\n")
+	sb.WriteString(dnsMasqConfigContent)
+	if err := os.WriteFile(UsgDNSMasqConfigPath, []byte(sb.String()), 0644); err != nil {
+		return fmt.Errorf("setupUSG: writing dnsmasq config: %w", err)
+	}
+
+	// Restart dnsmasq service.
+	if err := restartDNSMasq(); err != nil {
+		return fmt.Errorf("setupUSG: restartDNSMasq: %w", err)
 	}
 	return nil
 }
 
 func cleanupEdgeOS() error {
+	if isUSG {
+		return cleanupUSG()
+	}
+	return cleanupUDM()
+}
+
+func cleanupUDM() error {
 	// Remove the custom dnsmasq config
 	if err := os.Remove(edgeOSDNSMasqConfigPath); err != nil {
-		return err
+		return fmt.Errorf("cleanupUDM: os.Remove: %w", err)
 	}
 	// Restart dnsmasq service.
 	if err := restartDNSMasq(); err != nil {
-		return err
+		return fmt.Errorf("cleanupUDM: restartDNSMasq: %w", err)
+	}
+	return nil
+}
+
+func cleanupUSG() error {
+	if err := os.Rename(UsgDNSMasqBackupConfigPath, UsgDNSMasqConfigPath); err != nil {
+		return fmt.Errorf("cleanupUSG: os.Rename: %w", err)
+	}
+	// Restart dnsmasq service.
+	if err := restartDNSMasq(); err != nil {
+		return fmt.Errorf("cleanupUSG: restartDNSMasq: %w", err)
 	}
 	return nil
 }

internal/router/router.go

@@ -220,8 +220,10 @@ func distroName() string {
 	case bytes.HasPrefix(unameO(), []byte("Tomato")):
 		return Tomato
 	case haveDir("/config/scripts/post-config.d"):
+		checkUSG()
 		return EdgeOS
 	case haveFile("/etc/ubnt/init/vyatta-router"):
+		checkUSG()
 		return EdgeOS // For 2.x
 	case isPfsense():
 		return Pfsense
@@ -253,3 +255,8 @@ func isPfsense() bool {
 	b, err := os.ReadFile("/etc/platform")
 	return err == nil && bytes.HasPrefix(b, []byte("pfSense"))
 }
+
+func checkUSG() {
+	out, _ := exec.Command("mca-cli-op", "info").Output()
+	isUSG = bytes.Contains(out, []byte("UniFi-Gateway-"))
+}