Refactored all the things!

client/client.go

@@ -1,12 +0,0 @@
-package client
-
-import (
-	"net"
-
-	"github.com/9seconds/mtg/config"
-	"github.com/9seconds/mtg/mtproto"
-	"github.com/9seconds/mtg/wrappers"
-)
-
-// Init has to initialize client connection based on given config.
-type Init func(net.Conn, string, *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error)

client/direct.go

@@ -1,6 +1,7 @@
 package client
 
 import (
+	"context"
 	"net"
 	"time"
 
@@ -14,28 +15,29 @@ import (
 
 const handshakeTimeout = 10 * time.Second
 
-// DirectInit initializes client to access Telegram bypassing middleproxies.
-func DirectInit(conn net.Conn, socketID string, conf *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) {
-	if err := config.SetSocketOptions(conn); err != nil {
+func DirectInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn, connID string,
+	conf *config.Config) (wrappers.WrapStreamReadWriteCloser, *mtproto.ConnectionOpts, error) {
+	if err := config.SetSocketOptions(socket); err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot set socket options")
 	}
 
-	conn.SetReadDeadline(time.Now().Add(handshakeTimeout)) // nolint: errcheck
-	frame, err := obfuscated2.ExtractFrame(conn)
-	conn.SetReadDeadline(time.Time{}) // nolint: errcheck
+	socket.SetReadDeadline(time.Now().Add(handshakeTimeout))
+	frame, err := obfuscated2.ExtractFrame(socket)
 	if err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot extract frame")
 	}
+	socket.SetReadDeadline(time.Time{})
+	conn := wrappers.NewConn(socket, connID, wrappers.ConnPurposeClient, conf.PublicIPv4, conf.PublicIPv6)
 
 	obfs2, connOpts, err := obfuscated2.ParseObfuscated2ClientFrame(conf.Secret, frame)
 	if err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot parse obfuscated frame")
 	}
 	connOpts.ConnectionProto = mtproto.ConnectionProtocolAny
-	connOpts.ClientAddr = conn.RemoteAddr().(*net.TCPAddr)
+	connOpts.ClientAddr = conn.RemoteAddr()
 
-	socket := wrappers.NewTimeoutRWC(conn, socketID, conf.PublicIPv4, conf.PublicIPv6)
-	socket = wrappers.NewStreamCipherRWC(socket, obfs2.Encryptor, obfs2.Decryptor)
+	conn = wrappers.NewCtx(ctx, cancel, conn)
+	conn = wrappers.NewStreamCipher(conn, obfs2.Encryptor, obfs2.Decryptor)
 
-	return socket, connOpts, nil
+	return conn, connOpts, nil
 }

client/middle.go

@@ -1,30 +1,30 @@
 package client
 
 import (
+	"context"
 	"net"
 
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
-	mtwrappers "github.com/9seconds/mtg/mtproto/wrappers"
 	"github.com/9seconds/mtg/wrappers"
 )
 
-func MiddleInit(conn net.Conn, socketID string, conf *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) {
-	newConn, opts, err := DirectInit(conn, socketID, conf)
+func MiddleInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn, connID string,
+	conf *config.Config) (wrappers.WrapPacketReadWriteCloser, *mtproto.ConnectionOpts, error) {
+	conn, opts, err := DirectInit(ctx, cancel, socket, connID, conf)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	if opts.ConnectionType == mtproto.ConnectionTypeAbridged {
-		newConn = mtwrappers.NewAbridgedRWC(newConn, opts)
-	} else {
-		newConn = mtwrappers.NewIntermediateRWC(newConn, opts)
+	newConn := wrappers.NewMTProtoAbridged(conn, opts)
+	if opts.ConnectionType != mtproto.ConnectionTypeAbridged {
+		newConn = wrappers.NewMTProtoIntermediate(conn, opts)
 	}
 
 	opts.ConnectionProto = mtproto.ConnectionProtocolIPv4
-	if conn.LocalAddr().(*net.TCPAddr).IP.To4() == nil {
+	if socket.LocalAddr().(*net.TCPAddr).IP.To4() == nil {
 		opts.ConnectionProto = mtproto.ConnectionProtocolIPv6
 	}
 
-	return newConn, opts, nil
+	return newConn, opts, err
 }

config/config.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"net"
 	"strconv"
-	"time"
 
 	"github.com/juju/errors"
 )
@@ -14,10 +13,6 @@ import (
 const (
 	BufferWriteSize = 32 * 1024
 	BufferReadSize  = 32 * 1024
-	BufferSizeCopy  = 32 * 1024
-
-	TimeoutRead  = time.Minute
-	TimeoutWrite = time.Minute
 )
 
 // Config represents common configuration of mtg.

main.go

@@ -111,16 +111,21 @@ func main() {
 		zapcore.NewJSONEncoder(encoderCfg),
 		zapcore.Lock(os.Stderr),
 		atom,
-	)).Sugar()
+	))
+	zap.ReplaceGlobals(logger)
+	defer logger.Sync()
 
-	stat := proxy.NewStats(conf)
-	go stat.Serve()
+	var server *proxy.Proxy
+	if len(conf.AdTag) == 0 {
+		server = proxy.NewProxyDirect(conf)
+	} else {
+		server = proxy.NewProxyMiddle(conf)
+	}
 
-	srv := proxy.NewServer(conf, logger, stat)
 	printURLs(conf.GetURLs())
 
-	if err := srv.Serve(); err != nil {
-		logger.Fatal(err.Error())
+	if err := server.Serve(); err != nil {
+		zap.S().Fatalw("Server stopped", "error", err)
 	}
 }
 

obfuscated2/frame.go

@@ -9,6 +9,7 @@ import (
 	"github.com/juju/errors"
 
 	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/utils"
 )
 
 // [frameOffsetFirst:frameOffsetKey:frameOffsetIV:frameOffsetMagic:frameOffsetDC:frameOffsetEnd]
