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

Merge branch 'udp-session-timeout' into notsure2

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

17
internal/client/piper.go
View File

@ -4,6 +4,7 @@ import (
"github.com/cbeuw/Cloak/internal/common" "github.com/cbeuw/Cloak/internal/common"
"io" "io"
"net" "net"
"sync"
"time" "time"
mux "github.com/cbeuw/Cloak/internal/multiplex" mux "github.com/cbeuw/Cloak/internal/multiplex"
@ -17,7 +18,7 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration
log.Fatal(err) log.Fatal(err)
} }
streams := make(map[string]*mux.Stream) var streams sync.Map
data := make([]byte, 8192) data := make([]byte, 8192)
for { for {
@ -31,13 +32,15 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration
sesh = newSeshFunc() sesh = newSeshFunc()
} }
stream, ok := streams[addr.String()] var stream *mux.Stream
streamObj, ok := streams.Load(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()
@ -45,8 +48,9 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration
log.Errorf("Failed to open stream: %v", err) log.Errorf("Failed to open stream: %v", err)
continue continue
} }
_ = stream.SetReadDeadline(time.Now().Add(streamTimeout))
streams[addr.String()] = stream streams.Store(addr.String(), stream)
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)
@ -54,13 +58,16 @@ 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())
stream.Close() stream.Close()
return return
} }
_ = stream.SetReadDeadline(time.Now().Add(streamTimeout))
_, 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())
stream.Close() stream.Close()
return return
} }
@ -68,13 +75,15 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration
}(stream, localConn) }(stream, localConn)
} }
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)
delete(streams, addr.String()) streams.Delete(addr.String())
stream.Close() stream.Close()
continue continue
} }
_ = stream.SetReadDeadline(time.Now().Add(streamTimeout))
} }
} }