zhao/luci-app-mosdns
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
3fec661adb
luci-app-mosdns/mosdns/patches/117-pool-simplify-PackBuffer.patch
sbwml 3fec661adb mosdns: pacthes: sync upstream latest commit 
fixed #129
2023-09-23 15:17:07 +08:00

223 lines
6.9 KiB
Diff
Raw Blame History

From 64a83b8e28b3988df9eec4425130b57a09b15032 Mon Sep 17 00:00:00 2001
From: Irine Sistiana <49315432+IrineSistiana@users.noreply.github.com>
Date: Thu, 21 Sep 2023 22:06:49 +0800
Subject: [PATCH 1/9] pool: simplify PackBuffer
---
pkg/dnsutils/net_io.go | 7 +++---
pkg/pool/msg_buf.go | 41 ++++++++++++++++++------------------
pkg/pool/msg_buf_test.go | 6 +-----
pkg/server/http_handler.go | 6 +++---
pkg/server/tcp.go | 6 +++---
pkg/server/udp.go | 6 +++---
pkg/upstream/doh/upstream.go | 5 +++--
7 files changed, 37 insertions(+), 40 deletions(-)
diff --git a/pkg/dnsutils/net_io.go b/pkg/dnsutils/net_io.go
index f165446..26e6efb 100644
--- a/pkg/dnsutils/net_io.go
+++ b/pkg/dnsutils/net_io.go
@@ -101,13 +101,12 @@ func WriteRawMsgToTCP(c io.Writer, b []byte) (n int, err error) {
}
func WriteMsgToUDP(c io.Writer, m *dns.Msg) (int, error) {
- b, buf, err := pool.PackBuffer(m)
+ b, err := pool.PackBuffer(m)
if err != nil {
return 0, err
}
- defer pool.ReleaseBuf(buf)
-
- return c.Write(b)
+ defer pool.ReleaseBuf(b)
+ return c.Write(*b)
}
func ReadMsgFromUDP(c io.Reader, bufSize int) (*dns.Msg, int, error) {
diff --git a/pkg/pool/msg_buf.go b/pkg/pool/msg_buf.go
index 11faf7d..b5f861c 100644
--- a/pkg/pool/msg_buf.go
+++ b/pkg/pool/msg_buf.go
@@ -26,47 +26,48 @@ import (
"github.com/miekg/dns"
)
-// There is no such way to give dns.Msg.PackBuffer() a buffer
-// with a proper size.
-// Just give it a big buf and hope the buf will be reused in most scenes.
-const packBufSize = 4096
+// dns.Msg.PackBuffer requires a buffer with length of m.Len() + 1.
+// Don't know why it needs one more byte.
+func getPackBuffer(m *dns.Msg) int {
+ return m.Len() + 1
+}
// PackBuffer packs the dns msg m to wire format.
// Callers should release the buf by calling ReleaseBuf after they have done
// with the wire []byte.
-func PackBuffer(m *dns.Msg) (wire []byte, buf *[]byte, err error) {
- buf = GetBuf(packBufSize)
- wire, err = m.PackBuffer(*buf)
+func PackBuffer(m *dns.Msg) (*[]byte, error) {
+ b := GetBuf(getPackBuffer(m))
+ wire, err := m.PackBuffer(*b)
if err != nil {
- ReleaseBuf(buf)
- return nil, nil, err
+ ReleaseBuf(b)
+ return nil, err
}
- return wire, buf, nil
+ if &((*b)[0]) != &wire[0] { // reallocated
+ ReleaseBuf(b)
+ return nil, dns.ErrBuf
+ }
+ return b, nil
}
// PackBuffer packs the dns msg m to wire format, with to bytes length header.
// Callers should release the buf by calling ReleaseBuf.
-func PackTCPBuffer(m *dns.Msg) (buf *[]byte, err error) {
- b := GetBuf(packBufSize)
+func PackTCPBuffer(m *dns.Msg) (*[]byte, error) {
+ b := GetBuf(2 + getPackBuffer(m))
wire, err := m.PackBuffer((*b)[2:])
if err != nil {
ReleaseBuf(b)
return nil, err
}
+ if &((*b)[2]) != &wire[0] { // reallocated
+ ReleaseBuf(b)
+ return nil, dns.ErrBuf
+ }
l := len(wire)
if l > dns.MaxMsgSize {
ReleaseBuf(b)
return nil, fmt.Errorf("dns payload size %d is too large", l)
}
-
- if &((*b)[2]) != &wire[0] { // reallocated
- ReleaseBuf(b)
- b = GetBuf(l + 2)
- binary.BigEndian.PutUint16((*b)[:2], uint16(l))
- copy((*b)[2:], wire)
- return b, nil
- }
binary.BigEndian.PutUint16((*b)[:2], uint16(l))
*b = (*b)[:2+l]
return b, nil
diff --git a/pkg/pool/msg_buf_test.go b/pkg/pool/msg_buf_test.go
index 97d9a76..bfd98d1 100644
--- a/pkg/pool/msg_buf_test.go
+++ b/pkg/pool/msg_buf_test.go
@@ -28,12 +28,8 @@ import (
func TestPackBuffer_No_Allocation(t *testing.T) {
m := new(dns.Msg)
m.SetQuestion("123.", dns.TypeAAAA)
- wire, buf, err := PackBuffer(m)
+ _, err := PackBuffer(m)
if err != nil {
t.Fatal(err)
}
-
- if cap(wire) != cap(*buf) {
- t.Fatalf("wire and buf have different cap, wire %d, buf %d", cap(wire), cap(*buf))
- }
}
diff --git a/pkg/server/http_handler.go b/pkg/server/http_handler.go
index 58f5811..3e671e3 100644
--- a/pkg/server/http_handler.go
+++ b/pkg/server/http_handler.go
@@ -103,17 +103,17 @@ func (h *HttpHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
panic(err) // Force http server to close connection.
}
- b, buf, err := pool.PackBuffer(r)
+ b, err := pool.PackBuffer(r)
if err != nil {
w.WriteHeader(http.StatusInternalServerError)
h.warnErr(req, "failed to unpack handler's response", err)
return
}
- defer pool.ReleaseBuf(buf)
+ defer pool.ReleaseBuf(b)
w.Header().Set("Content-Type", "application/dns-message")
w.Header().Set("Cache-Control", fmt.Sprintf("max-age=%d", dnsutils.GetMinimalTTL(r)))
- if _, err := w.Write(b); err != nil {
+ if _, err := w.Write(*b); err != nil {
h.warnErr(req, "failed to write response", err)
return
}
diff --git a/pkg/server/tcp.go b/pkg/server/tcp.go
index 5f479b1..ddc4846 100644
--- a/pkg/server/tcp.go
+++ b/pkg/server/tcp.go
@@ -101,14 +101,14 @@ func ServeTCP(l net.Listener, h Handler, opts TCPServerOpts) error {
c.Close() // abort the connection
return
}
- b, buf, err := pool.PackBuffer(r)
+ b, err := pool.PackTCPBuffer(r)
if err != nil {
logger.Error("failed to unpack handler's response", zap.Error(err), zap.Stringer("msg", r))
return
}
- defer pool.ReleaseBuf(buf)
+ defer pool.ReleaseBuf(b)
- if _, err := dnsutils.WriteRawMsgToTCP(c, b); err != nil {
+ if _, err := c.Write(*b); err != nil {
logger.Warn("failed to write response", zap.Stringer("client", c.RemoteAddr()), zap.Error(err))
return
}
diff --git a/pkg/server/udp.go b/pkg/server/udp.go
index 4dc1087..22e8d2b 100644
--- a/pkg/server/udp.go
+++ b/pkg/server/udp.go
@@ -95,18 +95,18 @@ func ServeUDP(c *net.UDPConn, h Handler, opts UDPServerOpts) error {
}
if r != nil {
r.Truncate(getUDPSize(q))
- b, buf, err := pool.PackBuffer(r)
+ b, err := pool.PackBuffer(r)
if err != nil {
logger.Error("failed to unpack handler's response", zap.Error(err), zap.Stringer("msg", r))
return
}
- defer pool.ReleaseBuf(buf)
+ defer pool.ReleaseBuf(b)
var oob []byte
if oobWriter != nil && dstIpFromCm != nil {
oob = oobWriter(dstIpFromCm)
}
- if _, _, err := c.WriteMsgUDPAddrPort(b, oob, remoteAddr); err != nil {
+ if _, _, err := c.WriteMsgUDPAddrPort(*b, oob, remoteAddr); err != nil {
logger.Warn("failed to write response", zap.Stringer("client", remoteAddr), zap.Error(err))
}
}
diff --git a/pkg/upstream/doh/upstream.go b/pkg/upstream/doh/upstream.go
index abc124b..9cc72c4 100644
--- a/pkg/upstream/doh/upstream.go
+++ b/pkg/upstream/doh/upstream.go
@@ -54,11 +54,12 @@ var (
)
func (u *Upstream) ExchangeContext(ctx context.Context, q *dns.Msg) (*dns.Msg, error) {
- wire, buf, err := pool.PackBuffer(q)
+ bp, err := pool.PackBuffer(q)
if err != nil {
return nil, fmt.Errorf("failed to pack query msg, %w", err)
}
- defer pool.ReleaseBuf(buf)
+ defer pool.ReleaseBuf(bp)
+ wire := *bp
// In order to maximize HTTP cache friendliness, DoH clients using media
// formats that include the ID field from the DNS message header, such
--
2.34.8
Reference in New Issue View Git Blame Copy Permalink