From b319704282e0e45b47e242543dd5a6d33cb4a7d4 Mon Sep 17 00:00:00 2001
From: V2Ray <admin@v2ray.com>
Date: Sun, 20 Sep 2015 16:03:12 +0200
Subject: [PATCH] Add network type to destination to be ready for UDP
 connection

---
 common/net/destination.go        | 46 ++++++++++++++++++++++++++++++++
 point.go                         |  4 +--
 proxy/freedom/freedom.go         | 10 +++----
 proxy/freedom/freedomfactory.go  |  2 +-
 proxy/socks/protocol/socks.go    | 10 ++++---
 proxy/socks/socks.go             |  4 +--
 proxy/socks/socks_test.go        |  4 +--
 proxy/vmess/config.go            |  9 +++++--
 proxy/vmess/protocol/vmess.go    |  7 +++++
 proxy/vmess/vmess_test.go        |  2 +-
 proxy/vmess/vmessin.go           |  2 +-
 proxy/vmess/vmessout.go          | 22 +++++++--------
 testing/mocks/inboundhandler.go  |  2 +-
 testing/mocks/outboundhandler.go |  4 +--
 14 files changed, 94 insertions(+), 34 deletions(-)
 create mode 100644 common/net/destination.go

diff --git a/common/net/destination.go b/common/net/destination.go
new file mode 100644
index 000000000..bf88f753b
--- /dev/null
+++ b/common/net/destination.go
@@ -0,0 +1,46 @@
+package net
+
+import (
+	"github.com/v2ray/v2ray-core/common/log"
+)
+
+const (
+	NetTCP = byte(0x01)
+	NetUDP = byte(0x02)
+)
+
+type Destination struct {
+	network byte
+	address Address
+}
+
+func NewDestination(network byte, address Address) *Destination {
+	return &Destination{
+		network: network,
+		address: address,
+	}
+}
+
+func (dest *Destination) Network() string {
+	switch dest.network {
+	case NetTCP:
+		return "tcp"
+	case NetUDP:
+		return "udp"
+	default:
+		log.Warning("Unknown network %d", dest.network)
+		return "tcp"
+	}
+}
+
+func (dest *Destination) NetworkByte() byte {
+	return dest.network
+}
+
+func (dest *Destination) Address() Address {
+	return dest.address
+}
+
+func (dest *Destination) String() string {
+	return dest.address.String() + " (" + dest.Network() + ")"
+}
diff --git a/point.go b/point.go
index c03376023..2b0765b18 100644
--- a/point.go
+++ b/point.go
@@ -64,7 +64,7 @@ type InboundConnectionHandler interface {
 }
 
 type OutboundConnectionHandlerFactory interface {
-	Create(VP *Point, config []byte, dest v2net.Address) (OutboundConnectionHandler, error)
+	Create(VP *Point, config []byte, dest *v2net.Destination) (OutboundConnectionHandler, error)
 }
 
 type OutboundConnectionHandler interface {
@@ -85,7 +85,7 @@ func (vp *Point) Start() error {
 	return nil
 }
 
-func (vp *Point) NewInboundConnectionAccepted(destination v2net.Address) InboundRay {
+func (vp *Point) NewInboundConnectionAccepted(destination *v2net.Destination) InboundRay {
 	ray := NewRay()
 	// TODO: handle error
 	och, _ := vp.ochFactory.Create(vp, vp.ochConfig, destination)
diff --git a/proxy/freedom/freedom.go b/proxy/freedom/freedom.go
index 8b18844b4..aa2f95ed0 100644
--- a/proxy/freedom/freedom.go
+++ b/proxy/freedom/freedom.go
@@ -9,10 +9,10 @@ import (
 )
 
 type FreedomConnection struct {
-	dest v2net.Address
+	dest *v2net.Destination
 }
 
-func NewFreedomConnection(dest v2net.Address) *FreedomConnection {
+func NewFreedomConnection(dest *v2net.Destination) *FreedomConnection {
 	return &FreedomConnection{
 		dest: dest,
 	}
@@ -21,12 +21,12 @@ func NewFreedomConnection(dest v2net.Address) *FreedomConnection {
 func (vconn *FreedomConnection) Start(ray core.OutboundRay) error {
 	input := ray.OutboundInput()
 	output := ray.OutboundOutput()
-	conn, err := net.Dial("tcp", vconn.dest.String())
+	conn, err := net.Dial(vconn.dest.Network(), vconn.dest.Address().String())
+	log.Info("Freedom: Opening connection to %s", vconn.dest.String())
 	if err != nil {
 		close(output)
-		return log.Error("Freedom: Failed to open tcp connection: %s : %v", vconn.dest.String(), err)
+		return log.Error("Freedom: Failed to open connection: %s : %v", vconn.dest.String(), err)
 	}
-	log.Info("Freedom: Sending outbound tcp: %s", vconn.dest.String())
 
 	readFinish := make(chan bool)
 	writeFinish := make(chan bool)
diff --git a/proxy/freedom/freedomfactory.go b/proxy/freedom/freedomfactory.go
index a40bb3bbd..31d15d8aa 100644
--- a/proxy/freedom/freedomfactory.go
+++ b/proxy/freedom/freedomfactory.go
@@ -8,7 +8,7 @@ import (
 type FreedomFactory struct {
 }
 
-func (factory FreedomFactory) Create(vp *core.Point, config []byte, dest v2net.Address) (core.OutboundConnectionHandler, error) {
+func (factory FreedomFactory) Create(vp *core.Point, config []byte, dest *v2net.Destination) (core.OutboundConnectionHandler, error) {
 	return NewFreedomConnection(dest), nil
 }
 
diff --git a/proxy/socks/protocol/socks.go b/proxy/socks/protocol/socks.go
index 649346464..e4c71e945 100644
--- a/proxy/socks/protocol/socks.go
+++ b/proxy/socks/protocol/socks.go
@@ -278,17 +278,19 @@ func ReadRequest(reader io.Reader) (request *Socks5Request, err error) {
 	return
 }
 
-func (request *Socks5Request) Destination() v2net.Address {
+func (request *Socks5Request) Destination() *v2net.Destination {
+	var address v2net.Address
 	switch request.AddrType {
 	case AddrTypeIPv4:
-		return v2net.IPAddress(request.IPv4[:], request.Port)
+		address = v2net.IPAddress(request.IPv4[:], request.Port)
 	case AddrTypeIPv6:
-		return v2net.IPAddress(request.IPv6[:], request.Port)
+		address = v2net.IPAddress(request.IPv6[:], request.Port)
 	case AddrTypeDomain:
-		return v2net.DomainAddress(request.Domain, request.Port)
+		address = v2net.DomainAddress(request.Domain, request.Port)
 	default:
 		panic("Unknown address type")
 	}
+	return v2net.NewDestination(v2net.NetTCP, address)
 }
 
 const (
diff --git a/proxy/socks/socks.go b/proxy/socks/socks.go
index 62954a544..a21736af1 100644
--- a/proxy/socks/socks.go
+++ b/proxy/socks/socks.go
@@ -70,7 +70,7 @@ func (server *SocksServer) HandleConnection(connection net.Conn) error {
 		return err
 	}
 
-	var dest v2net.Address
+	var dest *v2net.Destination
 
 	// TODO refactor this part
 	if err == protocol.ErrorSocksVersion4 {
@@ -85,7 +85,7 @@ func (server *SocksServer) HandleConnection(connection net.Conn) error {
 			return ErrorCommandNotSupported
 		}
 
-		dest = v2net.IPAddress(auth4.IP[:], auth4.Port)
+		dest = v2net.NewDestination(v2net.NetTCP, v2net.IPAddress(auth4.IP[:], auth4.Port))
 	} else {
 		expectedAuthMethod := protocol.AuthNotRequired
 		if server.config.AuthMethod == JsonAuthMethodUserPass {
diff --git a/proxy/socks/socks_test.go b/proxy/socks/socks_test.go
index 626489bb2..7adf9d578 100644
--- a/proxy/socks/socks_test.go
+++ b/proxy/socks/socks_test.go
@@ -61,7 +61,7 @@ func TestSocksTcpConnect(t *testing.T) {
 
 	assert.Bytes([]byte(data2Send)).Equals(och.Data2Send.Bytes())
 	assert.Bytes(dataReturned).Equals(och.Data2Return)
-	assert.String(targetServer).Equals(och.Destination.String())
+	assert.String(targetServer).Equals(och.Destination.Address().String())
 }
 
 func TestSocksTcpConnectWithUserPass(t *testing.T) {
@@ -112,5 +112,5 @@ func TestSocksTcpConnectWithUserPass(t *testing.T) {
 
 	assert.Bytes([]byte(data2Send)).Equals(och.Data2Send.Bytes())
 	assert.Bytes(dataReturned).Equals(och.Data2Return)
-	assert.String(targetServer).Equals(och.Destination.String())
+	assert.String(targetServer).Equals(och.Destination.Address().String())
 }
diff --git a/proxy/vmess/config.go b/proxy/vmess/config.go
index 1eac61f21..43e7f0863 100644
--- a/proxy/vmess/config.go
+++ b/proxy/vmess/config.go
@@ -35,6 +35,7 @@ type VNextConfig struct {
 	Address string      `json:"address"`
 	Port    uint16      `json:"port"`
 	Users   []VMessUser `json:"users"`
+	Network string      `json:"network"`
 }
 
 func (config VNextConfig) ToVNextServer() VNextServer {
@@ -50,9 +51,13 @@ func (config VNextConfig) ToVNextServer() VNextServer {
 	if ip == nil {
 		panic(log.Error("Unable to parse VNext IP: %s", config.Address))
 	}
+	network := v2net.NetTCP
+	if config.Network == "udp" {
+		network = v2net.NetUDP
+	}
 	return VNextServer{
-		Address: v2net.IPAddress(ip, config.Port),
-		Users:   users,
+		Destination: v2net.NewDestination(network, v2net.IPAddress(ip, config.Port)),
+		Users:       users,
 	}
 }
 
diff --git a/proxy/vmess/protocol/vmess.go b/proxy/vmess/protocol/vmess.go
index 282db56e8..570c6b693 100644
--- a/proxy/vmess/protocol/vmess.go
+++ b/proxy/vmess/protocol/vmess.go
@@ -23,6 +23,9 @@ const (
 	addrTypeIPv6   = byte(0x03)
 	addrTypeDomain = byte(0x02)
 
+	CmdTCP = byte(0x01)
+	CmdUDP = byte(0x02)
+
 	Version = byte(0x01)
 
 	blockSize = 16
@@ -47,6 +50,10 @@ type VMessRequest struct {
 	Address        v2net.Address
 }
 
+func (request *VMessRequest) Destination() *v2net.Destination {
+	return v2net.NewDestination(request.Command, request.Address)
+}
+
 type VMessRequestReader struct {
 	vUserSet user.UserSet
 }
diff --git a/proxy/vmess/vmess_test.go b/proxy/vmess/vmess_test.go
index f68fdb075..2b577bd27 100644
--- a/proxy/vmess/vmess_test.go
+++ b/proxy/vmess/vmess_test.go
@@ -68,7 +68,7 @@ func TestVMessInAndOut(t *testing.T) {
 	err = pointB.Start()
 	assert.Error(err).IsNil()
 
-	dest := v2net.IPAddress([]byte{1, 2, 3, 4}, 80)
+	dest := v2net.NewDestination(v2net.NetTCP, v2net.IPAddress([]byte{1, 2, 3, 4}, 80))
 	ich.Communicate(dest)
 	assert.Bytes([]byte(data2Send)).Equals(och.Data2Send.Bytes())
 	assert.Bytes(ich.DataReturned.Bytes()).Equals(och.Data2Return)
diff --git a/proxy/vmess/vmessin.go b/proxy/vmess/vmessin.go
index 338809fce..ac142d999 100644
--- a/proxy/vmess/vmessin.go
+++ b/proxy/vmess/vmessin.go
@@ -61,7 +61,7 @@ func (handler *VMessInboundHandler) HandleConnection(connection net.Conn) error
 	}
 	log.Debug("VMessIn: Received request for %s", request.Address.String())
 
-	ray := handler.vPoint.NewInboundConnectionAccepted(request.Address)
+	ray := handler.vPoint.NewInboundConnectionAccepted(request.Destination())
 	input := ray.InboundInput()
 	output := ray.InboundOutput()
 
diff --git a/proxy/vmess/vmessout.go b/proxy/vmess/vmessout.go
index eb7ecce9f..5ed8ad54c 100644
--- a/proxy/vmess/vmessout.go
+++ b/proxy/vmess/vmessout.go
@@ -17,17 +17,17 @@ import (
 
 // VNext is the next Point server in the connection chain.
 type VNextServer struct {
-	Address v2net.Address // Address of VNext server
-	Users   []user.User   // User accounts for accessing VNext.
+	Destination *v2net.Destination // Address of VNext server
+	Users       []user.User        // User accounts for accessing VNext.
 }
 
 type VMessOutboundHandler struct {
 	vPoint    *core.Point
-	dest      v2net.Address
+	dest      *v2net.Destination
 	vNextList []VNextServer
 }
 
-func NewVMessOutboundHandler(vp *core.Point, vNextList []VNextServer, dest v2net.Address) *VMessOutboundHandler {
+func NewVMessOutboundHandler(vp *core.Point, vNextList []VNextServer, dest *v2net.Destination) *VMessOutboundHandler {
 	return &VMessOutboundHandler{
 		vPoint:    vp,
 		dest:      dest,
@@ -35,7 +35,7 @@ func NewVMessOutboundHandler(vp *core.Point, vNextList []VNextServer, dest v2net
 	}
 }
 
-func (handler *VMessOutboundHandler) pickVNext() (v2net.Address, user.User) {
+func (handler *VMessOutboundHandler) pickVNext() (*v2net.Destination, user.User) {
 	vNextLen := len(handler.vNextList)
 	if vNextLen == 0 {
 		panic("VMessOut: Zero vNext is configured.")
@@ -48,7 +48,7 @@ func (handler *VMessOutboundHandler) pickVNext() (v2net.Address, user.User) {
 	}
 	vNextUserIndex := mrand.Intn(vNextUserLen)
 	vNextUser := vNext.Users[vNextUserIndex]
-	return vNext.Address, vNextUser
+	return vNext.Destination, vNextUser
 }
 
 func (handler *VMessOutboundHandler) Start(ray core.OutboundRay) error {
@@ -57,8 +57,8 @@ func (handler *VMessOutboundHandler) Start(ray core.OutboundRay) error {
 	request := &protocol.VMessRequest{
 		Version: protocol.Version,
 		UserId:  vNextUser.Id,
-		Command: byte(0x01),
-		Address: handler.dest,
+		Command: handler.dest.NetworkByte(),
+		Address: handler.dest.Address(),
 	}
 	rand.Read(request.RequestIV[:])
 	rand.Read(request.RequestKey[:])
@@ -68,11 +68,11 @@ func (handler *VMessOutboundHandler) Start(ray core.OutboundRay) error {
 	return nil
 }
 
-func startCommunicate(request *protocol.VMessRequest, dest v2net.Address, ray core.OutboundRay) error {
+func startCommunicate(request *protocol.VMessRequest, dest *v2net.Destination, ray core.OutboundRay) error {
 	input := ray.OutboundInput()
 	output := ray.OutboundOutput()
 
-	conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{dest.IP, int(dest.Port), ""})
+	conn, err := net.DialTCP(dest.Network(), nil, &net.TCPAddr{dest.Address().IP, int(dest.Address().Port), ""})
 	if err != nil {
 		log.Error("Failed to open tcp (%s): %v", dest.String(), err)
 		close(output)
@@ -155,7 +155,7 @@ func handleResponse(conn *net.TCPConn, request *protocol.VMessRequest, output ch
 type VMessOutboundHandlerFactory struct {
 }
 
-func (factory *VMessOutboundHandlerFactory) Create(vp *core.Point, rawConfig []byte, destination v2net.Address) (core.OutboundConnectionHandler, error) {
+func (factory *VMessOutboundHandlerFactory) Create(vp *core.Point, rawConfig []byte, destination *v2net.Destination) (core.OutboundConnectionHandler, error) {
 	config, err := loadOutboundConfig(rawConfig)
 	if err != nil {
 		panic(log.Error("Failed to load VMess outbound config: %v", err))
diff --git a/testing/mocks/inboundhandler.go b/testing/mocks/inboundhandler.go
index 5a14bc295..70f746933 100644
--- a/testing/mocks/inboundhandler.go
+++ b/testing/mocks/inboundhandler.go
@@ -19,7 +19,7 @@ func (handler *InboundConnectionHandler) Listen(port uint16) error {
 	return nil
 }
 
-func (handler *InboundConnectionHandler) Communicate(dest v2net.Address) error {
+func (handler *InboundConnectionHandler) Communicate(dest *v2net.Destination) error {
 	ray := handler.Server.NewInboundConnectionAccepted(dest)
 
 	input := ray.InboundInput()
diff --git a/testing/mocks/outboundhandler.go b/testing/mocks/outboundhandler.go
index e5da7c723..a48593415 100644
--- a/testing/mocks/outboundhandler.go
+++ b/testing/mocks/outboundhandler.go
@@ -10,7 +10,7 @@ import (
 type OutboundConnectionHandler struct {
 	Data2Send   *bytes.Buffer
 	Data2Return []byte
-	Destination v2net.Address
+	Destination *v2net.Destination
 }
 
 func (handler *OutboundConnectionHandler) Start(ray core.OutboundRay) error {
@@ -32,7 +32,7 @@ func (handler *OutboundConnectionHandler) Start(ray core.OutboundRay) error {
 	return nil
 }
 
-func (handler *OutboundConnectionHandler) Create(point *core.Point, config []byte, dest v2net.Address) (core.OutboundConnectionHandler, error) {
+func (handler *OutboundConnectionHandler) Create(point *core.Point, config []byte, dest *v2net.Destination) (core.OutboundConnectionHandler, error) {
 	handler.Destination = dest
 	return handler, nil
 }