Remove contexts

client/client.go

 package client
 
 import (
-	"context"
 	"net"
 
 	"github.com/9seconds/mtg/config"
 	"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)

client/direct.go

 package client
 
 import (
-	"context"
 	"net"
 	"time"
 
 
 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")
 	}
 	connOpts.ClientAddr = conn.RemoteAddr()
 
 	conn = wrappers.NewStreamCipher(conn, obfs2.Encryptor, obfs2.Decryptor)
-	conn = wrappers.NewCtx(ctx, cancel, conn)
 
 	return conn, connOpts, nil
 }

client/middle.go

 package client
 
 import (
-	"context"
 	"net"
 
 	"github.com/9seconds/mtg/config"
 	"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
 	}

proxy/proxy.go

 package proxy
 
 import (
-	"context"
 	"io"
 	"net"
 	"sync"
 
 	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
 		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 {
 	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
 	}
 }
 
-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 {

wrappers/ctx.go

-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,
-	}
-}

wrappers/wrap.go

 	PacketReader
 }
 
-type BlockReadCloser interface {
+type PacketReadCloser interface {
 	io.Closer
 	PacketReader
 }