@@ -67,14 +68,7 @@ func (f Frame) ConnectionType() (mtproto.ConnectionType, error) {
 // Invert inverts frame for extracting encryption keys. Pkease check that link:
 // https://blog.susanka.eu/how-telegram-obfuscates-its-mtproto-traffic/
 func (f Frame) Invert() Frame {
-	reversed := make(Frame, FrameLen)
-	copy(reversed, f)
-
-	for i := 0; i < frameLenKey+frameLenIV; i++ {
-		reversed[frameOffsetFirst+i] = f[frameOffsetIV-1-i]
-	}
-
-	return reversed
+	return Frame(utils.ReverseBytes([]byte(f)))
 }
 
 // ExtractFrame extracts exact obfuscated2 handshake frame from given reader.

proxy/direct.go

@@ -0,0 +1,66 @@
+package proxy
+
+import (
+	"context"
+	"io"
+	"net"
+	"sync"
+
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/client"
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/telegram"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+func NewProxyDirect(conf *config.Config) *Proxy {
+	tg := telegram.NewDirectTelegram(conf)
+
+	return &Proxy{
+		conf: conf,
+		acceptCallback: func(ctx context.Context, cancel context.CancelFunc, clientSocket net.Conn,
+			connID string, wait *sync.WaitGroup, conf *config.Config) error {
+			client, opts, err := client.DirectInit(ctx, cancel, clientSocket, connID, conf)
+			if err != nil {
+				return errors.Annotate(err, "Cannot initialize client connection")
+			}
+			defer client.Close()
+
+			server, err := directTelegramStream(ctx, cancel, opts, connID, tg)
+			if err != nil {
+				return errors.Annotate(err, "Cannot initialize telegram connection")
+			}
+			defer server.Close()
+
+			wait.Add(2)
+
+			go directPipe(client, server, wait)
+			go directPipe(server, client, wait)
+
+			return nil
+		},
+	}
+}
+
+func directTelegramStream(ctx context.Context, cancel context.CancelFunc, opts *mtproto.ConnectionOpts,
+	connID string, tg *telegram.DirectTelegram) (wrappers.WrapStreamReadWriteCloser, error) {
+	streamConn, err := tg.Dial(connID, opts)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot dial to Telegram")
+	}
+	streamConn = wrappers.NewCtx(ctx, cancel, streamConn)
+
+	packetConn, err := tg.Init(opts, streamConn)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot handshake telegram")
+	}
+
+	return packetConn, nil
+}
+
+func directPipe(src io.Reader, dst io.Writer, wait *sync.WaitGroup) {
+	defer wait.Done()
+	io.Copy(dst, src)
+}

proxy/middle.go

@@ -0,0 +1,77 @@
+package proxy
+
+import (
+	"context"
+	"net"
+	"sync"
+
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/client"
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/telegram"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+func NewProxyMiddle(conf *config.Config) *Proxy {
+	tg := telegram.NewMiddleTelegram(conf)
+
+	return &Proxy{
+		conf: conf,
+		acceptCallback: func(ctx context.Context, cancel context.CancelFunc, clientSocket net.Conn,
+			connID string, wait *sync.WaitGroup, conf *config.Config) error {
+			client, opts, err := client.MiddleInit(ctx, cancel, clientSocket, connID, conf)
+			if err != nil {
+				return errors.Annotate(err, "Cannot initialize client connection")
+			}
+			defer client.Close()
+
+			server, err := middleTelegramStream(ctx, cancel, opts, connID, tg)
+			if err != nil {
+				return errors.Annotate(err, "Cannot initialize telegram connection")
+			}
+			defer server.Close()
+
+			wait.Add(2)
+
+			go middlePipe(client, server, wait, &opts.ReadHacks)
+			go middlePipe(server, client, wait, &opts.WriteHacks)
+
+			return nil
+		},
+	}
+}
+
+func middleTelegramStream(ctx context.Context, cancel context.CancelFunc, opts *mtproto.ConnectionOpts,
+	connID string, tg *telegram.MiddleTelegram) (wrappers.WrapPacketReadWriteCloser, error) {
+	streamConn, err := tg.Dial(connID, opts)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot dial to Telegram")
+	}
+	streamConn = wrappers.NewCtx(ctx, cancel, streamConn)
+
+	packetConn, err := tg.Init(opts, streamConn)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot handshake telegram")
+	}
+
+	return packetConn, nil
+}
+
+func middlePipe(src wrappers.WrapPacketReader, dst wrappers.WrapPacketWriter, wait *sync.WaitGroup, hacks *mtproto.Hacks) {
+	defer wait.Done()
+
+	for {
+		hacks.SimpleAck = false
+		hacks.QuickAck = false
+
+		packet, err := src.Read()
+		if err != nil {
+			return
+		}
+		if _, err = dst.Write(packet); err != nil {
+			return
+		}
+	}
+}

proxy/proxy.go

@@ -0,0 +1,63 @@
+package proxy
+
+import (
+	"context"
+	"net"
+	"sync"
+
+	"github.com/juju/errors"
+	uuid "github.com/satori/go.uuid"
+	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/config"
+)
+
+type proxyAcceptCallback func(context.Context, context.CancelFunc, net.Conn, string, *sync.WaitGroup, *config.Config) error
+
+type Proxy struct {
+	conf           *config.Config
+	acceptCallback proxyAcceptCallback
+}
+
+func (p *Proxy) Serve() error {
+	lsock, err := net.Listen("tcp", p.conf.BindAddr())
+	if err != nil {
+		return errors.Annotate(err, "Cannot create listen socket")
+	}
+
+	for {
+		if conn, err := lsock.Accept(); err != nil {
+			zap.S().Errorw("Cannot allocate incoming connection", "error", err)
+		} else {
+			go p.accept(conn)
+		}
+	}
+}
+
+func (p *Proxy) accept(conn net.Conn) {
+	connID := uuid.NewV4().String()
+	log := zap.S().With("connection_id", connID)
+
+	defer func() {
+		conn.Close()
+
+		if err := recover(); err != nil {
+			log.Errorw("Crash of accept handler", "error", err)
+		}
+	}()
+
+	log.Infow("Client connected", "addr", conn.RemoteAddr())
+
+	ctx, cancel := context.WithCancel(context.Background())
+	wait := &sync.WaitGroup{}
+
+	if err := p.acceptCallback(ctx, cancel, conn, connID, wait, p.conf); err != nil {
+		log.Errorw("Cannot initialize connection", "error", err)
+		cancel()
+	}
+
+	<-ctx.Done()
+	wait.Wait()
+
+	log.Infow("Client disconnected", "addr", conn.RemoteAddr())
+}

proxy/server.go

