diff --git a/client/client.go b/client/client.go index 379ead8..648fec2 100644 --- a/client/client.go +++ b/client/client.go @@ -1,7 +1,6 @@ package client import ( - "context" "net" "github.com/9seconds/mtg/config" @@ -9,4 +8,4 @@ import ( "github.com/9seconds/mtg/wrappers" ) -type Init func(context.Context, context.CancelFunc, net.Conn, string, *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) +type Init func(net.Conn, string, *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) diff --git a/client/direct.go b/client/direct.go index 09c0178..2c16436 100644 --- a/client/direct.go +++ b/client/direct.go @@ -1,7 +1,6 @@ package client import ( - "context" "net" "time" @@ -15,8 +14,7 @@ import ( const handshakeTimeout = 10 * time.Second -func DirectInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn, connID string, - conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) { +func DirectInit(socket net.Conn, connID string, conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) { if err := config.SetSocketOptions(socket); err != nil { return nil, nil, errors.Annotate(err, "Cannot set socket options") } @@ -37,7 +35,6 @@ func DirectInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn, connOpts.ClientAddr = conn.RemoteAddr() conn = wrappers.NewStreamCipher(conn, obfs2.Encryptor, obfs2.Decryptor) - conn = wrappers.NewCtx(ctx, cancel, conn) return conn, connOpts, nil } diff --git a/client/middle.go b/client/middle.go index 4f0be19..b85deef 100644 --- a/client/middle.go +++ b/client/middle.go @@ -1,7 +1,6 @@ package client import ( - "context" "net" "github.com/9seconds/mtg/config" @@ -9,9 +8,8 @@ import ( "github.com/9seconds/mtg/wrappers" ) -func MiddleInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn, connID string, - conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) { - conn, opts, err := DirectInit(ctx, cancel, socket, connID, conf) +func MiddleInit(socket net.Conn, connID string, conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) { + conn, opts, err := DirectInit(socket, connID, conf) if err != nil { return nil, nil, err } diff --git a/proxy/proxy.go b/proxy/proxy.go index f559f68..d11fc02 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -1,7 +1,6 @@ package proxy import ( - "context" "io" "net" "sync" @@ -52,15 +51,14 @@ func (p *Proxy) accept(conn net.Conn) { log.Infow("Client connected", "addr", conn.RemoteAddr()) - ctx, cancel := context.WithCancel(context.Background()) - client, opts, err := p.clientInit(ctx, cancel, conn, connID, p.conf) + client, opts, err := p.clientInit(conn, connID, p.conf) if err != nil { log.Errorw("Cannot initialize client connection", "error", err) return } defer client.(io.Closer).Close() - server, err := p.getTelegramConn(ctx, cancel, opts, connID) + server, err := p.getTelegramConn(opts, connID) if err != nil { log.Errorw("Cannot initialize server connection", "error", err) return @@ -82,19 +80,16 @@ func (p *Proxy) accept(conn net.Conn) { go p.directPipe(serverStream, clientStream, wait) } - <-ctx.Done() wait.Wait() log.Infow("Client disconnected", "addr", conn.RemoteAddr()) } -func (p *Proxy) getTelegramConn(ctx context.Context, cancel context.CancelFunc, opts *mtproto.ConnectionOpts, - connID string) (wrappers.Wrap, error) { +func (p *Proxy) getTelegramConn(opts *mtproto.ConnectionOpts, connID string) (wrappers.Wrap, error) { streamConn, err := p.tg.Dial(connID, opts) if err != nil { return nil, errors.Annotate(err, "Cannot dial to Telegram") } - streamConn = wrappers.NewCtx(ctx, cancel, streamConn) packetConn, err := p.tg.Init(opts, streamConn) if err != nil { @@ -104,8 +99,13 @@ func (p *Proxy) getTelegramConn(ctx context.Context, cancel context.CancelFunc, return packetConn, nil } -func (p *Proxy) middlePipe(src wrappers.PacketReader, dst wrappers.PacketWriter, wait *sync.WaitGroup, hacks *mtproto.Hacks) { - defer wait.Done() +func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst wrappers.PacketWriteCloser, wait *sync.WaitGroup, hacks *mtproto.Hacks) { + defer func() { + src.Close() + dst.Close() + wait.Done() + }() + for { hacks.SimpleAck = false hacks.QuickAck = false @@ -120,10 +120,14 @@ func (p *Proxy) middlePipe(src wrappers.PacketReader, dst wrappers.PacketWriter, } } -func (p *Proxy) directPipe(src io.Reader, dst io.Writer, wait *sync.WaitGroup) { - defer wait.Done() - io.Copy(dst, src) +func (p *Proxy) directPipe(src io.ReadCloser, dst io.WriteCloser, wait *sync.WaitGroup) { + defer func() { + src.Close() + dst.Close() + wait.Done() + }() + io.Copy(dst, src) } func NewProxy(conf *config.Config) *Proxy { diff --git a/wrappers/ctx.go b/wrappers/ctx.go deleted file mode 100644 index 0ea4af4..0000000 --- a/wrappers/ctx.go +++ /dev/null @@ -1,76 +0,0 @@ -package wrappers - -import ( - "context" - "net" - - "github.com/juju/errors" -) - -type Ctx struct { - cancel context.CancelFunc - conn StreamReadWriteCloser - ctx context.Context -} - -func (c *Ctx) Read(p []byte) (int, error) { - select { - case <-c.ctx.Done(): - return 0, errors.Annotate(c.ctx.Err(), "Read is failed because of closed context") - default: - n, err := c.conn.Read(p) - if err != nil { - c.cancel() - } - return n, err - } -} - -func (c *Ctx) Write(p []byte) (int, error) { - select { - case <-c.ctx.Done(): - return 0, errors.Annotate(c.ctx.Err(), "Write is failed because of closed context") - default: - n, err := c.conn.Write(p) - if err != nil { - c.cancel() - } - return n, err - } -} - -func (c *Ctx) LogDebug(msg string, data ...interface{}) { - c.conn.LogDebug(msg, data...) -} - -func (c *Ctx) LogInfo(msg string, data ...interface{}) { - c.conn.LogInfo(msg, data...) -} - -func (c *Ctx) LogWarn(msg string, data ...interface{}) { - c.conn.LogWarn(msg, data...) -} - -func (c *Ctx) LogError(msg string, data ...interface{}) { - c.conn.LogError(msg, data...) -} - -func (c *Ctx) LocalAddr() *net.TCPAddr { - return c.conn.LocalAddr() -} - -func (c *Ctx) RemoteAddr() *net.TCPAddr { - return c.conn.RemoteAddr() -} - -func (c *Ctx) Close() error { - return c.conn.Close() -} - -func NewCtx(ctx context.Context, cancel context.CancelFunc, conn StreamReadWriteCloser) StreamReadWriteCloser { - return &Ctx{ - ctx: ctx, - cancel: cancel, - conn: conn, - } -} diff --git a/wrappers/wrap.go b/wrappers/wrap.go index 05ba8ea..20cf7a0 100644 --- a/wrappers/wrap.go +++ b/wrappers/wrap.go @@ -71,7 +71,7 @@ type PacketReadWriter interface { PacketReader } -type BlockReadCloser interface { +type PacketReadCloser interface { io.Closer PacketReader }