massive refactoring for kcp

transport/internet/tcp/hub.go

@@ -0,0 +1,159 @@
+package tcp
+
+import (
+	"errors"
+	"net"
+	"sync"
+	"time"
+
+	v2net "github.com/v2ray/v2ray-core/common/net"
+	"github.com/v2ray/v2ray-core/transport/internet"
+)
+
+var (
+	ErrClosedListener = errors.New("Listener is closed.")
+)
+
+type ConnectionWithError struct {
+	conn net.Conn
+	err  error
+}
+
+type TCPListener struct {
+	sync.Mutex
+	acccepting    bool
+	listener      *net.TCPListener
+	awaitingConns chan *ConnectionWithError
+}
+
+func ListenTCP(address v2net.Address, port v2net.Port) (internet.Listener, error) {
+	listener, err := net.ListenTCP("tcp", &net.TCPAddr{
+		IP:   address.IP(),
+		Port: int(port),
+	})
+	if err != nil {
+		return nil, err
+	}
+	l := &TCPListener{
+		acccepting:    true,
+		listener:      listener,
+		awaitingConns: make(chan *ConnectionWithError, 32),
+	}
+	go l.KeepAccepting()
+	return l, nil
+}
+
+func (this *TCPListener) Accept() (internet.Connection, error) {
+	for this.acccepting {
+		select {
+		case connErr, open := <-this.awaitingConns:
+			if !open {
+				return nil, ErrClosedListener
+			}
+			if connErr.err != nil {
+				return nil, connErr.err
+			}
+			return NewConnection("", connErr.conn, this), nil
+		case <-time.After(time.Second * 2):
+		}
+	}
+	return nil, ErrClosedListener
+}
+
+func (this *TCPListener) KeepAccepting() {
+	for this.acccepting {
+		conn, err := this.listener.Accept()
+		this.Lock()
+		if !this.acccepting {
+			this.Unlock()
+			break
+		}
+		select {
+		case this.awaitingConns <- &ConnectionWithError{
+			conn: conn,
+			err:  err,
+		}:
+		default:
+			if conn != nil {
+				conn.Close()
+			}
+		}
+
+		this.Unlock()
+	}
+}
+
+func (this *TCPListener) Recycle(dest string, conn net.Conn) {
+	this.Lock()
+	defer this.Unlock()
+	if !this.acccepting {
+		return
+	}
+	select {
+	case this.awaitingConns <- &ConnectionWithError{conn: conn}:
+	default:
+		conn.Close()
+	}
+}
+
+func (this *TCPListener) Addr() net.Addr {
+	return this.listener.Addr()
+}
+
+func (this *TCPListener) Close() error {
+	this.Lock()
+	defer this.Unlock()
+	this.acccepting = false
+	this.listener.Close()
+	close(this.awaitingConns)
+	for connErr := range this.awaitingConns {
+		if connErr.conn != nil {
+			go connErr.conn.Close()
+		}
+	}
+	return nil
+}
+
+type RawTCPListener struct {
+	accepting bool
+	listener  *net.TCPListener
+}
+
+func (this *RawTCPListener) Accept() (internet.Connection, error) {
+	conn, err := this.listener.AcceptTCP()
+	if err != nil {
+		return nil, err
+	}
+	return &RawConnection{
+		TCPConn: *conn,
+	}, nil
+}
+
+func (this *RawTCPListener) Addr() net.Addr {
+	return this.listener.Addr()
+}
+
+func (this *RawTCPListener) Close() error {
+	this.accepting = false
+	this.listener.Close()
+	return nil
+}
+
+func ListenRawTCP(address v2net.Address, port v2net.Port) (internet.Listener, error) {
+	listener, err := net.ListenTCP("tcp", &net.TCPAddr{
+		IP:   address.IP(),
+		Port: int(port),
+	})
+	if err != nil {
+		return nil, err
+	}
+	return &RawTCPListener{
+		accepting: true,
+		listener:  listener,
+	}, nil
+}
+
+func init() {
+	internet.TCPListenFunc = ListenTCP
+	internet.RawTCPListenFunc = ListenRawTCP
+}