From 651e7a75d2120caf7672ce0ff44cbcfaab124d06 Mon Sep 17 00:00:00 2001
From: Mizzick <Mizzick@users.noreply.github.com>
Date: Fri, 23 Dec 2022 15:13:37 +0700
Subject: [PATCH] edns0 padding option (#42)

* edns0 padding option

* edns0 padding option

* edns0 padding option

* imp code

Co-authored-by: Dimitry Kolyshev <dkolyshev@adguard.com>
---
 README.md |  5 +++++
 main.go   | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/README.md b/README.md
index c70e49f..a5a20be 100644
--- a/README.md
+++ b/README.md
@@ -92,6 +92,11 @@ Specify the class of query (default IN):
 CLASS=CH ./dnslookup example.org tls://127.0.0.1
 ```
 
+Add EDNS0 Padding:
+```
+PAD=1 ./dnslookup example.org tls://127.0.0.1
+```
+
 Verbose-level logging:
 ```shell
 VERBOSE=1 ./dnslookup example.org tls://dns.adguard.com
diff --git a/main.go b/main.go
index 77a974b..edbc4c6 100644
--- a/main.go
+++ b/main.go
@@ -28,6 +28,7 @@ func main() {
 	timeoutStr := os.Getenv("TIMEOUT")
 	http3Enabled := os.Getenv("HTTP3") == "1"
 	verbose := os.Getenv("VERBOSE") == "1"
+	padding := os.Getenv("PAD") == "1"
 	class := getClass()
 	rrType := getRRType()
 
@@ -126,6 +127,11 @@ func main() {
 	req.Question = []dns.Question{
 		{Name: domain + ".", Qtype: rrType, Qclass: class},
 	}
+
+	if padding {
+		req.Extra = []dns.RR{newEDNS0Padding(req)}
+	}
+
 	reply, err := u.Exchange(&req)
 	if err != nil {
 		log.Fatalf("Cannot make the DNS request: %s", err)
@@ -186,3 +192,29 @@ func usage() {
 	os.Stdout.WriteString("<providerName>: optional, DNSCrypt provider name\n")
 	os.Stdout.WriteString("<serverPk>: optional, DNSCrypt server public key\n")
 }
+
+// requestPaddingBlockSize is used to pad responses over DoT and DoH according
+// to RFC 8467.
+const requestPaddingBlockSize = 128
+const uDPBufferSize = dns.DefaultMsgSize
+
+// newEDNS0Padding constructs a new OPT RR EDNS0 Padding for the extra section.
+func newEDNS0Padding(req dns.Msg) (extra dns.RR) {
+	msgLen := req.Len()
+	padLen := requestPaddingBlockSize - msgLen%requestPaddingBlockSize
+
+	// Truncate padding to fit in UDP buffer.
+	if msgLen+padLen > uDPBufferSize {
+		padLen = uDPBufferSize - msgLen
+		if padLen < 0 {
+			padLen = 0
+		}
+	}
+
+	return &dns.OPT{
+		Hdr: dns.RR_Header{Name: ".", Rrtype: dns.TypeOPT, Class: uDPBufferSize},
+		Option: []dns.EDNS0{
+			&dns.EDNS0_PADDING{Padding: make([]byte, padLen)},
+		},
+	}
+}