Code style improvements

This commit is contained in:
Andy Wang 2020-12-18 23:26:00 +00:00
parent 900387832d
commit be53fa37a0
No known key found for this signature in database
GPG Key ID: 181B49F9F38F3374
1 changed files with 14 additions and 26 deletions

View File

@ -11,7 +11,7 @@ import (
) )
func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration, singleplex bool, newSeshFunc func() *mux.Session) { func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration, singleplex bool, newSeshFunc func() *mux.Session) {
var multiplexSession *mux.Session var sesh *mux.Session
localConn, err := bindFunc() localConn, err := bindFunc()
if err != nil { if err != nil {
log.Fatal(err) log.Fatal(err)
@ -27,25 +27,19 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration
continue continue
} }
if !singleplex && (multiplexSession == nil || multiplexSession.IsClosed()) { if !singleplex && (sesh == nil || sesh.IsClosed()) {
multiplexSession = newSeshFunc() sesh = newSeshFunc()
} }
stream, ok := streams[addr.String()] stream, ok := streams[addr.String()]
if !ok { if !ok {
var session *mux.Session if singleplex {
if multiplexSession != nil { sesh = newSeshFunc()
session = multiplexSession
} else {
session = newSeshFunc()
} }
stream, err = session.OpenStream() stream, err = sesh.OpenStream()
if err != nil { if err != nil {
log.Errorf("Failed to open stream: %v", err) log.Errorf("Failed to open stream: %v", err)
if session.Singleplex {
session.Close()
}
continue continue
} }
@ -82,22 +76,19 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration
} }
func RouteTCP(listener net.Listener, streamTimeout time.Duration, singleplex bool, newSeshFunc func() *mux.Session) { func RouteTCP(listener net.Listener, streamTimeout time.Duration, singleplex bool, newSeshFunc func() *mux.Session) {
var multiplexSession *mux.Session var sesh *mux.Session
for { for {
localConn, err := listener.Accept() localConn, err := listener.Accept()
if err != nil { if err != nil {
log.Fatal(err) log.Fatal(err)
continue continue
} }
if !singleplex && (multiplexSession == nil || multiplexSession.IsClosed()) { if !singleplex && (sesh == nil || sesh.IsClosed()) {
multiplexSession = newSeshFunc() sesh = newSeshFunc()
} }
go func(multiplexSession *mux.Session, newSingleplexSeshFunc func() *mux.Session, localConn net.Conn, timeout time.Duration) { go func(sesh *mux.Session, localConn net.Conn, timeout time.Duration) {
var session *mux.Session if singleplex {
if multiplexSession != nil { sesh = newSeshFunc()
session = multiplexSession
} else {
session = newSingleplexSeshFunc()
} }
data := make([]byte, 10240) data := make([]byte, 10240)
@ -111,13 +102,10 @@ func RouteTCP(listener net.Listener, streamTimeout time.Duration, singleplex boo
var zeroTime time.Time var zeroTime time.Time
_ = localConn.SetReadDeadline(zeroTime) _ = localConn.SetReadDeadline(zeroTime)
stream, err := session.OpenStream() stream, err := sesh.OpenStream()
if err != nil { if err != nil {
log.Errorf("Failed to open stream: %v", err) log.Errorf("Failed to open stream: %v", err)
localConn.Close() localConn.Close()
if session.Singleplex {
session.Close()
}
return return
} }
@ -137,6 +125,6 @@ func RouteTCP(listener net.Listener, streamTimeout time.Duration, singleplex boo
if _, err = common.Copy(stream, localConn); err != nil { if _, err = common.Copy(stream, localConn); err != nil {
log.Tracef("copying proxy client to stream: %v", err) log.Tracef("copying proxy client to stream: %v", err)
} }
}(multiplexSession, newSeshFunc, localConn, streamTimeout) }(sesh, localConn, streamTimeout)
} }
} }