kirillius
/
Cloak
mirror of https://github.com/cbeuw/Cloak
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Use a mutex instead of sync.Map in RouteUDP.

This commit is contained in:
notsure2 2020-12-23 02:32:13 +02:00
parent 874d30c512
commit a26a2e84a4
1 changed files with 18 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
internal/client/piper.go
View File

@ -18,7 +18,8 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration
log.Fatal(err) log.Fatal(err)
} }
var streams sync.Map streams := make(map[string]*mux.Stream)
var streamsMutex sync.Mutex
data := make([]byte, 8192) data := make([]byte, 8192)
for { for {
@ -32,15 +33,14 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration
sesh = newSeshFunc() sesh = newSeshFunc()
} }
var stream *mux.Stream streamsMutex.Lock()
streamObj, ok := streams.Load(addr.String()) stream, ok := streams[addr.String()]
if !ok { if !ok {
if singleplex { if singleplex {
sesh = newSeshFunc() sesh = newSeshFunc()
} }
stream, err = sesh.OpenStream() stream, err = sesh.OpenStream()
streamObj = stream
if err != nil { if err != nil {
if singleplex { if singleplex {
sesh.Close() sesh.Close()
@ -50,7 +50,9 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration
} }
_ = stream.SetReadDeadline(time.Now().Add(streamTimeout)) _ = stream.SetReadDeadline(time.Now().Add(streamTimeout))
streams.Store(addr.String(), stream) streams[addr.String()] = stream
streamsMutex.Unlock()
proxyAddr := addr proxyAddr := addr
go func(stream *mux.Stream, localConn *net.UDPConn) { go func(stream *mux.Stream, localConn *net.UDPConn) {
buf := make([]byte, 8192) buf := make([]byte, 8192)
@ -58,7 +60,9 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration
n, err := stream.Read(buf) n, err := stream.Read(buf)
if err != nil { if err != nil {
log.Tracef("copying stream to proxy client: %v", err) log.Tracef("copying stream to proxy client: %v", err)
streams.Delete(addr.String()) streamsMutex.Lock()
delete(streams, addr.String())
streamsMutex.Unlock()
stream.Close() stream.Close()
return return
} }
@ -67,19 +71,24 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration
_, err = localConn.WriteTo(buf[:n], proxyAddr) _, err = localConn.WriteTo(buf[:n], proxyAddr)
if err != nil { if err != nil {
log.Tracef("copying stream to proxy client: %v", err) log.Tracef("copying stream to proxy client: %v", err)
streams.Delete(addr.String()) streamsMutex.Lock()
delete(streams, addr.String())
streamsMutex.Unlock()
stream.Close() stream.Close()
return return
} }
} }
}(stream, localConn) }(stream, localConn)
} else {
streamsMutex.Unlock()
} }
stream = streamObj.(*mux.Stream)
_, err = stream.Write(data[:i]) _, err = stream.Write(data[:i])
if err != nil { if err != nil {
log.Tracef("copying proxy client to stream: %v", err) log.Tracef("copying proxy client to stream: %v", err)
streams.Delete(addr.String()) streamsMutex.Lock()
delete(streams, addr.String())
streamsMutex.Unlock()
stream.Close() stream.Close()
continue continue
} }