@@ -1,184 +0,0 @@
-package proxy
-
-import (
-	"context"
-	"io"
-	"net"
-	"sync"
-
-	"github.com/juju/errors"
-	uuid "github.com/satori/go.uuid"
-	"go.uber.org/zap"
-
-	"github.com/9seconds/mtg/client"
-	"github.com/9seconds/mtg/config"
-	"github.com/9seconds/mtg/mtproto"
-	"github.com/9seconds/mtg/telegram"
-	"github.com/9seconds/mtg/utils"
-	"github.com/9seconds/mtg/wrappers"
-)
-
-// Server is an insgtance of MTPROTO proxy.
-type Server struct {
-	conf       *config.Config
-	logger     *zap.SugaredLogger
-	stats      *Stats
-	tg         telegram.Telegram
-	clientInit client.Init
-}
-
-// Serve does MTPROTO proxying.
-func (s *Server) Serve() error {
-	lsock, err := net.Listen("tcp", s.conf.BindAddr())
-	if err != nil {
-		return errors.Annotate(err, "Cannot create listen socket")
-	}
-
-	for {
-		if conn, err := lsock.Accept(); err != nil {
-			s.logger.Warn("Cannot allocate incoming connection", "error", err)
-		} else {
-			go s.accept(conn)
-		}
-	}
-}
-
-func (s *Server) accept(conn net.Conn) {
-	defer func() {
-		s.stats.closeConnection()
-		conn.Close() // nolint: errcheck
-
-		if r := recover(); r != nil {
-			s.logger.Errorw("Crash of accept handler", "error", r)
-		}
-	}()
-
-	s.stats.newConnection()
-	ctx, cancel := context.WithCancel(context.Background())
-	socketID := uuid.NewV4().String()
-
-	s.logger.Debugw("Client connected",
-		"addr", conn.RemoteAddr().String(),
-		"socketid", socketID,
-	)
-
-	connOpts, clientConn, err := s.getClientStream(ctx, cancel, conn, socketID)
-	if err != nil {
-		s.logger.Warnw("Cannot initialize client connection",
-			"addr", conn.RemoteAddr().String(),
-			"socketid", socketID,
-			"error", err,
-		)
-		return
-	}
-	defer clientConn.Close() // nolint: errcheck
-
-	tgConn, err := s.getTelegramStream(ctx, cancel, connOpts, socketID)
-	if err != nil {
-		s.logger.Warnw("Cannot initialize Telegram connection",
-			"socketid", socketID,
-			"error", err,
-		)
-		return
-	}
-	defer tgConn.Close() // nolint: errcheck
-
-	wait := &sync.WaitGroup{}
-	wait.Add(2)
-
-	go func() {
-		defer wait.Done()
-
-		for {
-			connOpts.ReadHacks.QuickAck = false
-			connOpts.ReadHacks.SimpleAck = false
-			if err := s.pump(clientConn, tgConn, socketID, "client"); err != nil {
-				s.logger.Infow("Client stream is aborted",
-					"socketid", socketID, "error", err)
-				return
-			}
-		}
-	}()
-	go func() {
-		defer wait.Done()
-
-		for {
-			connOpts.WriteHacks.QuickAck = false
-			connOpts.WriteHacks.SimpleAck = false
-			if err := s.pump(tgConn, clientConn, socketID, "telegram"); err != nil {
-				s.logger.Infow("Telegram stream is aborted",
-					"socketid", socketID, "error", err)
-				return
-			}
-		}
-	}()
-
-	<-ctx.Done()
-	wait.Wait()
-
-	s.logger.Debugw("Client disconnected",
-		"addr", conn.RemoteAddr().String(),
-		"socketid", socketID,
-	)
-}
-
-func (s *Server) getClientStream(ctx context.Context, cancel context.CancelFunc, conn net.Conn, socketID string) (*mtproto.ConnectionOpts, io.ReadWriteCloser, error) {
-	socket, connOpts, err := s.clientInit(conn, socketID, s.conf)
-	if err != nil {
-		return nil, nil, errors.Annotate(err, "Cannot init client connection")
-	}
-
-	socket = wrappers.NewTrafficRWC(socket, s.stats.addIncomingTraffic, s.stats.addOutgoingTraffic)
-	socket = wrappers.NewLogRWC(socket, s.logger, socketID, "client")
-	socket = wrappers.NewCtxRWC(ctx, cancel, socket)
-
-	return connOpts, socket, nil
-}
-
-func (s *Server) getTelegramStream(ctx context.Context, cancel context.CancelFunc, connOpts *mtproto.ConnectionOpts, socketID string) (io.ReadWriteCloser, error) {
-	conn, err := s.tg.Dial(socketID, connOpts)
-	if err != nil {
-		return nil, errors.Annotate(err, "Cannot connect to Telegram")
-	}
-
-	conn = wrappers.NewTrafficRWC(conn, s.stats.addIncomingTraffic, s.stats.addOutgoingTraffic)
-	conn, err = s.tg.Init(connOpts, conn)
-	if err != nil {
-		return nil, errors.Annotate(err, "Cannot handshake Telegram")
-	}
-
-	conn = wrappers.NewLogRWC(conn, s.logger, socketID, "telegram")
-	conn = wrappers.NewCtxRWC(ctx, cancel, conn)
-
-	return conn, nil
-}
-
-func (s *Server) pump(src io.Reader, dst io.Writer, socketID, name string) error {
-	buf, err := utils.ReadCurrentData(src)
-	if err != nil {
-		return errors.Annotate(err, "Cannot pump the socket")
-	}
-
-	_, err = dst.Write(buf)
-
-	return err
-}
-
-// NewServer creates new instance of MTPROTO proxy.
-func NewServer(conf *config.Config, logger *zap.SugaredLogger, stat *Stats) *Server {
-	clientInit := client.DirectInit
-	tg := telegram.NewDirectTelegram
-
-	if len(conf.AdTag) > 0 {
-		clientInit = client.MiddleInit
-		tg = telegram.NewMiddleTelegram
-	}
-
-	return &Server{
-		conf:       conf,
-		logger:     logger,
-		stats:      stat,
-		tg:         tg(conf, logger),
-		clientInit: clientInit,
-	}
-}

proxy/stats.go

@@ -1,74 +0,0 @@
-package proxy
-
-import (
-	"encoding/json"
-	"net/http"
-	"strconv"
-	"sync/atomic"
-	"time"
-
-	"github.com/9seconds/mtg/config"
-)
-
-type statsUptime time.Time
-
-func (s statsUptime) MarshalJSON() ([]byte, error) {
-	uptime := int(time.Since(time.Time(s)).Seconds())
-	return []byte(strconv.Itoa(uptime)), nil
-}
-
-// Stats is a datastructure for statistics on work of this proxy.
-type Stats struct {
-	AllConnections    uint64 `json:"all_connections"`
-	ActiveConnections uint32 `json:"active_connections"`
-	Traffic           struct {
-		Incoming uint64 `json:"incoming"`
-		Outgoing uint64 `json:"outgoing"`
-	} `json:"traffic"`
-	URLs   config.IPURLs `json:"urls"`
-	Uptime statsUptime   `json:"uptime"`
-
-	conf *config.Config
-}
-
-func (s *Stats) newConnection() {
-	atomic.AddUint64(&s.AllConnections, 1)
-	atomic.AddUint32(&s.ActiveConnections, 1)
-}
-
-func (s *Stats) closeConnection() {
-	atomic.AddUint32(&s.ActiveConnections, ^uint32(0))
-}
-
-func (s *Stats) addIncomingTraffic(n int) {
-	atomic.AddUint64(&s.Traffic.Incoming, uint64(n))
-}
-
-func (s *Stats) addOutgoingTraffic(n int) {
-	atomic.AddUint64(&s.Traffic.Outgoing, uint64(n))
-}
-
-// Serve runs statistics HTTP server.
-func (s *Stats) Serve() {
-	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Content-Type", "application/json")
-
-		encoder := json.NewEncoder(w)
-		encoder.SetEscapeHTML(false)
-		encoder.SetIndent("", "  ")
-		encoder.Encode(s) // nolint: errcheck, gas
-	})
-
-	http.ListenAndServe(s.conf.StatAddr(), nil) // nolint: errcheck, gas
-}
-
-// NewStats returns new instance of statistics datastructure.
-func NewStats(conf *config.Config) *Stats {
-	stat := &Stats{
-		Uptime: statsUptime(time.Now()),
-		conf:   conf,
-	}
-	stat.URLs = conf.GetURLs()
-
-	return stat
-}

