diff --git a/faketls/client_protocol.go b/faketls/client_protocol.go index b2fa024..1f6a207 100644 --- a/faketls/client_protocol.go +++ b/faketls/client_protocol.go @@ -8,7 +8,6 @@ import ( "io" "net" "strconv" - "sync" "time" "github.com/9seconds/mtg/antireplay" @@ -110,25 +109,7 @@ func (c *ClientProtocol) cloakHost(clientConn io.ReadWriteCloser) { return } - defer hostConn.Close() - - wg := &sync.WaitGroup{} - wg.Add(2) - - go c.pipe(hostConn, clientConn, wg) - - go c.pipe(clientConn, hostConn, wg) - - wg.Wait() -} - -func (c *ClientProtocol) pipe(dst io.WriteCloser, src io.Reader, wg *sync.WaitGroup) { - defer func() { - wg.Done() - dst.Close() - }() - - io.Copy(dst, src) // nolint: errcheck + cloak(clientConn, hostConn) } func MakeClientProtocol() protocol.ClientProtocol { diff --git a/faketls/cloak.go b/faketls/cloak.go new file mode 100644 index 0000000..a6ea80f --- /dev/null +++ b/faketls/cloak.go @@ -0,0 +1,62 @@ +package faketls + +import ( + "context" + "io" + "sync" + "time" + + "github.com/9seconds/mtg/wrappers/rwc" +) + +const cloakTimeout = 5 * time.Second + +func cloak(one, another io.ReadWriteCloser) { + defer func() { + one.Close() + another.Close() + }() + + channelPing := make(chan struct{}, 1) + ctx, cancel := context.WithCancel(context.Background()) + one = rwc.NewPing(ctx, one, channelPing) + another = rwc.NewPing(ctx, another, channelPing) + wg := &sync.WaitGroup{} + + wg.Add(2) + + go func() { + defer wg.Done() + io.Copy(one, another) // nolint: errcheck + }() + + go func() { + defer wg.Done() + io.Copy(another, one) // nolint: errcheck + }() + + go func() { + wg.Wait() + cancel() + }() + + go func() { + timer := time.NewTimer(cloakTimeout) + defer timer.Stop() + + for { + select { + case <-channelPing: + timer.Stop() + timer = time.NewTimer(cloakTimeout) + case <-ctx.Done(): + return + case <-timer.C: + cancel() + return + } + } + }() + + <-ctx.Done() +} diff --git a/wrappers/rwc/ping.go b/wrappers/rwc/ping.go new file mode 100644 index 0000000..2b224c2 --- /dev/null +++ b/wrappers/rwc/ping.go @@ -0,0 +1,48 @@ +package rwc + +import ( + "context" + "io" +) + +type wrapperPing struct { + parent io.ReadWriteCloser + ctx context.Context + channelPing chan<- struct{} +} + +func (w *wrapperPing) Read(p []byte) (int, error) { + n, err := w.parent.Read(p) + if err == nil { + select { + case <-w.ctx.Done(): + case w.channelPing <- struct{}{}: + } + } + + return n, err +} + +func (w *wrapperPing) Write(p []byte) (int, error) { + n, err := w.parent.Write(p) + if err == nil { + select { + case <-w.ctx.Done(): + case w.channelPing <- struct{}{}: + } + } + + return n, err +} + +func (w *wrapperPing) Close() error { + return w.parent.Close() +} + +func NewPing(ctx context.Context, parent io.ReadWriteCloser, channelPing chan<- struct{}) io.ReadWriteCloser { + return &wrapperPing{ + parent: parent, + ctx: ctx, + channelPing: channelPing, + } +}