From 63c6dd5617977819b76e47facc20a7bfced00594 Mon Sep 17 00:00:00 2001
From: SeanZH <i@dilfish.dev>
Date: Fri, 30 Apr 2021 18:37:01 +0800
Subject: [PATCH] add timeout and err handle

---
 .idea/workspace.xml | 13 ++-----------
 client.go           |  9 +++++++--
 dns.go              | 10 ++++++++--
 3 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/.idea/workspace.xml b/.idea/workspace.xml
index c90e754..8039a68 100644
--- a/.idea/workspace.xml
+++ b/.idea/workspace.xml
@@ -3,17 +3,8 @@
   <component name="ChangeListManager">
     <list default="true" id="824504b2-e674-493c-83af-08a5427caade" name="Default Changelist" comment="">
       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/a.go" beforeDir="false" afterPath="$PROJECT_DIR$/a.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/api.go" beforeDir="false" afterPath="$PROJECT_DIR$/api.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/api_add.go" beforeDir="false" afterPath="$PROJECT_DIR$/api_add.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/api_del.go" beforeDir="false" afterPath="$PROJECT_DIR$/api_del.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/api_list.go" beforeDir="false" afterPath="$PROJECT_DIR$/api_list.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/cname.go" beforeDir="false" afterPath="$PROJECT_DIR$/cname.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/ns.go" beforeDir="false" afterPath="$PROJECT_DIR$/ns.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/record.go" beforeDir="false" afterPath="$PROJECT_DIR$/record.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/svcb.go" beforeDir="false" afterPath="$PROJECT_DIR$/svcb.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/txt.go" beforeDir="false" afterPath="$PROJECT_DIR$/txt.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/type.go" beforeDir="false" afterPath="$PROJECT_DIR$/type.go" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/client.go" beforeDir="false" afterPath="$PROJECT_DIR$/client.go" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/dns.go" beforeDir="false" afterPath="$PROJECT_DIR$/dns.go" afterDir="false" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
diff --git a/client.go b/client.go
index 38696f6..1640b69 100644
--- a/client.go
+++ b/client.go
@@ -21,12 +21,17 @@ type MongoClientConfig struct {
 }
 
 func NewMongoClient(conf *MongoClientConfig) *MongoClient {
-	client, err := mongo.NewClient(options.Client().ApplyURI(conf.Addr))
+	client, err := mongo.NewClient(options.Client().
+		ApplyURI(conf.Addr).
+		SetConnectTimeout(time.Second * 2).
+		SetHeartbeatInterval(time.Second * 10).
+		SetSocketTimeout(time.Second * 2).
+		SetServerSelectionTimeout(time.Second *2))
 	if err != nil {
 		log.Println("new client error:", conf.Addr, err)
 		return nil
 	}
-	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
 	defer cancel()
 	err = client.Connect(ctx)
 	if err != nil {
diff --git a/dns.go b/dns.go
index 3a29994..1361470 100644
--- a/dns.go
+++ b/dns.go
@@ -19,25 +19,31 @@ func ParseReqInfo(r *dns.Msg) (err error) {
 }
 
 type Handler struct {
-	M *RecordManager
+	M           *RecordManager
+	ServFailMsg *dns.Msg
 }
 
 func NewHandler(conf *MongoClientConfig) *Handler {
 	m := NewRecordManager(conf)
-	return &Handler{M: m}
+	return &Handler{M: m, ServFailMsg: new(dns.Msg)}
 }
 
 func (h *Handler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 	log.Print("Client:", w.RemoteAddr())
 	log.Println("au:", r.Authoritative, "tr:", r.Truncated, "rd:", r.RecursionDesired, "ra:", r.RecursionAvailable, "ad:", r.AuthenticatedData, "cd:", r.CheckingDisabled)
+	h.ServFailMsg.SetReply(r)
 	err := ParseReqInfo(r)
 	if err != nil {
+		h.ServFailMsg.MsgHdr.Rcode = dns.RcodeBadName
 		log.Println("bad dns request:", err)
+		w.WriteMsg(h.ServFailMsg)
 		return
 	}
 	msg, err := h.GetRecord(r)
 	if err != nil {
+		h.ServFailMsg.MsgHdr.Rcode = dns.RcodeServerFailure
 		log.Println("get record error:", err)
+		w.WriteMsg(h.ServFailMsg)
 		return
 	}
 	w.WriteMsg(msg)