close method for inbound connection handler

proxy/socks/socks.go

@@ -23,9 +23,13 @@ var (
 
 // SocksServer is a SOCKS 5 proxy server
 type SocksServer struct {
-	accepting bool
-	space     app.Space
-	config    Config
+	sync.RWMutex
+	accepting   bool
+	space       app.Space
+	config      Config
+	tcpListener *net.TCPListener
+	udpConn     *net.UDPConn
+	udpAddress  v2net.Destination
 }
 
 func NewSocksServer(space app.Space, config Config) *SocksServer {
@@ -35,6 +39,26 @@ func NewSocksServer(space app.Space, config Config) *SocksServer {
 	}
 }
 
+func (this *SocksServer) Close() {
+	this.accepting = false
+	if this.tcpListener != nil {
+		this.Lock()
+		if this.tcpListener != nil {
+			this.tcpListener.Close()
+			this.tcpListener = nil
+		}
+		this.Unlock()
+	}
+	if this.udpConn != nil {
+		this.Lock()
+		if this.udpConn != nil {
+			this.udpConn.Close()
+			this.udpConn = nil
+		}
+		this.Unlock()
+	}
+}
+
 func (this *SocksServer) Listen(port v2net.Port) error {
 	listener, err := net.ListenTCP("tcp", &net.TCPAddr{
 		IP:   []byte{0, 0, 0, 0},
@@ -46,17 +70,23 @@ func (this *SocksServer) Listen(port v2net.Port) error {
 		return err
 	}
 	this.accepting = true
-	go this.AcceptConnections(listener)
+	this.tcpListener = listener
+	go this.AcceptConnections()
 	if this.config.UDPEnabled() {
 		this.ListenUDP(port)
 	}
 	return nil
 }
 
-func (this *SocksServer) AcceptConnections(listener *net.TCPListener) {
+func (this *SocksServer) AcceptConnections() {
 	for this.accepting {
 		retry.Timed(100 /* times */, 100 /* ms */).On(func() error {
-			connection, err := listener.AcceptTCP()
+			this.RLock()
+			defer this.RUnlock()
+			if !this.accepting {
+				return nil
+			}
+			connection, err := this.tcpListener.AcceptTCP()
 			if err != nil {
 				log.Error("Socks failed to accept new connection %v", err)
 				return err
@@ -64,7 +94,6 @@ func (this *SocksServer) AcceptConnections(listener *net.TCPListener) {
 			go this.HandleConnection(connection)
 			return nil
 		})
-
 	}
 }
 
@@ -187,7 +216,7 @@ func (this *SocksServer) handleUDP(reader *v2net.TimeOutReader, writer io.Writer
 	response := protocol.NewSocks5Response()
 	response.Error = protocol.ErrorSuccess
 
-	udpAddr := this.getUDPAddr()
+	udpAddr := this.udpAddress
 
 	response.Port = udpAddr.Port()
 	switch {

proxy/socks/udp.go

@@ -9,8 +9,6 @@ import (
 	"github.com/v2ray/v2ray-core/proxy/socks/protocol"
 )
 
-var udpAddress v2net.Destination
-
 func (this *SocksServer) ListenUDP(port v2net.Port) error {
 	addr := &net.UDPAddr{
 		IP:   net.IP{0, 0, 0, 0},
@@ -22,20 +20,23 @@ func (this *SocksServer) ListenUDP(port v2net.Port) error {
 		log.Error("Socks failed to listen UDP on port %d: %v", port, err)
 		return err
 	}
-	udpAddress = v2net.UDPDestination(v2net.IPAddress(this.config.IP()), port)
+	this.udpAddress = v2net.UDPDestination(v2net.IPAddress(this.config.IP()), port)
+	this.udpConn = conn
 
-	go this.AcceptPackets(conn)
+	go this.AcceptPackets()
 	return nil
 }
 
-func (this *SocksServer) getUDPAddr() v2net.Destination {
-	return udpAddress
-}
-
-func (this *SocksServer) AcceptPackets(conn *net.UDPConn) error {
-	for {
+func (this *SocksServer) AcceptPackets() error {
+	for this.accepting {
 		buffer := alloc.NewBuffer()
-		nBytes, addr, err := conn.ReadFromUDP(buffer.Value)
+		this.RLock()
+		if !this.accepting {
+			this.RUnlock()
+			return nil
+		}
+		nBytes, addr, err := this.udpConn.ReadFromUDP(buffer.Value)
+		this.RUnlock()
 		if err != nil {
 			log.Error("Socks failed to read UDP packets: %v", err)
 			buffer.Release()
@@ -60,11 +61,12 @@ func (this *SocksServer) AcceptPackets(conn *net.UDPConn) error {
 
 		udpPacket := v2net.NewPacket(request.Destination(), request.Data, false)
 		log.Info("Send packet to %s with %d bytes", udpPacket.Destination().String(), request.Data.Len())
-		go this.handlePacket(conn, udpPacket, addr, request.Address, request.Port)
+		go this.handlePacket(udpPacket, addr, request.Address, request.Port)
 	}
+	return nil
 }
 
-func (this *SocksServer) handlePacket(conn *net.UDPConn, packet v2net.Packet, clientAddr *net.UDPAddr, targetAddr v2net.Address, port v2net.Port) {
+func (this *SocksServer) handlePacket(packet v2net.Packet, clientAddr *net.UDPAddr, targetAddr v2net.Address, port v2net.Port) {
 	ray := this.space.PacketDispatcher().DispatchToOutbound(packet)
 	close(ray.InboundInput())
 
@@ -80,7 +82,13 @@ func (this *SocksServer) handlePacket(conn *net.UDPConn, packet v2net.Packet, cl
 		udpMessage := alloc.NewSmallBuffer().Clear()
 		response.Write(udpMessage)
 
-		nBytes, err := conn.WriteToUDP(udpMessage.Value, clientAddr)
+		this.RLock()
+		if !this.accepting {
+			this.RUnlock()
+			return
+		}
+		nBytes, err := this.udpConn.WriteToUDP(udpMessage.Value, clientAddr)
+		this.RUnlock()
 		udpMessage.Release()
 		response.Data.Release()
 		if err != nil {