dns cache command

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"))

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
+}

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)
+}