From d9d14dfad21e24d03d6c8a88a04ef5f1d579d776 Mon Sep 17 00:00:00 2001
From: wencaiwulue <895703375@qq.com>
Date: Wed, 18 Aug 2021 23:23:47 +0800
Subject: [PATCH] support specific namespace short dns resolve

---
 TODO.MD            |  4 +++-
 dns/dns_unix.go    | 10 ++++++++--
 dns/dns_windows.go | 18 +++++++++++++++++-
 pkg/main.go        |  2 +-
 4 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/TODO.MD b/TODO.MD
index eca20002..ed4c6b00 100644
--- a/TODO.MD
+++ b/TODO.MD
@@ -2,4 +2,6 @@
 
 ## 域名解析功能 --完成
 
-## 多个service inbound -- 完成
\ No newline at end of file
+## 多个service inbound -- 完成
+
+## 短域名解析
\ No newline at end of file
diff --git a/dns/dns_unix.go b/dns/dns_unix.go
index ed8610d5..9310078e 100644
--- a/dns/dns_unix.go
+++ b/dns/dns_unix.go
@@ -10,12 +10,18 @@ import (
 	"path/filepath"
 )
 
-func DNS(ip string) error {
+func DNS(ip string, namespace string) error {
 	var err error
 	if err = os.MkdirAll(filepath.Join("/", "etc", "resolver"), fs.ModePerm); err != nil {
 		log.Error(err)
 	}
 	filename := filepath.Join("/", "etc", "resolver", "local")
 	fileContent := "nameserver " + ip
-	return ioutil.WriteFile(filename, []byte(fileContent), fs.ModePerm)
+	_ = ioutil.WriteFile(filename, []byte(fileContent), fs.ModePerm)
+
+	filename = filepath.Join("/", "etc", "resolver", namespace)
+	fileContent = "nameserver " + ip + "\nsearch " + namespace + ".svc.cluster.local svc.cluster.local cluster.local\noptions ndots:5"
+	_ = ioutil.WriteFile(filename, []byte(fileContent), fs.ModePerm)
+
+	return nil
 }
diff --git a/dns/dns_windows.go b/dns/dns_windows.go
index 8c78afaa..ce83f9bb 100644
--- a/dns/dns_windows.go
+++ b/dns/dns_windows.go
@@ -9,7 +9,7 @@ import (
 	"os/exec"
 )
 
-func DNS(ip string) error {
+func DNS(ip string, namespace string) error {
 	tunName := os.Getenv("tunName")
 	fmt.Println("tun name: " + tunName)
 	args := []string{
@@ -24,9 +24,25 @@ func DNS(ip string) error {
 	output, err := exec.Command("netsh", args...).CombinedOutput()
 	fmt.Println(exec.Command("netsh", args...).Args)
 	log.Info(string(output))
+	_ = addNicSuffix(namespace)
 	if err != nil {
 		log.Error(err)
 		return nil
 	}
 	return nil
 }
+
+// @see https://docs.microsoft.com/en-us/powershell/module/dnsclient/set-dnsclientglobalsetting?view=windowsserver2019-ps#example-1--set-the-dns-suffix-search-list
+func addNicSuffix(namespace string) error {
+	cmd := exec.Command("PowerShell", []string{
+		"Set-DnsClientGlobalSetting",
+		"-SuffixSearchList",
+		fmt.Sprintf("@(\"%s.svc.cluster.local\", \"svc.cluster.local\")", namespace),
+	}...)
+	output, err := cmd.CombinedOutput()
+	if err != nil {
+		log.Warn(err)
+	}
+	log.Info(string(output))
+	return err
+}
diff --git a/pkg/main.go b/pkg/main.go
index 3256ff46..7899e4fd 100644
--- a/pkg/main.go
+++ b/pkg/main.go
@@ -174,7 +174,7 @@ func main() {
 	}
 	//time.Sleep(time.Second * 5)
 	dnsServiceIp := util.GetDNSServiceIpFromPod(clientset, restclient, config, util.TrafficManager, namespace)
-	if err := dns.DNS(dnsServiceIp); err != nil {
+	if err := dns.DNS(dnsServiceIp, namespace); err != nil {
 		log.Fatal(err)
 	}
 	if runtime.GOOS == "windows" {