[ndns] auto reload record file

Author: ccsexyz

ndns/ndns.go

@@ -9,6 +9,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/fsnotify/fsnotify"
 	"github.com/miekg/dns"
 )
 
@@ -21,23 +22,33 @@ type record struct {
 }
 
 var (
-	file string
-	tcp  bool
-	udp  bool
-	addr string
-
-	recordMap map[string]*record
+	file     string
+	tcp      bool
+	udp      bool
+	addr     string
+	autoload bool
+
+	recordLock sync.Mutex
+	recordMap  map[string]*record
 )
 
 func init() {
 	flag.StringVar(&file, "f", "record.json", "record file")
 	flag.BoolVar(&tcp, "t", false, "listen tcp")
 	flag.BoolVar(&udp, "u", true, "listen udp")
 	flag.StringVar(&addr, "l", ":53", "listen address")
+	flag.BoolVar(&autoload, "a", true, "auto reload record file")
 
 	recordMap = make(map[string]*record)
 }
 
+func lookUpRecord(k string) (r *record, ok bool) {
+	recordLock.Lock()
+	r, ok = recordMap[k]
+	recordLock.Unlock()
+	return
+}
+
 type handler struct {
 	isTcpServer bool
 }
@@ -51,7 +62,7 @@ func (this *handler) packAnswers(msg *dns.Msg, qtype uint16, domain string) {
 		return
 	}
 
-	r, ok := recordMap[domain]
+	r, ok := lookUpRecord(domain)
 	if !ok {
 		log.Println("domain", domain, "not found")
 		return
@@ -124,18 +135,55 @@ func runUdpServer() {
 	}
 }
 
-func main() {
-	flag.Parse()
-
+func loadFile() error {
 	records, err := readRecords(file)
 	if err != nil {
-		log.Fatal(err)
+		return err
 	}
 
+	recordLock.Lock()
+	defer recordLock.Unlock()
+
+	recordMap = make(map[string]*record)
+
 	for _, v := range records {
 		recordMap[v.Host+"."] = v
 	}
 
+	return nil
+}
+
+func updater() {
+	watcher, err := fsnotify.NewWatcher()
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer watcher.Close()
+
+	err = watcher.Add(file)
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer watcher.Remove(file)
+
+	for event := range watcher.Events {
+		log.Println("updater go event", event)
+		err = loadFile()
+		if err != nil {
+			log.Println("try to load file, but failed", err)
+		}
+		watcher.Remove(file)
+		watcher.Add(file)
+	}
+}
+
+func main() {
+	flag.Parse()
+
+	if err := loadFile(); err != nil {
+		log.Fatal(err)
+	}
+
 	var wg sync.WaitGroup
 
 	if tcp {
@@ -154,5 +202,7 @@ func main() {
 		}()
 	}
 
+	go updater()
+
 	wg.Wait()
 }

ndns/record.json

@@ -1,7 +1,7 @@
 [
     {
         "host": "www.baidu.com",
-        "tc": true,
+        "tc": false,
         "ips": [
             "127.0.0.2",
             "127.0.0.3",

util.h

@@ -664,7 +664,7 @@ time_format(char *buf, size_t cap, time_t now)
     return 0;
 }
 
-int
+static int
 read_command_output(const char *command, char *buf, size_t n)
 {
     if (buf == NULL || n == 0) {