From fdc1b81e4224dbed8fd7372f79679de98504ecec Mon Sep 17 00:00:00 2001
From: Star Brilliant <coder@poorlab.com>
Date: Wed, 25 Jan 2023 06:47:13 +0000
Subject: [PATCH] Properly truncate DNS packets

This should fix issue #144.
---
 doh-client/google.go | 14 +++++---------
 doh-client/ietf.go   | 14 +++++---------
 2 files changed, 10 insertions(+), 18 deletions(-)

diff --git a/doh-client/google.go b/doh-client/google.go
index e7a8db9..f298ca3 100644
--- a/doh-client/google.go
+++ b/doh-client/google.go
@@ -163,6 +163,11 @@ func (c *Client) parseResponseGoogle(ctx context.Context, w dns.ResponseWriter,
 	fixEmptyNames(&respJSON)
 
 	fullReply := jsondns.Unmarshal(req.reply, &respJSON, req.udpSize, req.ednsClientNetmask)
+	if isTCP {
+		fullReply.Truncate(dns.MaxMsgSize)
+	} else {
+		fullReply.Truncate(int(req.udpSize))
+	}
 	buf, err := fullReply.Pack()
 	if err != nil {
 		log.Println(err)
@@ -170,15 +175,6 @@ func (c *Client) parseResponseGoogle(ctx context.Context, w dns.ResponseWriter,
 		w.WriteMsg(req.reply)
 		return
 	}
-	if !isTCP && len(buf) > int(req.udpSize) {
-		fullReply.Truncated = true
-		buf, err = fullReply.Pack()
-		if err != nil {
-			log.Println(err)
-			return
-		}
-		buf = buf[:req.udpSize]
-	}
 	w.Write(buf)
 }
 
diff --git a/doh-client/ietf.go b/doh-client/ietf.go
index 30baf50..bf9545c 100644
--- a/doh-client/ietf.go
+++ b/doh-client/ietf.go
@@ -231,6 +231,11 @@ func (c *Client) parseResponseIETF(ctx context.Context, w dns.ResponseWriter, r
 		_ = fixRecordTTL(rr, timeDelta)
 	}
 
+	if isTCP {
+		fullReply.Truncate(dns.MaxMsgSize)
+	} else {
+		fullReply.Truncate(int(req.udpSize))
+	}
 	buf, err := fullReply.Pack()
 	if err != nil {
 		log.Printf("packing error with upstream %s: %v\n", req.currentUpstream, err)
@@ -238,15 +243,6 @@ func (c *Client) parseResponseIETF(ctx context.Context, w dns.ResponseWriter, r
 		w.WriteMsg(req.reply)
 		return
 	}
-	if !isTCP && len(buf) > int(req.udpSize) {
-		fullReply.Truncated = true
-		buf, err = fullReply.Pack()
-		if err != nil {
-			log.Printf("re-packing error with upstream %s: %v\n", req.currentUpstream, err)
-			return
-		}
-		buf = buf[:req.udpSize]
-	}
 	_, err = w.Write(buf)
 	if err != nil {
 		log.Printf("failed to write to client: %v\n", err)