From 8d6fdd014a120d9131b479c7cdd23efa1fa5cb67 Mon Sep 17 00:00:00 2001
From: v2ray <admin@v2ray.com>
Date: Fri, 25 Dec 2015 01:07:42 +0100
Subject: [PATCH] dns cache command

---
 common/net/testing/assert/address.go | 12 +++++
 proxy/vmess/command/dns.go           | 69 ++++++++++++++++++++++++++++
 proxy/vmess/command/dns_test.go      | 32 +++++++++++++
 3 files changed, 113 insertions(+)
 create mode 100644 proxy/vmess/command/dns.go
 create mode 100644 proxy/vmess/command/dns_test.go

diff --git a/common/net/testing/assert/address.go b/common/net/testing/assert/address.go
index 2c82f205f..4783c152d 100644
--- a/common/net/testing/assert/address.go
+++ b/common/net/testing/assert/address.go
@@ -24,6 +24,18 @@ func (subject *AddressSubject) DisplayString() string {
 	return subject.Subject.DisplayString(subject.value.String())
 }
 
+func (subject *AddressSubject) Equals(another v2net.Address) {
+	if subject.value.IsIPv4() && another.IsIPv4() {
+		IP(subject.value.IP()).Equals(another.IP())
+	} else if subject.value.IsIPv6() && another.IsIPv6() {
+		IP(subject.value.IP()).Equals(another.IP())
+	} else if subject.value.IsDomain() && another.IsDomain() {
+		assert.StringLiteral(subject.value.Domain()).Equals(another.Domain())
+	} else {
+		subject.Fail(subject.DisplayString(), "equals to", another)
+	}
+}
+
 func (subject *AddressSubject) IsIPv4() {
 	if !subject.value.IsIPv4() {
 		subject.Fail(subject.DisplayString(), "is", serial.StringLiteral("an IPv4 address"))
diff --git a/proxy/vmess/command/dns.go b/proxy/vmess/command/dns.go
new file mode 100644
index 000000000..4e4a5e5d3
--- /dev/null
+++ b/proxy/vmess/command/dns.go
@@ -0,0 +1,69 @@
+package command
+
+import (
+	"errors"
+	"io"
+
+	v2net "github.com/v2ray/v2ray-core/common/net"
+	"github.com/v2ray/v2ray-core/transport"
+)
+
+func init() {
+	RegisterResponseCommand(2, func() Command { return new(CacheDns) })
+}
+
+const (
+	typeIPv4 byte = 1
+	typeIPv6 byte = 2
+)
+
+var (
+	ErrDomainAddress = errors.New("Unexpected domain address")
+)
+
+// Size: 1 byte type + 4 or 16 byte IP addr
+type CacheDns struct {
+	Address v2net.Address
+}
+
+func (this *CacheDns) Marshal(writer io.Writer) (int, error) {
+	if this.Address.IsIPv4() {
+		writer.Write([]byte{typeIPv4})
+		writer.Write(this.Address.IP())
+		return 5, nil
+	}
+
+	if this.Address.IsIPv6() {
+		writer.Write([]byte{typeIPv6})
+		writer.Write(this.Address.IP())
+		return 17, nil
+	}
+
+	return 0, ErrDomainAddress
+}
+
+func (this *CacheDns) Unmarshal(data []byte) error {
+	if len(data) == 0 {
+		return transport.CorruptedPacket
+	}
+	typeIP := data[0]
+	data = data[1:]
+
+	if typeIP == typeIPv4 {
+		if len(data) < 4 {
+			return transport.CorruptedPacket
+		}
+		this.Address = v2net.IPAddress(data[0:4])
+		return nil
+	}
+
+	if typeIP == typeIPv6 {
+		if len(data) < 16 {
+			return transport.CorruptedPacket
+		}
+		this.Address = v2net.IPAddress(data[0:16])
+		return nil
+	}
+
+	return transport.CorruptedPacket
+}
diff --git a/proxy/vmess/command/dns_test.go b/proxy/vmess/command/dns_test.go
new file mode 100644
index 000000000..ee1390adb
--- /dev/null
+++ b/proxy/vmess/command/dns_test.go
@@ -0,0 +1,32 @@
+package command_test
+
+import (
+	"testing"
+
+	"github.com/v2ray/v2ray-core/common/alloc"
+	v2net "github.com/v2ray/v2ray-core/common/net"
+	netassert "github.com/v2ray/v2ray-core/common/net/testing/assert"
+	. "github.com/v2ray/v2ray-core/proxy/vmess/command"
+	v2testing "github.com/v2ray/v2ray-core/testing"
+	"github.com/v2ray/v2ray-core/testing/assert"
+)
+
+func TestCacheDnsIPv4(t *testing.T) {
+	v2testing.Current(t)
+
+	cd := &CacheDns{
+		Address: v2net.IPAddress([]byte{1, 2, 3, 4}),
+	}
+
+	buffer := alloc.NewBuffer().Clear()
+	defer buffer.Release()
+
+	nBytes, err := cd.Marshal(buffer)
+	assert.Error(err).IsNil()
+	assert.Int(nBytes).Equals(buffer.Len())
+
+	cd2 := &CacheDns{}
+	err = cd2.Unmarshal(buffer.Value)
+	assert.Error(err).IsNil()
+	netassert.Address(cd.Address).Equals(cd2.Address)
+}