telegram/dialer.go

@@ -30,11 +30,14 @@ func (t *tgDialer) dial(addr string) (net.Conn, error) {
 	return conn, nil
 }
 
-func (t *tgDialer) dialRWC(addr, sock string) (wrappers.ReadWriteCloserWithAddr, error) {
+func (t *tgDialer) dialRWC(addr, connID string) (wrappers.WrapStreamReadWriteCloser, error) {
 	conn, err := t.dial(addr)
 	if err != nil {
 		return nil, err
 	}
 
-	return wrappers.NewTimeoutRWC(conn, sock, t.conf.PublicIPv4, t.conf.PublicIPv6), nil
+	tgConn := wrappers.NewConn(conn, connID, wrappers.ConnPurposeTelegram,
+		t.conf.PublicIPv4, t.conf.PublicIPv6)
+
+	return tgConn, nil
 }

telegram/direct.go

@@ -4,7 +4,6 @@ import (
 	"net"
 
 	"github.com/juju/errors"
-	"go.uber.org/zap"
 
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
@@ -29,11 +28,11 @@ var (
 	}
 )
 
-type directTelegram struct {
+type DirectTelegram struct {
 	baseTelegram
 }
 
-func (t *directTelegram) Dial(sock string, connOpts *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error) {
+func (t *DirectTelegram) Dial(connID string, connOpts *mtproto.ConnectionOpts) (wrappers.WrapStreamReadWriteCloser, error) {
 	dc := connOpts.DC
 	if dc < 0 {
 		dc = -dc
@@ -41,23 +40,23 @@ func (t *directTelegram) Dial(sock string, connOpts *mtproto.ConnectionOpts) (wr
 		dc = 1
 	}
 
-	return t.baseTelegram.dial(dc-1, sock, connOpts.ConnectionProto)
+	return t.baseTelegram.dial(dc-1, connID, connOpts.ConnectionProto)
 }
 
-func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.ReadWriteCloserWithAddr) (wrappers.ReadWriteCloserWithAddr, error) {
+func (t *DirectTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.WrapStreamReadWriteCloser) (wrappers.WrapStreamReadWriteCloser, error) {
 	obfs2, frame := obfuscated2.MakeTelegramObfuscated2Frame(connOpts)
 
-	if n, err := conn.Write(frame); err != nil || n != obfuscated2.FrameLen {
+	if _, err := conn.Write(frame); err != nil {
 		return nil, errors.Annotate(err, "Cannot write hadnshake frame")
 	}
 
-	return wrappers.NewStreamCipherRWC(conn, obfs2.Encryptor, obfs2.Decryptor), nil
+	return wrappers.NewStreamCipher(conn, obfs2.Encryptor, obfs2.Decryptor), nil
 }
 
 // NewDirectTelegram returns Telegram instance which connects directly
 // to Telegram bypassing middleproxies.
-func NewDirectTelegram(conf *config.Config, _ *zap.SugaredLogger) Telegram {
-	return &directTelegram{baseTelegram{
+func NewDirectTelegram(conf *config.Config) *DirectTelegram {
+	return &DirectTelegram{baseTelegram{
 		dialer: tgDialer{
 			Dialer: net.Dialer{Timeout: telegramDialTimeout},
 			conf:   conf,

telegram/middle.go

@@ -1,29 +1,26 @@
 package telegram
 
 import (
-	"io"
 	"net"
 	"net/http"
 	"sync"
 
 	"github.com/juju/errors"
-	"go.uber.org/zap"
 
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
 	"github.com/9seconds/mtg/mtproto/rpc"
-	mtwrappers "github.com/9seconds/mtg/mtproto/wrappers"
 	"github.com/9seconds/mtg/wrappers"
 )
 
-type middleTelegram struct {
+type MiddleTelegram struct {
 	middleTelegramCaller
 
 	conf *config.Config
 }
 
-func NewMiddleTelegram(conf *config.Config, logger *zap.SugaredLogger) Telegram {
-	tg := &middleTelegram{
+func NewMiddleTelegram(conf *config.Config) *MiddleTelegram {
+	tg := &MiddleTelegram{
 		middleTelegramCaller: middleTelegramCaller{
 			baseTelegram: baseTelegram{
 				dialer: tgDialer{
@@ -31,7 +28,6 @@ func NewMiddleTelegram(conf *config.Config, logger *zap.SugaredLogger) Telegram
 					conf:   conf,
 				},
 			},
-			logger: logger,
 			httpClient: &http.Client{
 				Timeout: middleTelegramHTTPClientTimeout,
 			},
@@ -48,8 +44,8 @@ func NewMiddleTelegram(conf *config.Config, logger *zap.SugaredLogger) Telegram
 	return tg
 }
 
-func (t *middleTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.ReadWriteCloserWithAddr) (wrappers.ReadWriteCloserWithAddr, error) {
-	rpcNonceConn := mtwrappers.NewFrameRWC(conn, rpc.SeqNoNonce)
+func (t *MiddleTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.WrapStreamReadWriteCloser) (wrappers.WrapPacketReadWriteCloser, error) {
+	rpcNonceConn := wrappers.NewMTProtoFrame(conn, rpc.SeqNoNonce)
 
 	rpcNonceReq, err := t.sendRPCNonceRequest(rpcNonceConn)
 	if err != nil {
@@ -60,22 +56,22 @@ func (t *middleTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.Re
 		return nil, err
 	}
 
-	secureConn := mtwrappers.NewMiddleProxyCipherRWC(conn, rpcNonceReq, rpcNonceResp, t.proxySecret)
-	secureConn = mtwrappers.NewFrameRWC(secureConn, rpc.SeqNoHandshake)
+	secureConn := wrappers.NewMiddleProxyCipher(conn, rpcNonceReq, rpcNonceResp, t.proxySecret)
+	frameConn := wrappers.NewMTProtoFrame(secureConn, rpc.SeqNoHandshake)
 
-	rpcHandshakeReq, err := t.sendRPCHandshakeRequest(secureConn)
+	rpcHandshakeReq, err := t.sendRPCHandshakeRequest(frameConn)
 	if err != nil {
 		return nil, err
 	}
-	_, err = t.receiveRPCHandshakeResponse(secureConn, rpcHandshakeReq)
+	_, err = t.receiveRPCHandshakeResponse(frameConn, rpcHandshakeReq)
 	if err != nil {
 		return nil, err
 	}
 
-	return mtwrappers.NewProxyRequestRWC(secureConn, connOpts, t.conf.AdTag)
+	return wrappers.NewMTProtoProxy(frameConn, connOpts, t.conf.AdTag)
 }
 
-func (t *middleTelegram) sendRPCNonceRequest(conn io.Writer) (*rpc.NonceRequest, error) {
+func (t *MiddleTelegram) sendRPCNonceRequest(conn wrappers.WrapPacketWriter) (*rpc.NonceRequest, error) {
 	rpcNonceReq, err := rpc.NewNonceRequest(t.proxySecret)
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot create RPC nonce request")
@@ -87,15 +83,13 @@ func (t *middleTelegram) sendRPCNonceRequest(conn io.Writer) (*rpc.NonceRequest,
 	return rpcNonceReq, nil
 }
 
-func (t *middleTelegram) receiveRPCNonceResponse(conn io.Reader, req *rpc.NonceRequest) (*rpc.NonceResponse, error) {
-	var ans [128]byte
-
-	n, err := conn.Read(ans[:])
+func (t *MiddleTelegram) receiveRPCNonceResponse(conn wrappers.WrapPacketReader, req *rpc.NonceRequest) (*rpc.NonceResponse, error) {
+	packet, err := conn.Read()
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot read RPC nonce response")
 	}
 
-	rpcNonceResp, err := rpc.NewNonceResponse(ans[:n])
+	rpcNonceResp, err := rpc.NewNonceResponse(packet)
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot initialize RPC nonce response")
 	}
@@ -106,7 +100,7 @@ func (t *middleTelegram) receiveRPCNonceResponse(conn io.Reader, req *rpc.NonceR
 	return rpcNonceResp, nil
 }
 
-func (t *middleTelegram) sendRPCHandshakeRequest(conn io.Writer) (*rpc.HandshakeRequest, error) {
+func (t *MiddleTelegram) sendRPCHandshakeRequest(conn wrappers.WrapPacketWriter) (*rpc.HandshakeRequest, error) {
 	req := rpc.NewHandshakeRequest()
 	if _, err := conn.Write(req.Bytes()); err != nil {
 		return nil, errors.Annotate(err, "Cannot send RPC handshake request")
@@ -115,15 +109,13 @@ func (t *middleTelegram) sendRPCHandshakeRequest(conn io.Writer) (*rpc.Handshake
 	return req, nil
 }
 
-func (t *middleTelegram) receiveRPCHandshakeResponse(conn io.Reader, req *rpc.HandshakeRequest) (*rpc.HandshakeResponse, error) {
-	var ans [128]byte
-
-	n, err := conn.Read(ans[:])
+func (t *MiddleTelegram) receiveRPCHandshakeResponse(conn wrappers.WrapPacketReader, req *rpc.HandshakeRequest) (*rpc.HandshakeResponse, error) {
+	packet, err := conn.Read()
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot read RPC handshake response")
 	}
 
-	rpcHandshakeResp, err := rpc.NewHandshakeResponse(ans[:n])
+	rpcHandshakeResp, err := rpc.NewHandshakeResponse(packet)
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot initialize RPC handshake response")
 	}

telegram/middle_caller.go

@@ -35,11 +35,10 @@ type middleTelegramCaller struct {
 
 	proxySecret []byte
 	dialerMutex *sync.RWMutex
-	logger      *zap.SugaredLogger
 	httpClient  *http.Client
 }
 
-func (t *middleTelegramCaller) Dial(sock string, connOpts *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error) {
+func (t *middleTelegramCaller) Dial(connID string, connOpts *mtproto.ConnectionOpts) (wrappers.WrapStreamReadWriteCloser, error) {
 	dc := connOpts.DC
 	if dc == 0 {
 		dc = 1
@@ -47,13 +46,13 @@ func (t *middleTelegramCaller) Dial(sock string, connOpts *mtproto.ConnectionOpt
 	t.dialerMutex.RLock()
 	defer t.dialerMutex.RUnlock()
 
-	return t.baseTelegram.dial(dc, sock, connOpts.ConnectionProto)
+	return t.baseTelegram.dial(dc, connID, connOpts.ConnectionProto)
 }
 
 func (t *middleTelegramCaller) autoUpdate() {
 	for range time.Tick(middleTelegramAutoUpdateInterval) {
 		if err := t.update(); err != nil {
-			t.logger.Warnw("Cannot update from Telegram", "error", err)
+			zap.S().Warnw("Cannot update from Telegram", "error", err)
 		}
 	}
 }
@@ -80,7 +79,7 @@ func (t *middleTelegramCaller) update() error {
 	t.v6Addresses = v6Addresses
 	t.dialerMutex.Unlock()
 
-	t.logger.Infow("Telegram middle proxy data has been updated")
+	zap.S().Infow("Telegram middle proxy data has been updated")
 
 	return nil
 }

telegram/telegram.go

@@ -9,14 +9,6 @@ import (
 	"github.com/9seconds/mtg/wrappers"
 )
 
-// Telegram defines an interface to connect to Telegram. This
-// encapsulates logic of working with middleproxies or direct
-// connections.
-type Telegram interface {
-	Dial(string, *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error)
-	Init(*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr) (wrappers.ReadWriteCloserWithAddr, error)
-}
-
 type baseTelegram struct {
 	dialer tgDialer
 
@@ -24,7 +16,7 @@ type baseTelegram struct {
 	v6Addresses map[int16][]string
 }
 
-func (b *baseTelegram) dial(dcIdx int16, sock string, proto mtproto.ConnectionProtocol) (wrappers.ReadWriteCloserWithAddr, error) {
+func (b *baseTelegram) dial(dcIdx int16, connID string, proto mtproto.ConnectionProtocol) (wrappers.WrapStreamReadWriteCloser, error) {
 	addrs := make([]string, 2)
 
 	if proto&mtproto.ConnectionProtocolIPv6 != 0 {
@@ -39,7 +31,7 @@ func (b *baseTelegram) dial(dcIdx int16, sock string, proto mtproto.ConnectionPr
 	}
 
 	for _, addr := range addrs {
-		if conn, err := b.dialer.dialRWC(addr, sock); err == nil {
+		if conn, err := b.dialer.dialRWC(addr, connID); err == nil {
 			return conn, err
 		}
 	}

wrappers/blockcipher.go

@@ -1,6 +1,7 @@
 package wrappers
 
 import (
+	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
 	"net"
@@ -9,77 +10,89 @@ import (
 	"github.com/juju/errors"
 )
 
-type WrapBlockCipher struct {
-	BufferedReader
+type BlockCipher struct {
+	buf *bytes.Buffer
 
 	conn      WrapStreamReadWriteCloser
 	encryptor cipher.BlockMode
 	decryptor cipher.BlockMode
 }
 
-func (w *WrapBlockCipher) Read(p []byte) (int, error) {
-	return w.BufferedRead(p, func() error {
-		var buf []byte
+func (b *BlockCipher) Read(p []byte) (int, error) {
+	if b.buf.Len() > 0 {
+		return b.flush(p)
+	}
 
-		for len(buf) == 0 || len(buf)%aes.BlockSize != 0 {
-			rv, err := utils.ReadCurrentData(w.conn)
-			if err != nil {
-				return errors.Annotate(err, "Cannot read from socket")
-			}
-			buf = append(buf, rv...)
+	buf := []byte{}
+	for len(buf) == 0 || len(buf)%aes.BlockSize != 0 {
+		rv, err := utils.ReadCurrentData(b.conn)
+		if err != nil {
+			return 0, errors.Annotate(err, "Cannot read from socket")
 		}
+		buf = append(buf, rv...)
+	}
+
+	b.decryptor.CryptBlocks(buf, buf)
+	b.buf.Write(buf)
 
-		w.decryptor.CryptBlocks(buf, buf)
-		w.Buffer.Write(buf)
+	return b.flush(p)
+}
+
+func (b *BlockCipher) flush(p []byte) (int, error) {
+	if b.buf.Len() <= len(p) {
+		sizeToReturn := b.buf.Len()
+		copy(p, b.buf.Bytes())
+		b.buf.Reset()
+		return sizeToReturn, nil
+	}
 
-		return nil
-	})
+	return b.buf.Read(p)
 }
 
-func (w *WrapBlockCipher) Write(p []byte) (int, error) {
+func (b *BlockCipher) Write(p []byte) (int, error) {
 	if len(p)%aes.BlockSize > 0 {
 		return 0, errors.Errorf("Incorrect block size %d", len(p))
 	}
 
 	encrypted := make([]byte, len(p))
-	w.encryptor.CryptBlocks(encrypted, p)
+	b.encryptor.CryptBlocks(encrypted, p)
 
-	return w.conn.Write(encrypted)
+	return b.conn.Write(encrypted)
 }
 
-func (w *WrapBlockCipher) LogDebug(msg string, data ...interface{}) {
-	w.conn.LogDebug(msg, data...)
+func (b *BlockCipher) LogDebug(msg string, data ...interface{}) {
+	b.conn.LogDebug(msg, data...)
 }
 
-func (w *WrapBlockCipher) LogInfo(msg string, data ...interface{}) {
-	w.conn.LogInfo(msg, data...)
+func (b *BlockCipher) LogInfo(msg string, data ...interface{}) {
+	b.conn.LogInfo(msg, data...)
 }
 
-func (w *WrapBlockCipher) LogWarn(msg string, data ...interface{}) {
-	w.conn.LogWarn(msg, data...)
+func (b *BlockCipher) LogWarn(msg string, data ...interface{}) {
+	b.conn.LogWarn(msg, data...)
 }
 
-func (w *WrapBlockCipher) LogError(msg string, data ...interface{}) {
-	w.conn.LogError(msg, data...)
+func (b *BlockCipher) LogError(msg string, data ...interface{}) {
+	b.conn.LogError(msg, data...)
 }
 
-func (w *WrapBlockCipher) LocalAddr() *net.TCPAddr {
-	return w.conn.LocalAddr()
+func (b *BlockCipher) LocalAddr() *net.TCPAddr {
+	return b.conn.LocalAddr()
 }
 
-func (w *WrapBlockCipher) RemoteAddr() *net.TCPAddr {
-	return w.conn.RemoteAddr()
+func (b *BlockCipher) RemoteAddr() *net.TCPAddr {
+	return b.conn.RemoteAddr()
 }
 
-func (w *WrapBlockCipher) Close() error {
-	return w.conn.Close()
+func (b *BlockCipher) Close() error {
+	return b.conn.Close()
 }
 
-func NewWrapBlockCipher(conn WrapStreamReadWriteCloser, encryptor, decryptor cipher.BlockMode) WrapStreamReadWriteCloser {
-	return &WrapBlockCipher{
-		BufferedReader: NewBufferedReader(),
-		conn:           conn,
-		encryptor:      encryptor,
-		decryptor:      decryptor,
+func NewBlockCipher(conn WrapStreamReadWriteCloser, encryptor, decryptor cipher.BlockMode) WrapStreamReadWriteCloser {
+	return &BlockCipher{
+		buf:       &bytes.Buffer{},
+		conn:      conn,
+		encryptor: encryptor,
+		decryptor: decryptor,
 	}
 }

wrappers/buffered_reader.go

@@ -1,32 +0,0 @@
-package wrappers
-
-import "bytes"
-
-type BufferedReader struct {
-	Buffer *bytes.Buffer
-}
-
-func (b *BufferedReader) BufferedRead(p []byte, callback func() error) (int, error) {
-	if b.Buffer.Len() > 0 {
-		return b.flush(p)
-	}
-	if err := callback(); err != nil {
-		return 0, err
-	}
-	return b.flush(p)
-}
-
-func (b *BufferedReader) flush(p []byte) (int, error) {
-	if b.Buffer.Len() <= len(p) {
-		sizeToReturn := b.Buffer.Len()
-		copy(p, b.Buffer.Bytes())
-		b.Buffer.Reset()
-		return sizeToReturn, nil
-	}
-
-	return b.Buffer.Read(p)
-}
-
-func NewBufferedReader() BufferedReader {
-	return BufferedReader{Buffer: &bytes.Buffer{}}
-}

wrappers/conn.go

@@ -30,8 +30,7 @@ const (
 	connTimeoutWrite = 5 * time.Minute
 )
 
-type WrapConn struct {
-	purpose    ConnPurpose
+type Conn struct {
 	connID     string
 	conn       net.Conn
 	logger     *zap.SugaredLogger
@@ -39,77 +38,80 @@ type WrapConn struct {
 	publicIPv6 net.IP
 }
 
-func (w *WrapConn) Write(p []byte) (int, error) {
-	w.conn.SetWriteDeadline(time.Now().Add(connTimeoutWrite))
-	n, err := w.conn.Write(p)
+func (c *Conn) Write(p []byte) (int, error) {
+	c.conn.SetWriteDeadline(time.Now().Add(connTimeoutWrite))
+	n, err := c.conn.Write(p)
 
-	w.logger.Debugw("Write to stream", "bytes", n, "error", err)
+	c.logger.Debugw("Write to stream", "bytes", n, "error", err)
 
 	return n, err
 }
 
-func (w *WrapConn) Read(p []byte) (int, error) {
-	w.conn.SetReadDeadline(time.Now().Add(connTimeoutRead))
-	n, err := w.conn.Read(p)
+func (c *Conn) Read(p []byte) (int, error) {
+	c.conn.SetReadDeadline(time.Now().Add(connTimeoutRead))
+	n, err := c.conn.Read(p)
 
-	w.logger.Debugw("Read from stream", "bytes", n, "error", err)
+	c.logger.Debugw("Read from stream", "bytes", n, "error", err)
 
 	return n, err
 }
 
-func (w *WrapConn) Close() error {
-	defer w.LogDebug("Closed connection")
-	return w.conn.Close()
+func (c *Conn) Close() error {
+	defer c.LogDebug("Closed connection")
+	return c.conn.Close()
 }
 
-func (w *WrapConn) LocalAddr() *net.TCPAddr {
-	addr := w.conn.LocalAddr().(*net.TCPAddr)
+func (c *Conn) LocalAddr() *net.TCPAddr {
+	addr := c.conn.LocalAddr().(*net.TCPAddr)
 	newAddr := *addr
 
-	if w.RemoteAddr().IP.To4() != nil {
-		if w.publicIPv4 != nil {
-			newAddr.IP = w.publicIPv4
+	if c.RemoteAddr().IP.To4() != nil {
+		if c.publicIPv4 != nil {
+			newAddr.IP = c.publicIPv4
 		}
-	} else if w.publicIPv6 != nil {
-		newAddr.IP = w.publicIPv6
+	} else if c.publicIPv6 != nil {
+		newAddr.IP = c.publicIPv6
 	}
 
 	return &newAddr
 }
 
-func (w *WrapConn) RemoteAddr() *net.TCPAddr {
-	return w.conn.RemoteAddr().(*net.TCPAddr)
+func (c *Conn) RemoteAddr() *net.TCPAddr {
+	return c.conn.RemoteAddr().(*net.TCPAddr)
 }
 
-func (w *WrapConn) LogDebug(msg string, data ...interface{}) {
-	w.logger.Debugw(msg, data...)
+func (c *Conn) LogDebug(msg string, data ...interface{}) {
+	c.logger.Debugw(msg, data...)
 }
 
-func (w *WrapConn) LogInfo(msg string, data ...interface{}) {
-	w.logger.Infow(msg, data...)
+func (c *Conn) LogInfo(msg string, data ...interface{}) {
+	c.logger.Infow(msg, data...)
 }
 
-func (w *WrapConn) LogWarn(msg string, data ...interface{}) {
-	w.logger.Warnw(msg, data...)
+func (c *Conn) LogWarn(msg string, data ...interface{}) {
+	c.logger.Warnw(msg, data...)
 }
 
-func (w *WrapConn) LogError(msg string, data ...interface{}) {
-	w.logger.Errorw(msg, data...)
+func (c *Conn) LogError(msg string, data ...interface{}) {
+	c.logger.Errorw(msg, data...)
 }
 
-func NewConn(connID string, purpose ConnPurpose, conn net.Conn, publicIPv4, publicIPv6 net.IP) WrapStreamReadWriteCloser {
+func NewConn(conn net.Conn, connID string, purpose ConnPurpose, publicIPv4, publicIPv6 net.IP) WrapStreamReadWriteCloser {
 	logger := zap.S().With(
 		"connection_id", connID,
 		"local_address", conn.LocalAddr(),
 		"remote_address", conn.RemoteAddr(),
+		"purpose", purpose,
 	)
 
-	return &WrapConn{
+	wrapper := Conn{
 		logger:     logger,
-		purpose:    purpose,
 		connID:     connID,
 		conn:       conn,
 		publicIPv4: publicIPv4,
 		publicIPv6: publicIPv6,
 	}
+	wrapper.logger = logger.With("faked_local_addr", wrapper.LocalAddr())
+
+	return &wrapper
 }

wrappers/ctx.go

@@ -7,68 +7,68 @@ import (
 	"github.com/juju/errors"
 )
 
-type WrapCtx struct {
+type Ctx struct {
 	cancel context.CancelFunc
 	conn   WrapStreamReadWriteCloser
 	ctx    context.Context
 }
 
-func (w *WrapCtx) Read(p []byte) (int, error) {
+func (c *Ctx) Read(p []byte) (int, error) {
 	select {
-	case <-w.ctx.Done():
-		return 0, errors.Annotate(w.ctx.Err(), "Read is failed because of closed context")
+	case <-c.ctx.Done():
+		return 0, errors.Annotate(c.ctx.Err(), "Read is failed because of closed context")
 	default:
-		n, err := w.conn.Read(p)
+		n, err := c.conn.Read(p)
 		if err != nil {
-			w.cancel()
+			c.cancel()
 		}
 		return n, err
 	}
 }
 
-func (w *WrapCtx) Write(p []byte) (int, error) {
+func (c *Ctx) Write(p []byte) (int, error) {
 	select {
-	case <-w.ctx.Done():
-		return 0, errors.Annotate(w.ctx.Err(), "Write is failed because of closed context")
+	case <-c.ctx.Done():
+		return 0, errors.Annotate(c.ctx.Err(), "Write is failed because of closed context")
 	default:
-		n, err := w.conn.Write(p)
+		n, err := c.conn.Write(p)
 		if err != nil {
-			w.cancel()
+			c.cancel()
 		}
 		return n, err
 	}
 }
 
-func (w *WrapCtx) LogDebug(msg string, data ...interface{}) {
-	w.conn.LogDebug(msg, data...)
+func (c *Ctx) LogDebug(msg string, data ...interface{}) {
+	c.conn.LogDebug(msg, data...)
 }
 
-func (w *WrapCtx) LogInfo(msg string, data ...interface{}) {
-	w.conn.LogInfo(msg, data...)
+func (c *Ctx) LogInfo(msg string, data ...interface{}) {
+	c.conn.LogInfo(msg, data...)
 }
 
-func (w *WrapCtx) LogWarn(msg string, data ...interface{}) {
-	w.conn.LogWarn(msg, data...)
+func (c *Ctx) LogWarn(msg string, data ...interface{}) {
+	c.conn.LogWarn(msg, data...)
 }
 
-func (w *WrapCtx) LogError(msg string, data ...interface{}) {
-	w.conn.LogError(msg, data...)
+func (c *Ctx) LogError(msg string, data ...interface{}) {
+	c.conn.LogError(msg, data...)
 }
 
-func (w *WrapCtx) LocalAddr() *net.TCPAddr {
-	return w.conn.LocalAddr()
+func (c *Ctx) LocalAddr() *net.TCPAddr {
+	return c.conn.LocalAddr()
 }
 
-func (w *WrapCtx) RemoteAddr() *net.TCPAddr {
-	return w.conn.RemoteAddr()
+func (c *Ctx) RemoteAddr() *net.TCPAddr {
+	return c.conn.RemoteAddr()
 }
 
-func (w *WrapCtx) Close() error {
-	return w.conn.Close()
+func (c *Ctx) Close() error {
+	return c.conn.Close()
 }
 
 func NewCtx(ctx context.Context, cancel context.CancelFunc, conn WrapStreamReadWriteCloser) WrapStreamReadWriteCloser {
-	return &WrapCtx{
+	return &Ctx{
 		ctx:    ctx,
 		cancel: cancel,
 		conn:   conn,

wrappers/mtproto_cipher.go

@@ -32,7 +32,7 @@ func NewMiddleProxyCipher(conn WrapStreamReadWriteCloser, req *rpc.NonceRequest,
 	enc, _ := makeEncrypterDecrypter(encKey, encIV)
 	_, dec := makeEncrypterDecrypter(decKey, decIV)
 
-	return NewWrapBlockCipher(conn, enc, dec)
+	return NewBlockCipher(conn, enc, dec)
 }
 
 func deriveKeys(purpose CipherPurpose, req *rpc.NonceRequest, resp *rpc.NonceResponse, client *net.TCPAddr, remote *net.TCPAddr, secret []byte) ([]byte, []byte) {

wrappers/mtproto_proxy.go

@@ -4,8 +4,10 @@ import (
 	"bytes"
 	"net"
 
-	"github.com/9seconds/mtg/mtproto/rpc"
 	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/mtproto/rpc"
 )
 
 type MTProtoProxy struct {
@@ -127,9 +129,14 @@ func (m *MTProtoProxy) Close() error {
 	return m.conn.Close()
 }
 
-func NewMTProtoProxy(conn WrapPacketReadWriteCloser, req *rpc.ProxyRequest) WrapPacketReadWriteCloser {
+func NewMTProtoProxy(conn WrapPacketReadWriteCloser, connOpts *mtproto.ConnectionOpts, adTag []byte) (WrapPacketReadWriteCloser, error) {
+	req, err := rpc.NewProxyRequest(connOpts.ClientAddr, conn.LocalAddr(), connOpts, adTag)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot create new RPC proxy request")
+	}
+
 	return &MTProtoProxy{
 		conn: conn,
 		req:  req,
-	}
+	}, nil
 }

wrappers/streamcipher.go

@@ -7,59 +7,59 @@ import (
 	"github.com/juju/errors"
 )
 
-type WrapStreamCipher struct {
+type StreamCipher struct {
 	encryptor cipher.Stream
 	decryptor cipher.Stream
 	conn      WrapStreamReadWriteCloser
 }
 
-func (w *WrapStreamCipher) Read(p []byte) (int, error) {
-	n, err := w.conn.Read(p)
+func (s *StreamCipher) Read(p []byte) (int, error) {
+	n, err := s.conn.Read(p)
 	if err != nil {
 		return 0, errors.Annotate(err, "Cannot read stream ciphered data")
 	}
-	w.decryptor.XORKeyStream(p, p[:n])
+	s.decryptor.XORKeyStream(p, p[:n])
 
 	return n, nil
 }
 
-func (w *WrapStreamCipher) Write(p []byte) (int, error) {
+func (s *StreamCipher) Write(p []byte) (int, error) {
 	encrypted := make([]byte, len(p))
-	w.encryptor.XORKeyStream(encrypted, p)
+	s.encryptor.XORKeyStream(encrypted, p)
 
-	return w.conn.Write(encrypted)
+	return s.conn.Write(encrypted)
 }
 
-func (w *WrapStreamCipher) LogDebug(msg string, data ...interface{}) {
-	w.conn.LogDebug(msg, data...)
+func (s *StreamCipher) LogDebug(msg string, data ...interface{}) {
+	s.conn.LogDebug(msg, data...)
 }
 
-func (w *WrapStreamCipher) LogInfo(msg string, data ...interface{}) {
-	w.conn.LogInfo(msg, data...)
+func (s *StreamCipher) LogInfo(msg string, data ...interface{}) {
+	s.conn.LogInfo(msg, data...)
 }
 
-func (w *WrapStreamCipher) LogWarn(msg string, data ...interface{}) {
-	w.conn.LogWarn(msg, data...)
+func (s *StreamCipher) LogWarn(msg string, data ...interface{}) {
+	s.conn.LogWarn(msg, data...)
 }
 
-func (w *WrapStreamCipher) LogError(msg string, data ...interface{}) {
-	w.conn.LogError(msg, data...)
+func (s *StreamCipher) LogError(msg string, data ...interface{}) {
+	s.conn.LogError(msg, data...)
 }
 
-func (w *WrapStreamCipher) LocalAddr() *net.TCPAddr {
-	return w.conn.LocalAddr()
+func (s *StreamCipher) LocalAddr() *net.TCPAddr {
+	return s.conn.LocalAddr()
 }
 
-func (w *WrapStreamCipher) RemoteAddr() *net.TCPAddr {
-	return w.conn.RemoteAddr()
+func (s *StreamCipher) RemoteAddr() *net.TCPAddr {
+	return s.conn.RemoteAddr()
 }
 
-func (w *WrapStreamCipher) Close() error {
-	return w.conn.Close()
+func (s *StreamCipher) Close() error {
+	return s.conn.Close()
 }
 
 func NewStreamCipher(conn WrapStreamReadWriteCloser, encryptor, decryptor cipher.Stream) WrapStreamReadWriteCloser {
-	return &WrapStreamCipher{
+	return &StreamCipher{
 		conn:      conn,
 		encryptor: encryptor,
 		decryptor: decryptor,

wrappers/wrap.go

@@ -56,6 +56,11 @@ type WrapPacketReader interface {
 	Wrap
 }
 
+type WrapPacketWriter interface {
+	io.Writer
+	Wrap
+}
+
 type WrapPacketReadWriter interface {
 	io.Writer
 	WrapPacketReader