From 838cb8950ac7ee783bc0fc07acd4ef1e8615b579 Mon Sep 17 00:00:00 2001
From: V2Ray <admin@v2ray.com>
Date: Sun, 4 Oct 2015 00:21:06 +0200
Subject: [PATCH] Enable UDP in Socks proxy

---
 common/net/timed_io.go |  8 ++++++++
 proxy/socks/socks.go   | 40 ++++++++++++++++++++++++++++++++++++++--
 proxy/socks/udp.go     |  9 ++++++++-
 3 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/common/net/timed_io.go b/common/net/timed_io.go
index cc9fb5194..795165ca5 100644
--- a/common/net/timed_io.go
+++ b/common/net/timed_io.go
@@ -28,3 +28,11 @@ func (reader *TimeOutReader) Read(p []byte) (n int, err error) {
 	reader.connection.SetReadDeadline(emptyTime)
 	return
 }
+
+func (reader *TimeOutReader) GetTimeOut() int {
+	return reader.timeout
+}
+
+func (reader *TimeOutReader) SetTimeOut(value int) {
+	reader.timeout = value
+}
diff --git a/proxy/socks/socks.go b/proxy/socks/socks.go
index d33879f8f..92d2fe389 100644
--- a/proxy/socks/socks.go
+++ b/proxy/socks/socks.go
@@ -76,7 +76,7 @@ func (server *SocksServer) HandleConnection(connection net.Conn) error {
 	}
 }
 
-func (server *SocksServer) handleSocks5(reader io.Reader, writer io.Writer, auth protocol.Socks5AuthenticationRequest) error {
+func (server *SocksServer) handleSocks5(reader *v2net.TimeOutReader, writer io.Writer, auth protocol.Socks5AuthenticationRequest) error {
 	expectedAuthMethod := protocol.AuthNotRequired
 	if server.config.IsPassword() {
 		expectedAuthMethod = protocol.AuthUserPass
@@ -130,7 +130,11 @@ func (server *SocksServer) handleSocks5(reader io.Reader, writer io.Writer, auth
 
 	response := protocol.NewSocks5Response()
 
-	if request.Command == protocol.CmdBind || (!server.config.UDPEnabled && request.Command == protocol.CmdUdpAssociate) {
+	if request.Command == protocol.CmdUdpAssociate && server.config.UDPEnabled {
+		return server.handleUDP(reader, writer)
+	}
+
+	if request.Command == protocol.CmdBind || request.Command == protocol.CmdUdpAssociate {
 		response := protocol.NewSocks5Response()
 		response.Error = protocol.ErrorCommandNotSupported
 		err = protocol.WriteResponse(writer, response)
@@ -143,6 +147,7 @@ func (server *SocksServer) handleSocks5(reader io.Reader, writer io.Writer, auth
 	}
 
 	response.Error = protocol.ErrorSuccess
+
 	response.Port = request.Port
 	response.AddrType = request.AddrType
 	switch response.AddrType {
@@ -170,6 +175,37 @@ func (server *SocksServer) handleSocks5(reader io.Reader, writer io.Writer, auth
 	return nil
 }
 
+func (server *SocksServer) handleUDP(reader *v2net.TimeOutReader, writer io.Writer) error {
+	response := protocol.NewSocks5Response()
+	response.Error = protocol.ErrorSuccess
+
+	udpAddr := server.getUDPAddr()
+
+	response.Port = udpAddr.Port()
+	switch {
+	case udpAddr.IsIPv4():
+		response.AddrType = protocol.AddrTypeIPv4
+		copy(response.IPv4[:], udpAddr.IP())
+	case udpAddr.IsIPv6():
+		response.AddrType = protocol.AddrTypeIPv6
+		copy(response.IPv6[:], udpAddr.IP())
+	case udpAddr.IsDomain():
+		response.AddrType = protocol.AddrTypeDomain
+		response.Domain = udpAddr.Domain()
+	}
+	err := protocol.WriteResponse(writer, response)
+	if err != nil {
+		log.Error("Socks failed to write response: %v", err)
+		return err
+	}
+
+	reader.SetTimeOut(300) /* 5 minutes */
+	buffer := make([]byte, 1024)
+	reader.Read(buffer)
+
+	return nil
+}
+
 func (server *SocksServer) handleSocks4(reader io.Reader, writer io.Writer, auth protocol.Socks4AuthenticationRequest) error {
 	result := protocol.Socks4RequestGranted
 	if auth.Command == protocol.CmdBind {
diff --git a/proxy/socks/udp.go b/proxy/socks/udp.go
index 28715299f..e8842b97b 100644
--- a/proxy/socks/udp.go
+++ b/proxy/socks/udp.go
@@ -12,6 +12,8 @@ const (
 	bufferSize = 2 * 1024
 )
 
+var udpAddress v2net.Address
+
 func (server *SocksServer) ListenUDP(port uint16) error {
 	addr := &net.UDPAddr{
 		IP:   net.IP{0, 0, 0, 0},
@@ -23,14 +25,19 @@ func (server *SocksServer) ListenUDP(port uint16) error {
 		log.Error("Socks failed to listen UDP on port %d: %v", port, err)
 		return err
 	}
+	udpAddress = v2net.IPAddress([]byte{127, 0, 0, 1}, port)
 
 	go server.AcceptPackets(conn)
 	return nil
 }
 
+func (server *SocksServer) getUDPAddr() v2net.Address {
+	return udpAddress
+}
+
 func (server *SocksServer) AcceptPackets(conn *net.UDPConn) error {
 	for {
-		buffer := make([]byte, 0, bufferSize)
+		buffer := make([]byte, bufferSize)
 		nBytes, addr, err := conn.ReadFromUDP(buffer)
 		if err != nil {
 			log.Error("Socks failed to read UDP packets: %v", err)