Skip to content

Commit

Permalink
fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
ginuerzh committed Feb 8, 2020
1 parent 94dcfca commit ece7994
Showing 1 changed file with 20 additions and 4 deletions.
24 changes: 20 additions & 4 deletions forward.go
Original file line number Diff line number Diff line change
Expand Up @@ -361,6 +361,10 @@ func TCPRemoteForwardListener(addr string, chain *Chain) (Listener, error) {
}

func (l *tcpRemoteForwardListener) isChainValid() bool {
if l.chain.IsEmpty() {
return false
}

lastNode := l.chain.LastNode()
if (lastNode.Protocol == "forward" && lastNode.Transport == "ssh") ||
lastNode.Protocol == "socks5" || lastNode.Protocol == "" {
Expand Down Expand Up @@ -429,7 +433,7 @@ func (l *tcpRemoteForwardListener) accept() (conn net.Conn, err error) {
return l.chain.Dial(l.addr.String())
}

if lastNode.Protocol == "socks5" || lastNode.Protocol == "" {
if l.isChainValid() {
if lastNode.GetBool("mbind") {
return l.muxAccept() // multiplexing support for binding.
}
Expand Down Expand Up @@ -588,6 +592,8 @@ type udpRemoteForwardListener struct {
ln *net.UDPConn
ttl time.Duration
closed chan struct{}
ready chan struct{}
once sync.Once
closeMux sync.Mutex
config *UDPListenConfig
}
Expand All @@ -613,18 +619,25 @@ func UDPRemoteForwardListener(addr string, chain *Chain, cfg *UDPListenConfig) (
chain: chain,
connMap: new(udpConnMap),
connChan: make(chan net.Conn, backlog),
ready: make(chan struct{}),
closed: make(chan struct{}),
config: cfg,
}

go ln.listenLoop()

<-ln.ready

return ln, err
}

func (l *udpRemoteForwardListener) isChainValid() bool {
if l.chain.IsEmpty() {
return false
}

lastNode := l.chain.LastNode()
return lastNode.Protocol == "socks5"
return lastNode.Protocol == "socks5" || lastNode.Protocol == ""
}

func (l *udpRemoteForwardListener) listenLoop() {
Expand All @@ -635,6 +648,10 @@ func (l *udpRemoteForwardListener) listenLoop() {
return
}

l.once.Do(func() {
close(l.ready)
})

func() {
defer conn.Close()

Expand Down Expand Up @@ -691,8 +708,7 @@ func (l *udpRemoteForwardListener) connect() (conn net.PacketConn, err error) {
default:
}

lastNode := l.chain.LastNode()
if lastNode.Protocol == "socks5" || lastNode.Protocol == "" {
if l.isChainValid() {
var cc net.Conn
cc, err = getSocks5UDPTunnel(l.chain, l.addr)
if err != nil {
Expand Down

0 comments on commit ece7994

Please sign in to comment.