diff --git a/proxy/http/server.go b/proxy/http/server.go index db57b249..d95d91ad 100644 --- a/proxy/http/server.go +++ b/proxy/http/server.go @@ -14,21 +14,25 @@ var ( tun = tunnel.Instance() ) -func NewHttpProxy(addr string) (chan<- struct{}, <-chan struct{}, error) { +type httpListener struct { + net.Listener + address string + closed bool +} + +func NewHttpProxy(addr string) (*httpListener, error) { l, err := net.Listen("tcp", addr) if err != nil { - return nil, nil, err + return nil, err } - - done := make(chan struct{}) - closed := make(chan struct{}) + hl := &httpListener{l, addr, false} go func() { log.Infoln("HTTP proxy listening at: %s", addr) for { - c, err := l.Accept() + c, err := hl.Accept() if err != nil { - if _, open := <-done; !open { + if hl.closed { break } continue @@ -37,13 +41,16 @@ func NewHttpProxy(addr string) (chan<- struct{}, <-chan struct{}, error) { } }() - go func() { - <-done - l.Close() - closed <- struct{}{} - }() + return hl, nil +} - return done, closed, nil +func (l *httpListener) Close() { + l.closed = true + l.Listener.Close() +} + +func (l *httpListener) Address() string { + return l.address } func handleConn(conn net.Conn) { diff --git a/proxy/listener.go b/proxy/listener.go index 5e635ba5..82922aab 100644 --- a/proxy/listener.go +++ b/proxy/listener.go @@ -13,15 +13,14 @@ import ( var ( allowLan = false - socksListener *listener - httpListener *listener - redirListener *listener + socksListener listener + httpListener listener + redirListener listener ) -type listener struct { - Address string - Done chan<- struct{} - Closed <-chan struct{} +type listener interface { + Close() + Address() string } type Ports struct { @@ -42,11 +41,10 @@ func ReCreateHTTP(port int) error { addr := genAddr(port, allowLan) if httpListener != nil { - if httpListener.Address == addr { + if httpListener.Address() == addr { return nil } - httpListener.Done <- struct{}{} - <-httpListener.Closed + httpListener.Close() httpListener = nil } @@ -54,16 +52,12 @@ func ReCreateHTTP(port int) error { return nil } - done, closed, err := http.NewHttpProxy(addr) + var err error + httpListener, err = http.NewHttpProxy(addr) if err != nil { return err } - httpListener = &listener{ - Address: addr, - Done: done, - Closed: closed, - } return nil } @@ -71,11 +65,10 @@ func ReCreateSocks(port int) error { addr := genAddr(port, allowLan) if socksListener != nil { - if socksListener.Address == addr { + if socksListener.Address() == addr { return nil } - socksListener.Done <- struct{}{} - <-socksListener.Closed + socksListener.Close() socksListener = nil } @@ -83,16 +76,12 @@ func ReCreateSocks(port int) error { return nil } - done, closed, err := socks.NewSocksProxy(addr) + var err error + socksListener, err = socks.NewSocksProxy(addr) if err != nil { return err } - socksListener = &listener{ - Address: addr, - Done: done, - Closed: closed, - } return nil } @@ -100,11 +89,10 @@ func ReCreateRedir(port int) error { addr := genAddr(port, allowLan) if redirListener != nil { - if redirListener.Address == addr { + if redirListener.Address() == addr { return nil } - redirListener.Done <- struct{}{} - <-redirListener.Closed + redirListener.Close() redirListener = nil } @@ -112,16 +100,12 @@ func ReCreateRedir(port int) error { return nil } - done, closed, err := redir.NewRedirProxy(addr) + var err error + redirListener, err = redir.NewRedirProxy(addr) if err != nil { return err } - redirListener = &listener{ - Address: addr, - Done: done, - Closed: closed, - } return nil } @@ -130,19 +114,19 @@ func GetPorts() *Ports { ports := &Ports{} if httpListener != nil { - _, portStr, _ := net.SplitHostPort(httpListener.Address) + _, portStr, _ := net.SplitHostPort(httpListener.Address()) port, _ := strconv.Atoi(portStr) ports.Port = port } if socksListener != nil { - _, portStr, _ := net.SplitHostPort(socksListener.Address) + _, portStr, _ := net.SplitHostPort(socksListener.Address()) port, _ := strconv.Atoi(portStr) ports.SocksPort = port } if redirListener != nil { - _, portStr, _ := net.SplitHostPort(redirListener.Address) + _, portStr, _ := net.SplitHostPort(redirListener.Address()) port, _ := strconv.Atoi(portStr) ports.RedirPort = port } diff --git a/proxy/redir/tcp.go b/proxy/redir/tcp.go index bb6d3c55..d54a17c4 100644 --- a/proxy/redir/tcp.go +++ b/proxy/redir/tcp.go @@ -13,21 +13,25 @@ var ( tun = tunnel.Instance() ) -func NewRedirProxy(addr string) (chan<- struct{}, <-chan struct{}, error) { +type redirListener struct { + net.Listener + address string + closed bool +} + +func NewRedirProxy(addr string) (*redirListener, error) { l, err := net.Listen("tcp", addr) if err != nil { - return nil, nil, err + return nil, err } - - done := make(chan struct{}) - closed := make(chan struct{}) + rl := &redirListener{l, addr, false} go func() { log.Infof("Redir proxy listening at: %s", addr) for { c, err := l.Accept() if err != nil { - if _, open := <-done; !open { + if rl.closed { break } continue @@ -36,13 +40,16 @@ func NewRedirProxy(addr string) (chan<- struct{}, <-chan struct{}, error) { } }() - go func() { - <-done - l.Close() - closed <- struct{}{} - }() + return rl, nil +} - return done, closed, nil +func (l *redirListener) Close() { + l.closed = true + l.Listener.Close() +} + +func (l *redirListener) Address() string { + return l.address } func handleRedir(conn net.Conn) { diff --git a/proxy/socks/tcp.go b/proxy/socks/tcp.go index 646050be..a4d4db4b 100644 --- a/proxy/socks/tcp.go +++ b/proxy/socks/tcp.go @@ -14,21 +14,25 @@ var ( tun = tunnel.Instance() ) -func NewSocksProxy(addr string) (chan<- struct{}, <-chan struct{}, error) { +type sockListener struct { + net.Listener + address string + closed bool +} + +func NewSocksProxy(addr string) (*sockListener, error) { l, err := net.Listen("tcp", addr) if err != nil { - return nil, nil, err + return nil, err } - done := make(chan struct{}) - closed := make(chan struct{}) - + sl := &sockListener{l, addr, false} go func() { log.Infof("SOCKS proxy listening at: %s", addr) for { c, err := l.Accept() if err != nil { - if _, open := <-done; !open { + if sl.closed { break } continue @@ -37,13 +41,16 @@ func NewSocksProxy(addr string) (chan<- struct{}, <-chan struct{}, error) { } }() - go func() { - <-done - l.Close() - closed <- struct{}{} - }() + return sl, nil +} - return done, closed, nil +func (l *sockListener) Close() { + l.closed = true + l.Listener.Close() +} + +func (l *sockListener) Address() string { + return l.address } func handleSocks(conn net.Conn) {