From 1c830472b997ce03772bbc328b23379d35b87bda Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Sat, 29 Dec 2018 09:03:32 +0100
Subject: [PATCH] dns protocol package

---
 app/dns/udpns.go          | 16 +-------
 common/protocol/dns/io.go | 77 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 14 deletions(-)
 create mode 100644 common/protocol/dns/io.go

diff --git a/app/dns/udpns.go b/app/dns/udpns.go
index 3a246886..2b2ba0f7 100644
--- a/app/dns/udpns.go
+++ b/app/dns/udpns.go
@@ -8,10 +8,10 @@ import (
 	"time"
 
 	"golang.org/x/net/dns/dnsmessage"
-
 	"v2ray.com/core/common"
 	"v2ray.com/core/common/buf"
 	"v2ray.com/core/common/net"
+	"v2ray.com/core/common/protocol/dns"
 	"v2ray.com/core/common/session"
 	"v2ray.com/core/common/signal/pubsub"
 	"v2ray.com/core/common/task"
@@ -293,25 +293,13 @@ func (s *ClassicNameServer) buildMsgs(domain string, option IPOption) []*dnsmess
 	return msgs
 }
 
-func msgToBuffer2(msg *dnsmessage.Message) (*buf.Buffer, error) {
-	buffer := buf.New()
-	rawBytes := buffer.Extend(buf.Size)
-	packed, err := msg.AppendPack(rawBytes[:0])
-	if err != nil {
-		buffer.Release()
-		return nil, err
-	}
-	buffer.Resize(0, int32(len(packed)))
-	return buffer, nil
-}
-
 func (s *ClassicNameServer) sendQuery(ctx context.Context, domain string, option IPOption) {
 	newError("querying DNS for: ", domain).AtDebug().WriteToLog(session.ExportIDToError(ctx))
 
 	msgs := s.buildMsgs(domain, option)
 
 	for _, msg := range msgs {
-		b, err := msgToBuffer2(msg)
+		b, err := dns.PackMessage(msg)
 		common.Must(err)
 		s.udpServer.Dispatch(context.Background(), s.address, b)
 	}
diff --git a/common/protocol/dns/io.go b/common/protocol/dns/io.go
new file mode 100644
index 00000000..6b236112
--- /dev/null
+++ b/common/protocol/dns/io.go
@@ -0,0 +1,77 @@
+package dns
+
+import (
+	"sync"
+
+	"golang.org/x/net/dns/dnsmessage"
+	"v2ray.com/core/common"
+	"v2ray.com/core/common/buf"
+)
+
+func PackMessage(msg *dnsmessage.Message) (*buf.Buffer, error) {
+	buffer := buf.New()
+	rawBytes := buffer.Extend(buf.Size)
+	packed, err := msg.AppendPack(rawBytes[:0])
+	if err != nil {
+		buffer.Release()
+		return nil, err
+	}
+	buffer.Resize(0, int32(len(packed)))
+	return buffer, nil
+}
+
+type MessageReader interface {
+	ReadMessage() (*buf.Buffer, error)
+}
+
+type UDPReader struct {
+	buf.Reader
+
+	access sync.Mutex
+	cache  buf.MultiBuffer
+}
+
+func (r *UDPReader) readCache() *buf.Buffer {
+	r.access.Lock()
+	defer r.access.Unlock()
+
+	mb, b := buf.SplitFirst(r.cache)
+	r.cache = mb
+	return b
+}
+
+func (r *UDPReader) refill() error {
+	mb, err := r.Reader.ReadMultiBuffer()
+	if err != nil {
+		return err
+	}
+	r.access.Lock()
+	r.cache = mb
+	r.access.Unlock()
+	return nil
+}
+
+// ReadMessage implements MessageReader.
+func (r *UDPReader) ReadMessage() (*buf.Buffer, error) {
+	for {
+		b := r.readCache()
+		if b != nil {
+			return b, nil
+		}
+		if err := r.refill(); err != nil {
+			return nil, err
+		}
+	}
+}
+
+// Close implements common.Closable.
+func (r *UDPReader) Close() error {
+	defer func() {
+		r.access.Lock()
+		buf.ReleaseMulti(r.cache)
+		r.cache = nil
+		r.access.Unlock()
+	}()
+
+	return common.Close(r.Reader)
+}