wip

conntypes/packet.go

@@ -0,0 +1,3 @@
+package conntypes
+
+type Packet []byte

conntypes/wrappers.go

@@ -1,4 +1,4 @@
-package wrappers
+package conntypes
 
 import (
 	"io"
@@ -8,8 +8,6 @@ import (
 	"go.uber.org/zap"
 )
 
-type Packet []byte
-
 // Wrap is a base interface for all wrappers in this package.
 type Wrap interface {
 	Conn() net.Conn

go.mod

@@ -9,6 +9,8 @@ require (
 	github.com/allegro/bigcache v1.2.1
 	github.com/beevik/ntp v0.2.0
 	github.com/cespare/xxhash v1.1.0
+	github.com/dustin/go-humanize v1.0.0
+	github.com/gammazero/deque v0.0.0-20190521012701-46e4ffb7a622
 	github.com/juju/errors v0.0.0-20190806202954-0232dcc7464d
 	github.com/kr/pretty v0.1.0 // indirect
 	github.com/pkg/errors v0.8.1

go.sum

@@ -24,6 +24,10 @@ github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghf
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo=
+github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
+github.com/gammazero/deque v0.0.0-20190521012701-46e4ffb7a622 h1:lxbhOGZ9pU3Kf8P6lFluUcE82yVZn2EqEf4+mWRNPV0=
+github.com/gammazero/deque v0.0.0-20190521012701-46e4ffb7a622/go.mod h1:D90+MBHVc9Sk1lJAbEVgws0eYEurY4mv2TDso3Nxh3w=
 github.com/go-kit/kit v0.8.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as=
 github.com/go-logfmt/logfmt v0.3.0/go.mod h1:Qt1PoO58o5twSAckw1HlFXLmHsOX5/0LbT9GBnD5lWE=
 github.com/go-logfmt/logfmt v0.4.0/go.mod h1:3RMwSq7FuexP4Kalkev3ejPJsZTpXXBr9+V4qmtdjCk=

hub/closeable_channel.go

@@ -0,0 +1,59 @@
+package hub
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"github.com/9seconds/mtg/conntypes"
+)
+
+const closeableChannelReadTimeout = 2 * time.Minute
+
+type ChannelReadCloser interface {
+	Read() (conntypes.Packet, error)
+	Close()
+}
+
+type closeableChannel struct {
+	channel chan conntypes.Packet
+	ctx     context.Context
+	cancel  context.CancelFunc
+}
+
+func (c *closeableChannel) Read() (conntypes.Packet, error) {
+	timer := time.NewTimer(closeableChannelReadTimeout)
+	defer timer.Stop()
+
+	select {
+	case <-timer.C:
+		return nil, errors.New("timeout")
+	case <-c.ctx.Done():
+		return nil, errors.New("channel was closed")
+	case packet := <-c.channel:
+		return packet, nil
+	}
+}
+
+func (c *closeableChannel) write(packet conntypes.Packet) error {
+	select {
+	case <-c.ctx.Done():
+		return errors.New("channel was closed")
+	case c.channel <- packet:
+		return nil
+	}
+}
+
+func (c *closeableChannel) Close() {
+	c.cancel()
+	c.channel = nil
+}
+
+func newCloseableChannel(ctx context.Context) *closeableChannel {
+	ctx, cancel := context.WithCancel(ctx)
+	return &closeableChannel{
+		channel: make(chan conntypes.Packet),
+		ctx:     ctx,
+		cancel:  cancel,
+	}
+}

hub/connection.go

@@ -0,0 +1,100 @@
+package hub
+
+import (
+	"fmt"
+	"math/rand"
+	"sync"
+
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/protocol"
+)
+
+type connectionID int
+
+type connection struct {
+	conn    conntypes.PacketReadWriteCloser
+	mutex   sync.RWMutex
+	id      connectionID
+	hub     *connectionHub
+	pending uint
+	closing bool
+}
+
+func (c *connection) Write(packet conntypes.Packet) error {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	err := c.conn.Write(packet)
+	if err != nil {
+		// if we tried to write into a socket and it was broken, it is
+		// a time to reconsider the prescence of this socket at all.
+		//
+		// probably we need to remove it completely because it seems
+		// that connection is broken.
+		c.pending = 0
+	}
+	return err
+}
+
+func (c *connection) Read() (conntypes.Packet, error) {
+	packet, err := c.conn.Read()
+
+	c.mutex.Lock()
+	if err != nil {
+		c.pending--
+	} else {
+		c.pending = 0
+	}
+	c.mutex.Unlock()
+
+	return packet, err
+}
+
+func (c *connection) Stats() (bool, uint) {
+	c.mutex.RLock()
+	defer c.mutex.RUnlock()
+
+	return c.closing, c.pending
+}
+
+func (c *connection) Close() error {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	c.closing = true
+	return c.conn.Close()
+}
+
+func (c *connection) run() {
+	for {
+		packet, err := c.conn.Read()
+		if err != nil {
+			c.Close()
+			c.hub.brokenSocketsChan <- c.id
+			c.hub = nil
+			return
+		}
+
+		// TODO
+		if channel, ok := Registry.getChannel(conntypes.ConnID{}); ok {
+			go channel.write(packet) // nolint: errcheck
+		}
+	}
+}
+
+func newConnection(hub *connectionHub, req *protocol.TelegramRequest) (*connection, error) {
+	conn, err := mtproto.TelegramProtocol(req)
+	if err != nil {
+		return nil, fmt.Errorf("cannot create a new connection: %w", err)
+	}
+
+	rv := &connection{
+		conn: conn,
+		hub:  hub,
+		id:   connectionID(rand.Int()),
+	}
+	go rv.run()
+
+	return rv, nil
+}

hub/connection_hub.go

@@ -0,0 +1,84 @@
+package hub
+
+import "time"
+
+const hubGCEvery = time.Minute
+
+type connectionHub struct {
+	sockets map[connectionID]*connection
+
+	brokenSocketsChan      chan connectionID
+	connectionRequestsChan chan *connectionHubRequest
+	returnConnectionsChan  chan *connection
+}
+
+func (h *connectionHub) run() {
+	gcTicker := time.NewTicker(hubGCEvery)
+	defer gcTicker.Stop()
+
+	for {
+		select {
+		case <-gcTicker.C:
+			h.runGC()
+		case id := <-h.brokenSocketsChan:
+			h.runBrokenConnection(id)
+		case request := <-h.connectionRequestsChan:
+			h.runConnectionRequest(request)
+		case conn := <-h.returnConnectionsChan:
+			h.runReturnConnection(conn)
+		}
+	}
+}
+
+func (h *connectionHub) runBrokenConnection(id connectionID) {
+	delete(h.sockets, id)
+}
+
+func (h *connectionHub) runGC() {
+	for key, conn := range h.sockets {
+		closing, pending := conn.Stats()
+		switch {
+		case closing:
+			delete(h.sockets, key)
+		case pending == 0:
+			conn.Close()
+			delete(h.sockets, key)
+			return
+		}
+
+	}
+}
+
+func (h *connectionHub) runConnectionRequest(req *connectionHubRequest) {
+	for key, conn := range h.sockets {
+		closing, _ := conn.Stats()
+		delete(h.sockets, key)
+
+		if !closing {
+			req.responseChan <- conn
+			return
+		}
+	}
+
+	newConn, err := newConnection(h, req.req)
+	if err != nil {
+		close(req.responseChan)
+		return
+	}
+
+	req.responseChan <- newConn
+}
+
+func (h *connectionHub) runReturnConnection(conn *connection) {
+	h.sockets[conn.id] = conn
+}
+
+func newConnectionHub() *connectionHub {
+	return &connectionHub{
+		sockets: map[connectionID]*connection{},
+
+		brokenSocketsChan:      make(chan connectionID, 1),
+		connectionRequestsChan: make(chan *connectionHubRequest),
+		returnConnectionsChan:  make(chan *connection, 1),
+	}
+}

hub/connection_hub_request.go

@@ -0,0 +1,8 @@
+package hub
+
+import "github.com/9seconds/mtg/protocol"
+
+type connectionHubRequest struct {
+	req          *protocol.TelegramRequest
+	responseChan chan<- *connection
+}

hub/hub.go

@@ -0,0 +1,48 @@
+package hub
+
+import (
+	"errors"
+	"sync"
+
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/protocol"
+)
+
+type Concentrator struct {
+	hubs sync.Map
+}
+
+func (c *Concentrator) Write(packet conntypes.Packet, req *protocol.TelegramRequest) error {
+	hub := c.getHub(req)
+	connectionChan := make(chan *connection)
+	hub.connectionRequestsChan <- &connectionHubRequest{
+		req:          req,
+		responseChan: connectionChan,
+	}
+
+	conn, ok := <-connectionChan
+	if !ok {
+		return errors.New("cannot establish connection to telegram")
+	}
+}
+
+func (c *Concentrator) getHub(req *protocol.TelegramRequest) *connectionHub {
+	dcMapRaw, ok := c.hubs.Load(req.ClientProtocol.DC())
+	if !ok {
+		dcMapRaw, _ = c.hubs.LoadOrStore(req.ClientProtocol.DC(), &sync.Map{})
+	}
+	dcMap := dcMapRaw.(*sync.Map)
+
+	loaded := true
+	hubRaw, ok := dcMap.Load(req.ClientProtocol.ConnectionProtocol())
+	if !ok {
+		hubRaw, loaded = dcMap.LoadOrStore(req.ClientProtocol.ConnectionProtocol(),
+			newConnectionHub())
+	}
+	hub := hubRaw.(*connectionHub)
+	if !loaded {
+		go hub.run()
+	}
+
+	return hub
+}

hub/registry.go

@@ -0,0 +1,53 @@
+package hub
+
+import (
+	"context"
+	"sync"
+
+	"github.com/9seconds/mtg/conntypes"
+)
+
+var Registry *RegistryStruct
+
+type RegistryStruct struct {
+	conns map[string]*closeableChannel
+	ctx   context.Context
+	mutex sync.RWMutex
+}
+
+func (r *RegistryStruct) Register(id conntypes.ConnID) ChannelReadCloser {
+	channel := newCloseableChannel(r.ctx)
+
+	r.mutex.Lock()
+	r.conns[string(id[:])] = channel
+	r.mutex.Unlock()
+
+	return channel
+}
+
+func (r *RegistryStruct) Unregister(id conntypes.ConnID) {
+	r.mutex.Lock()
+	defer r.mutex.Unlock()
+
+	if channel, ok := r.conns[string(id[:])]; ok {
+		channel.Close()
+		delete(r.conns, string(id[:]))
+	}
+}
+
+func (r *RegistryStruct) getChannel(id conntypes.ConnID) (*closeableChannel, bool) {
+	r.mutex.RLock()
+	defer r.mutex.RUnlock()
+
+	if value, ok := r.conns[string(id[:])]; ok {
+		return value, true
+	}
+	return nil, false
+}
+
+func InitRegistry(ctx context.Context) {
+	Registry = &RegistryStruct{
+		ctx:   ctx,
+		conns: map[string]*closeableChannel{},
+	}
+}

mtproto/protocol.go

@@ -3,16 +3,15 @@ package mtproto
 import (
 	"fmt"
 
+	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/mtproto/rpc"
 	"github.com/9seconds/mtg/protocol"
 	"github.com/9seconds/mtg/telegram"
 	"github.com/9seconds/mtg/wrappers"
 )
 
-func TelegramProtocol(req *protocol.TelegramRequest) (wrappers.Wrap, error) {
+func TelegramProtocol(req *protocol.TelegramRequest) (conntypes.PacketReadWriteCloser, error) {
-	conn, err := telegram.Middle.Dial(req.Ctx,
+	conn, err := telegram.Middle.Dial(req.ClientProtocol.DC(),
-		req.Cancel,
-		req.ClientProtocol.DC(),
 		req.ClientProtocol.ConnectionProtocol())
 	if err != nil {
 		return nil, fmt.Errorf("cannot connect to telegram: %w", err)
@@ -42,7 +41,7 @@ func TelegramProtocol(req *protocol.TelegramRequest) (wrappers.Wrap, error) {
 	return frameConn, nil
 }
 
-func doRPCNonceRequest(conn wrappers.PacketWriter) (*rpc.NonceRequest, error) {
+func doRPCNonceRequest(conn conntypes.PacketWriter) (*rpc.NonceRequest, error) {
 	rpcNonceReq, err := rpc.NewNonceRequest(telegram.Middle.Secret())
 	if err != nil {
 		panic(err)
@@ -54,7 +53,7 @@ func doRPCNonceRequest(conn wrappers.PacketWriter) (*rpc.NonceRequest, error) {
 	return rpcNonceReq, nil
 }
 
-func getRPCNonceResponse(conn wrappers.PacketReader, req *rpc.NonceRequest) (*rpc.NonceResponse, error) {
+func getRPCNonceResponse(conn conntypes.PacketReader, req *rpc.NonceRequest) (*rpc.NonceResponse, error) {
 	packet, err := conn.Read()
 	if err != nil {
 		return nil, fmt.Errorf("cannot read from connection: %w", err)
@@ -71,14 +70,14 @@ func getRPCNonceResponse(conn wrappers.PacketReader, req *rpc.NonceRequest) (*rp
 	return resp, nil
 }
 
-func doRPCHandshakeRequest(conn wrappers.PacketWriter) error {
+func doRPCHandshakeRequest(conn conntypes.PacketWriter) error {
 	if err := conn.Write(rpc.HandshakeRequest); err != nil {
 		return fmt.Errorf("cannot make a request: %w", err)
 	}
 	return nil
 }
 
-func getRPCHandshakeResponse(conn wrappers.PacketReader) error {
+func getRPCHandshakeResponse(conn conntypes.PacketReader) error {
 	packet, err := conn.Read()
 	if err != nil {
 		return fmt.Errorf("cannot read a response: %w", err)

obfuscated2/client_protocol.go

@@ -37,7 +37,7 @@ func (c *ClientProtocol) DC() conntypes.DC {
 	return c.dc
 }
 
-func (c *ClientProtocol) Handshake(socket wrappers.StreamReadWriteCloser) (wrappers.StreamReadWriteCloser, error) {
+func (c *ClientProtocol) Handshake(socket conntypes.StreamReadWriteCloser) (conntypes.StreamReadWriteCloser, error) {
 	fm, err := c.ReadFrame(socket)
 	if err != nil {
 		return nil, fmt.Errorf("cannot make a client handshake: %w", err)
@@ -88,7 +88,7 @@ func (c *ClientProtocol) Handshake(socket wrappers.StreamReadWriteCloser) (wrapp
 	return wrappers.NewObfuscated2(socket, encryptor, decryptor), nil
 }
 
-func (c *ClientProtocol) ReadFrame(socket wrappers.StreamReader) (fm Frame, err error) {
+func (c *ClientProtocol) ReadFrame(socket conntypes.StreamReader) (fm Frame, err error) {
 	if _, err = io.ReadFull(handshakeReader{socket}, fm.Bytes()); err != nil {
 		err = fmt.Errorf("cannot extract obfuscated2 frame: %w", err)
 	}
@@ -96,7 +96,7 @@ func (c *ClientProtocol) ReadFrame(socket wrappers.StreamReader) (fm Frame, err
 }
 
 type handshakeReader struct {
-	parent wrappers.StreamReader
+	parent conntypes.StreamReader
 }
 
 func (h handshakeReader) Read(p []byte) (int, error) {

obfuscated2/telegram_protocol.go

@@ -4,20 +4,21 @@ import (
 	"crypto/rand"
 	"fmt"
 
+	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/protocol"
 	"github.com/9seconds/mtg/telegram"
 	"github.com/9seconds/mtg/utils"
 	"github.com/9seconds/mtg/wrappers"
 )
 
-func TelegramProtocol(req *protocol.TelegramRequest) (wrappers.Wrap, error) {
+func TelegramProtocol(req *protocol.TelegramRequest) (conntypes.StreamReadWriteCloser, error) {
-	socket, err := telegram.Direct.Dial(req.Ctx,
+	conn, err := telegram.Direct.Dial(req.ClientProtocol.DC(),
-		req.Cancel,
-		req.ClientProtocol.DC(),
 		req.ClientProtocol.ConnectionProtocol())
 	if err != nil {
 		return nil, fmt.Errorf("cannot dial to telegram: %w", err)
 	}
+	conn = wrappers.NewTimeout(conn)
+	conn = wrappers.NewCtx(req.Ctx, req.Cancel, conn)
 	fm := generateFrame(req.ClientProtocol)
 	data := fm.Bytes()
 
@@ -30,11 +31,11 @@ func TelegramProtocol(req *protocol.TelegramRequest) (wrappers.Wrap, error) {
 	encryptor.XORKeyStream(data, data)
 	copy(data[:frameOffsetIV], copyFrame[:frameOffsetIV])
 
-	if _, err := socket.Write(data); err != nil {
+	if _, err := conn.Write(data); err != nil {
 		return nil, fmt.Errorf("cannot write handshake frame to telegram: %w", err)
 	}
 
-	return wrappers.NewObfuscated2(socket, encryptor, decryptor), nil
+	return wrappers.NewObfuscated2(conn, encryptor, decryptor), nil
 }
 
 func generateFrame(cp protocol.ClientProtocol) (fm Frame) {

protocol/interfaces.go

@@ -1,16 +1,12 @@
 package protocol
 
-import (
+import "github.com/9seconds/mtg/conntypes"
-	"github.com/9seconds/mtg/conntypes"
-	"github.com/9seconds/mtg/wrappers"
-)
 
 type ClientProtocol interface {
-	Handshake(wrappers.StreamReadWriteCloser) (wrappers.StreamReadWriteCloser, error)
+	Handshake(conntypes.StreamReadWriteCloser) (conntypes.StreamReadWriteCloser, error)
 	ConnectionType() conntypes.ConnectionType
 	ConnectionProtocol() conntypes.ConnectionProtocol
 	DC() conntypes.DC
 }
 
-type TelegramProtocol func(*TelegramRequest) (wrappers.Wrap, error)
 type ClientProtocolMaker func() ClientProtocol

protocol/request.go

@@ -6,12 +6,11 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/9seconds/mtg/conntypes"
-	"github.com/9seconds/mtg/wrappers"
 )
 
 type TelegramRequest struct {
 	Logger         *zap.SugaredLogger
-	ClientConn     wrappers.StreamReadWriteCloser
+	ClientConn     conntypes.StreamReadWriteCloser
 	ConnID         conntypes.ConnID
 	Ctx            context.Context
 	Cancel         context.CancelFunc

proxy/proxy.go

@@ -63,25 +63,26 @@ func (p *Proxy) accept(conn net.Conn) {
 	ctx, cancel := context.WithCancel(p.Context)
 	defer cancel()
 
-	wrappedConn := wrappers.NewClientConn(ctx, cancel, conn, connID)
+	clientConn := wrappers.NewClientConn(conn, connID)
-	wrappedConn = wrappers.NewTraffic(wrappedConn)
+	clientConn = wrappers.NewCtx(ctx, cancel, clientConn)
-	defer wrappedConn.Close()
+	clientConn = wrappers.NewTimeout(clientConn)
+	clientConn = wrappers.NewTraffic(clientConn)
+	defer clientConn.Close()
 
 	clientProtocol := p.ClientProtocolMaker()
-	wrappedConn, err := clientProtocol.Handshake(wrappedConn)
+	clientConn, err := clientProtocol.Handshake(clientConn)
 	if err != nil {
 		logger.Warnw("Cannot perform client handshake", "error", err)
 		return
 	}
-	defer wrappedConn.Close()
 
-	stats.S.ClientConnected(clientProtocol.ConnectionType(), wrappedConn.RemoteAddr())
+	stats.S.ClientConnected(clientProtocol.ConnectionType(), clientConn.RemoteAddr())
-	defer stats.S.ClientDisconnected(clientProtocol.ConnectionType(), wrappedConn.RemoteAddr())
+	defer stats.S.ClientDisconnected(clientProtocol.ConnectionType(), clientConn.RemoteAddr())
 	logger.Infow("Client connected", "addr", conn.RemoteAddr())
 
 	req := &protocol.TelegramRequest{
 		Logger:         logger,
-		ClientConn:     wrappedConn,
+		ClientConn:     clientConn,
 		ConnID:         connID,
 		Ctx:            ctx,
 		Cancel:         cancel,
@@ -102,7 +103,7 @@ func (p *Proxy) acceptDirectConnection(request *protocol.TelegramRequest) error
 	if err != nil {
 		return err
 	}
-	telegramConn := telegramConnRaw.(wrappers.StreamReadWriteCloser)
+	telegramConn := telegramConnRaw.(conntypes.StreamReadWriteCloser)
 	defer telegramConn.Close()
 
 	wg := &sync.WaitGroup{}

telegram/base.go

@@ -1,7 +1,6 @@
 package telegram
 
 import (
-	"context"
 	"fmt"
 	"math/rand"
 	"net"
@@ -28,25 +27,8 @@ func (b *baseTelegram) Secret() []byte {
 	return b.secret
 }
 
-func (b *baseTelegram) dialToAddress(ctx context.Context,
+func (b *baseTelegram) dial(dc conntypes.DC,
-	cancel context.CancelFunc,
+	protocol conntypes.ConnectionProtocol) (conntypes.StreamReadWriteCloser, error) {
-	addr string) (wrappers.StreamReadWriteCloser, error) {
-	conn, err := b.dialer.Dial("tcp", addr)
-	if err != nil {
-		return nil, fmt.Errorf("dial has failed: %w", err)
-	}
-
-	if err := utils.InitTCP(conn); err != nil {
-		return nil, fmt.Errorf("cannot initialize tcp socket: %w", err)
-	}
-
-	return wrappers.NewTelegramConn(ctx, cancel, conn), nil
-}
-
-func (b *baseTelegram) dial(ctx context.Context,
-	cancel context.CancelFunc,
-	dc conntypes.DC,
-	protocol conntypes.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error) {
 	addr := ""
 
 	switch protocol {
@@ -56,7 +38,16 @@ func (b *baseTelegram) dial(ctx context.Context,
 		addr = b.chooseAddress(b.v6Addresses, dc, b.V6DefaultDC)
 	}
 
-	return b.dialToAddress(ctx, cancel, addr)
+	conn, err := b.dialer.Dial("tcp", addr)
+	if err != nil {
+		return nil, fmt.Errorf("dial has failed: %w", err)
+	}
+
+	if err := utils.InitTCP(conn); err != nil {
+		return nil, fmt.Errorf("cannot initialize tcp socket: %w", err)
+	}
+
+	return wrappers.NewTelegramConn(conn), nil
 }
 
 func (b *baseTelegram) chooseAddress(addresses map[conntypes.DC][]string,
@@ -66,7 +57,10 @@ func (b *baseTelegram) chooseAddress(addresses map[conntypes.DC][]string,
 		addrs, _ = addresses[defaultDC]
 	}
 
-	if len(addrs) > 0 {
+	switch {
+	case len(addrs) == 1:
+		return addrs[0]
+	case len(addrs) > 1:
 		return addrs[rand.Intn(len(addrs))]
 	}
 

telegram/direct.go

@@ -1,15 +1,11 @@
 package telegram
 
 import (
-	"context"
 	"net"
 
 	"github.com/9seconds/mtg/conntypes"
-	"github.com/9seconds/mtg/wrappers"
 )
 
-var Direct = newDirectTelegram()
-
 const (
 	directV4DefaultIdx conntypes.DC = 1
 	directV6DefaultIdx conntypes.DC = 1
@@ -36,10 +32,8 @@ type directTelegram struct {
 	baseTelegram
 }
 
-func (d *directTelegram) Dial(ctx context.Context,
+func (d *directTelegram) Dial(dc conntypes.DC,
-	cancel context.CancelFunc,
+	protocol conntypes.ConnectionProtocol) (conntypes.StreamReadWriteCloser, error) {
-	dc conntypes.DC,
-	protocol conntypes.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error) {
 	switch {
 	case dc < 0:
 		dc = -dc
@@ -47,11 +41,10 @@ func (d *directTelegram) Dial(ctx context.Context,
 		dc = conntypes.DCDefaultIdx
 	}
 
-	return d.baseTelegram.dial(ctx, cancel, dc-1, protocol)
+	return d.baseTelegram.dial(dc-1, protocol)
 }
 
-func newDirectTelegram() Telegram {
+var Direct = &directTelegram{
-	return &directTelegram{
 	baseTelegram: baseTelegram{
 		dialer:      net.Dialer{Timeout: telegramDialTimeout},
 		v4DefaultDC: directV4DefaultIdx,
@@ -59,5 +52,4 @@ func newDirectTelegram() Telegram {
 		v4Addresses: directV4Addresses,
 		v6Addresses: directV6Addresses,
 	},
-	}
 }

telegram/interfaces.go

@@ -1,16 +1,8 @@
 package telegram
 
-import (
+import  "github.com/9seconds/mtg/conntypes"
-	"context"
-
-	"github.com/9seconds/mtg/conntypes"
-	"github.com/9seconds/mtg/wrappers"
-)
 
 type Telegram interface {
-	Dial(context.Context,
+	Dial(conntypes.DC, conntypes.ConnectionProtocol) (conntypes.StreamReadWriteCloser, error)
-		context.CancelFunc,
-		conntypes.DC,
-		conntypes.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error)
 	Secret() []byte
 }

telegram/middle.go

@@ -1,7 +1,6 @@
 package telegram
 
 import (
-	"context"
 	"fmt"
 	"net"
 	"sync"
@@ -11,7 +10,6 @@ import (
 
 	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/telegram/api"
-	"github.com/9seconds/mtg/wrappers"
 )
 
 const middleTelegramBackgroundUpdateEvery = time.Hour
@@ -67,10 +65,8 @@ func (m *middleTelegram) backgroundUpdate() {
 	}
 }
 
-func (m *middleTelegram) Dial(ctx context.Context,
+func (m *middleTelegram) Dial(dc conntypes.DC,
-	cancel context.CancelFunc,
+	protocol conntypes.ConnectionProtocol) (conntypes.StreamReadWriteCloser, error) {
-	dc conntypes.DC,
-	protocol conntypes.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error) {
 	if dc == 0 {
 		dc = conntypes.DCDefaultIdx
 	}
@@ -78,7 +74,7 @@ func (m *middleTelegram) Dial(ctx context.Context,
 	m.mutex.RLock()
 	defer m.mutex.RUnlock()
 
-	return m.baseTelegram.dial(ctx, cancel, dc, protocol)
+	return m.baseTelegram.dial(dc, protocol)
 }
 
 func MiddleInit() {

wrappers/blockcipher.go

@@ -10,6 +10,8 @@ import (
 	"time"
 
 	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/conntypes"
 )
 
 const blockCipherReadCurrentDataBufferSize = 1024 + 1 // +1 because telegram operates with blocks mod 4
@@ -17,7 +19,7 @@ const blockCipherReadCurrentDataBufferSize = 1024 + 1 // +1 because telegram ope
 type wrapperBlockCipher struct {
 	buf bytes.Buffer
 
-	parent    StreamReadWriteCloser
+	parent    conntypes.StreamReadWriteCloser
 	encryptor cipher.BlockMode
 	decryptor cipher.BlockMode
 }
@@ -47,7 +49,8 @@ func (w *wrapperBlockCipher) ReadTimeout(p []byte, timeout time.Duration) (int,
 	return w.read(p, readAllTimeout(timeout))
 }
 
-func (w *wrapperBlockCipher) read(p []byte, reader func(StreamReadWriteCloser) ([]byte, error)) (int, error) {
+func (w *wrapperBlockCipher) read(p []byte,
+	reader func(conntypes.StreamReadWriteCloser) ([]byte, error)) (int, error) {
 	if w.buf.Len() > 0 {
 		return w.flush(p)
 	}
@@ -90,7 +93,7 @@ func (w *wrapperBlockCipher) encrypt(p []byte) ([]byte, error) {
 	return encrypted, nil
 }
 
-func readAll(src StreamReadWriteCloser) (rv []byte, err error) {
+func readAll(src conntypes.StreamReadWriteCloser) (rv []byte, err error) {
 	buf := make([]byte, blockCipherReadCurrentDataBufferSize)
 	n := blockCipherReadCurrentDataBufferSize
 
@@ -105,8 +108,8 @@ func readAll(src StreamReadWriteCloser) (rv []byte, err error) {
 	return rv, nil
 }
 
-func readAllTimeout(timeout time.Duration) func(StreamReadWriteCloser) ([]byte, error) {
+func readAllTimeout(timeout time.Duration) func(conntypes.StreamReadWriteCloser) ([]byte, error) {
-	return func(src StreamReadWriteCloser) (rv []byte, err error) {
+	return func(src conntypes.StreamReadWriteCloser) (rv []byte, err error) {
 		tmo := timeout
 		buf := make([]byte, blockCipherReadCurrentDataBufferSize)
 		n := blockCipherReadCurrentDataBufferSize
@@ -148,7 +151,8 @@ func (w *wrapperBlockCipher) RemoteAddr() *net.TCPAddr {
 	return w.parent.RemoteAddr()
 }
 
-func newBlockCipher(parent StreamReadWriteCloser, encryptor, decryptor cipher.BlockMode) StreamReadWriteCloser {
+func newBlockCipher(parent conntypes.StreamReadWriteCloser,
+	encryptor, decryptor cipher.BlockMode) conntypes.StreamReadWriteCloser {
 	return &wrapperBlockCipher{
 		parent:    parent,
 		encryptor: encryptor,

wrappers/conn.go

@@ -1,7 +1,6 @@
 package wrappers
 
 import (
-	"context"
 	"fmt"
 	"net"
 	"time"
@@ -19,15 +18,8 @@ const (
 	connPurposeTelegram
 )
 
-const (
-	connTimeoutRead  = 2 * time.Minute
-	connTimeoutWrite = 2 * time.Minute
-)
-
 type wrapperConn struct {
 	parent     net.Conn
-	ctx        context.Context
-	cancel     context.CancelFunc
 	connID     conntypes.ConnID
 	logger     *zap.SugaredLogger
 	localAddr  *net.TCPAddr
@@ -35,43 +27,34 @@ type wrapperConn struct {
 }
 
 func (w *wrapperConn) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
-	select {
-	case <-w.ctx.Done():
-		w.Close()
-		return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err())
-
-	default:
 	if err := w.parent.SetWriteDeadline(time.Now().Add(timeout)); err != nil {
-			w.Close() // nolint: gosec
+		w.Close()
 		return 0, fmt.Errorf("cannot set write deadline to the socket: %w", err)
 	}
 
+	return w.Write(p)
+}
+
+func (w *wrapperConn) Write(p []byte) (int, error) {
 	n, err := w.parent.Write(p)
-		w.logger.Debugw("Write to stream", "bytes", n, "error", err)
+	w.logger.Debugw("write to stream", "bytes", n, "error", err)
 	if err != nil {
 		w.Close() // nolint: gosec
 	}
 
 	return n, err
-	}
-}
-
-func (w *wrapperConn) Write(p []byte) (int, error) {
-	return w.WriteTimeout(p, connTimeoutWrite)
 }
 
 func (w *wrapperConn) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
-	select {
-	case <-w.ctx.Done():
-		w.Close()
-		return 0, fmt.Errorf("cannot read because context was closed: %w", w.ctx.Err())
-
-	default:
 	if err := w.parent.SetReadDeadline(time.Now().Add(timeout)); err != nil {
 		w.Close()
 		return 0, fmt.Errorf("cannot set read deadline to the socket: %w", err)
 	}
 
+	return w.Read(p)
+}
+
+func (w *wrapperConn) Read(p []byte) (int, error) {
 	n, err := w.parent.Read(p)
 	w.logger.Debugw("Read from stream", "bytes", n, "error", err)
 	if err != nil {
@@ -79,17 +62,10 @@ func (w *wrapperConn) ReadTimeout(p []byte, timeout time.Duration) (int, error)
 	}
 
 	return n, err
-	}
-}
-
-func (w *wrapperConn) Read(p []byte) (int, error) {
-	return w.ReadTimeout(p, connTimeoutRead)
 }
 
 func (w *wrapperConn) Close() error {
 	w.logger.Debugw("Close connection")
-	w.cancel()
-
 	return w.parent.Close()
 }
 
@@ -109,11 +85,9 @@ func (w *wrapperConn) RemoteAddr() *net.TCPAddr {
 	return w.remoteAddr
 }
 
-func newConn(ctx context.Context,
+func newConn(parent net.Conn,
-	cancel context.CancelFunc,
-	parent net.Conn,
 	connID conntypes.ConnID,
-	purpose connPurpose) StreamReadWriteCloser {
+	purpose connPurpose) conntypes.StreamReadWriteCloser {
 	localAddr := *parent.LocalAddr().(*net.TCPAddr)
 
 	if parent.RemoteAddr().(*net.TCPAddr).IP.To4() != nil {
@@ -135,8 +109,6 @@ func newConn(ctx context.Context,
 
 	return &wrapperConn{
 		parent:     parent,
-		ctx:        ctx,
-		cancel:     cancel,
 		connID:     connID,
 		logger:     logger,
 		remoteAddr: parent.RemoteAddr().(*net.TCPAddr),
@@ -144,15 +116,11 @@ func newConn(ctx context.Context,
 	}
 }
 
-func NewClientConn(ctx context.Context,
+func NewClientConn(parent net.Conn,
-	cancel context.CancelFunc,
+	connID conntypes.ConnID) conntypes.StreamReadWriteCloser {
-	parent net.Conn,
+	return newConn(parent, connID, connPurposeClient)
-	connID conntypes.ConnID) StreamReadWriteCloser {
-	return newConn(ctx, cancel, parent, connID, connPurposeClient)
 }
 
-func NewTelegramConn(ctx context.Context,
+func NewTelegramConn(parent net.Conn) conntypes.StreamReadWriteCloser {
-	cancel context.CancelFunc,
+	return newConn(parent, conntypes.ConnID{}, connPurposeTelegram)
-	parent net.Conn) StreamReadWriteCloser {
-	return newConn(ctx, cancel, parent, conntypes.ConnID{}, connPurposeTelegram)
 }

wrappers/ctx.go

@@ -0,0 +1,89 @@
+package wrappers
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"time"
+
+	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/conntypes"
+)
+
+type wrapperCtx struct {
+	parent conntypes.StreamReadWriteCloser
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func (w *wrapperCtx) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
+	select {
+	case <-w.ctx.Done():
+		w.Close()
+		return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err())
+	default:
+		return w.parent.WriteTimeout(p, timeout)
+	}
+}
+
+func (w *wrapperCtx) Write(p []byte) (int, error) {
+	select {
+	case <-w.ctx.Done():
+		w.Close()
+		return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err())
+	default:
+		return w.parent.Write(p)
+	}
+}
+
+func (w *wrapperCtx) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
+	select {
+	case <-w.ctx.Done():
+		w.Close()
+		return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err())
+	default:
+		return w.parent.ReadTimeout(p, timeout)
+	}
+}
+
+func (w *wrapperCtx) Read(p []byte) (int, error) {
+	select {
+	case <-w.ctx.Done():
+		w.Close()
+		return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err())
+	default:
+		return w.parent.Read(p)
+	}
+}
+
+func (w *wrapperCtx) Close() error {
+	w.cancel()
+	return w.parent.Close()
+}
+
+func (w *wrapperCtx) Conn() net.Conn {
+	return w.parent.Conn()
+}
+
+func (w *wrapperCtx) Logger() *zap.SugaredLogger {
+	return w.parent.Logger().Named("ctx")
+}
+
+func (w *wrapperCtx) LocalAddr() *net.TCPAddr {
+	return w.parent.LocalAddr()
+}
+
+func (w *wrapperCtx) RemoteAddr() *net.TCPAddr {
+	return w.parent.RemoteAddr()
+}
+
+func NewCtx(ctx context.Context,
+	cancel context.CancelFunc,
+	parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser {
+	return &wrapperCtx{
+		parent: parent,
+		ctx:    ctx,
+		cancel: cancel,
+	}
+}

wrappers/mtproto_cipher.go

@@ -9,6 +9,7 @@ import (
 	"encoding/binary"
 	"net"
 
+	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/mtproto/rpc"
 	"github.com/9seconds/mtg/utils"
 )
@@ -22,10 +23,10 @@ const (
 
 var mtprotoEmptyIP = [4]byte{0x00, 0x00, 0x00, 0x00}
 
-func NewMiddleProxyCipher(parent StreamReadWriteCloser,
+func NewMiddleProxyCipher(parent conntypes.StreamReadWriteCloser,
 	req *rpc.NonceRequest,
 	resp *rpc.NonceResponse,
-	secret []byte) StreamReadWriteCloser {
+	secret []byte) conntypes.StreamReadWriteCloser {
 	localAddr := parent.LocalAddr()
 	remoteAddr := parent.RemoteAddr()
 

wrappers/mtproto_frame.go

@@ -11,6 +11,8 @@ import (
 	"net"
 
 	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/conntypes"
 )
 
 const (
@@ -33,13 +35,13 @@ var mtprotoFramePadding = []byte{0x04, 0x00, 0x00, 0x00}
 // PADDING is custom padding schema to complete frame length to such that
 //    len(frame) % 16 == 0
 type wrapperMtprotoFrame struct {
-	parent     StreamReadWriteCloser
+	parent     conntypes.StreamReadWriteCloser
 	logger     *zap.SugaredLogger
 	readSeqNo  int32
 	writeSeqNo int32
 }
 
-func (w *wrapperMtprotoFrame) Read() (Packet, error) {
+func (w *wrapperMtprotoFrame) Read() (conntypes.Packet, error) {
 	buf := &bytes.Buffer{}
 	sum := crc32.NewIEEE()
 	writer := io.MultiWriter(buf, sum)
@@ -101,7 +103,7 @@ func (w *wrapperMtprotoFrame) Read() (Packet, error) {
 	return data, nil
 }
 
-func (w *wrapperMtprotoFrame) Write(p Packet) error {
+func (w *wrapperMtprotoFrame) Write(p conntypes.Packet) error {
 	messageLength := 4 + 4 + len(p) + 4
 	paddingLength := (aes.BlockSize - messageLength%aes.BlockSize) % aes.BlockSize
 
@@ -149,7 +151,7 @@ func (w *wrapperMtprotoFrame) RemoteAddr() *net.TCPAddr {
 	return w.parent.RemoteAddr()
 }
 
-func NewMtprotoFrame(parent StreamReadWriteCloser, seqNo int32) PacketReadWriteCloser {
+func NewMtprotoFrame(parent conntypes.StreamReadWriteCloser, seqNo int32) conntypes.PacketReadWriteCloser {
 	return &wrapperMtprotoFrame{
 		parent:     parent,
 		logger:     parent.Logger().Named("mtproto-frame"),

wrappers/obfuscated2.go

@@ -7,12 +7,14 @@ import (
 	"time"
 
 	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/conntypes"
 )
 
 type wrapperObfuscated2 struct {
 	encryptor cipher.Stream
 	decryptor cipher.Stream
-	parent    StreamReadWriteCloser
+	parent    conntypes.StreamReadWriteCloser
 }
 
 func (w *wrapperObfuscated2) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
@@ -71,7 +73,8 @@ func (w *wrapperObfuscated2) Close() error {
 	return w.parent.Close()
 }
 
-func NewObfuscated2(socket StreamReadWriteCloser, encryptor, decryptor cipher.Stream) StreamReadWriteCloser {
+func NewObfuscated2(socket conntypes.StreamReadWriteCloser,
+	encryptor, decryptor cipher.Stream) conntypes.StreamReadWriteCloser {
 	return &wrapperObfuscated2{
 		parent:    socket,
 		encryptor: encryptor,

wrappers/stats.go

@@ -6,11 +6,12 @@ import (
 
 	"go.uber.org/zap"
 
+	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/stats"
 )
 
 type wrapperStats struct {
-	parent StreamReadWriteCloser
+	parent conntypes.StreamReadWriteCloser
 }
 
 func (w *wrapperStats) Write(p []byte) (int, error) {
@@ -61,6 +62,6 @@ func (w *wrapperStats) Close() error {
 	return w.parent.Close()
 }
 
-func NewTraffic(parent StreamReadWriteCloser) StreamReadWriteCloser {
+func NewTraffic(parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser {
 	return &wrapperStats{parent}
 }

wrappers/timeout.go

@@ -0,0 +1,61 @@
+package wrappers
+
+import (
+	"net"
+	"time"
+
+	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/conntypes"
+)
+
+const (
+	timeoutRead  = 2 * time.Minute
+	timeoutWrite = 2 * time.Minute
+)
+
+type wrapperTimeout struct {
+	parent conntypes.StreamReadWriteCloser
+}
+
+func (w *wrapperTimeout) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
+	return w.parent.WriteTimeout(p, timeout)
+}
+
+func (w *wrapperTimeout) Write(p []byte) (int, error) {
+	return w.parent.WriteTimeout(p, timeoutWrite)
+}
+
+func (w *wrapperTimeout) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
+	return w.parent.ReadTimeout(p, timeout)
+}
+
+func (w *wrapperTimeout) Read(p []byte) (int, error) {
+	return w.parent.ReadTimeout(p, timeoutRead)
+}
+
+func (w *wrapperTimeout) Close() error {
+	return w.parent.Close()
+}
+
+func (w *wrapperTimeout) Conn() net.Conn {
+	return w.parent.Conn()
+}
+
+func (w *wrapperTimeout) Logger() *zap.SugaredLogger {
+	return w.parent.Logger().Named("timeout")
+}
+
+func (w *wrapperTimeout) LocalAddr() *net.TCPAddr {
+	return w.parent.LocalAddr()
+}
+
+func (w *wrapperTimeout) RemoteAddr() *net.TCPAddr {
+	return w.parent.RemoteAddr()
+}
+
+func NewTimeout(parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser {
+	return &wrapperTimeout{
+		parent: parent,
+	}
+}