Change rwc to rwc with addresses

client/client.go

@@ -1,12 +1,12 @@
 package client
 
 import (
-	"io"
 	"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, *config.Config) (*mtproto.ConnectionOpts, io.ReadWriteCloser, error)
+type Init func(net.Conn, *config.Config) (*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr, error)

client/direct.go

@@ -1,7 +1,6 @@
 package client
 
 import (
-	"io"
 	"net"
 	"time"
 
@@ -16,7 +15,7 @@ import (
 const handshakeTimeout = 10 * time.Second
 
 // DirectInit initializes client to access Telegram bypassing middleproxies.
-func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, io.ReadWriteCloser, error) {
+func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr, error) {
 	if err := config.SetSocketOptions(conn); err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot set socket options")
 	}

mtproto/wrappers/crypt.go

@@ -7,7 +7,6 @@ import (
 	"crypto/md5"
 	"crypto/sha1"
 	"encoding/binary"
-	"io"
 	"net"
 
 	"github.com/9seconds/mtg/mtproto/rpc"
@@ -23,14 +22,17 @@ const (
 
 var emptyIP = [4]byte{0x00, 0x00, 0x00, 0x00}
 
-func NewMiddleProxyCipherRWC(conn io.ReadWriteCloser, req *rpc.RPCNonceRequest, resp *rpc.RPCNonceResponse, client *net.TCPAddr, remote *net.TCPAddr, secret []byte) io.ReadWriteCloser {
+func NewMiddleProxyCipherRWC(conn wrappers.ReadWriteCloserWithAddr, req *rpc.RPCNonceRequest,
+	resp *rpc.RPCNonceResponse, client *net.TCPAddr, remote *net.TCPAddr,
+	secret []byte) wrappers.ReadWriteCloserWithAddr {
 	encryptor := newCBCCipher(CipherPurposeClient, req, resp, client, remote, secret)
 	decryptor := newCBCCipher(CipherPurposeServer, req, resp, client, remote, secret)
 
 	return wrappers.NewBlockCipherRWC(conn, encryptor, decryptor)
 }
 
-func newCBCCipher(purpose CipherPurpose, req *rpc.RPCNonceRequest, resp *rpc.RPCNonceResponse, client *net.TCPAddr, remote *net.TCPAddr, secret []byte) cipher.BlockMode {
+func newCBCCipher(purpose CipherPurpose, req *rpc.RPCNonceRequest, resp *rpc.RPCNonceResponse,
+	client *net.TCPAddr, remote *net.TCPAddr, secret []byte) cipher.BlockMode {
 	message := bytes.Buffer{}
 	message.Write(resp.Nonce[:])
 	message.Write(req.Nonce[:])

mtproto/wrappers/frame.go

@@ -6,8 +6,11 @@ import (
 	"encoding/binary"
 	"hash/crc32"
 	"io"
+	"net"
 
 	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/wrappers"
 )
 
 // Frame: { MessageLength(4) | SequenceNumber(4) | Message(???) | CRC32(4) [| padding(4), ...] }
@@ -19,7 +22,7 @@ const (
 var frameRWCPadding = [4]byte{0x04, 0x00, 0x00, 0x00}
 
 type FrameRWC struct {
-	conn io.ReadWriteCloser
+	conn wrappers.ReadWriteCloserWithAddr
 
 	readSeqNo  int32
 	writeSeqNo int32
@@ -102,6 +105,10 @@ func (f *FrameRWC) Close() error {
 	return f.conn.Close()
 }
 
+func (f *FrameRWC) Addr() *net.TCPAddr {
+	return f.conn.Addr()
+}
+
 func (f *FrameRWC) flush(p []byte) (int, error) {
 	sizeToRead := len(p)
 	if f.readBuf.Len() < sizeToRead {
@@ -119,7 +126,7 @@ func (f *FrameRWC) flush(p []byte) (int, error) {
 	return sizeToRead, nil
 }
 
-func NewFrameRWC(conn io.ReadWriteCloser, seqNo int32) io.ReadWriteCloser {
+func NewFrameRWC(conn wrappers.ReadWriteCloserWithAddr, seqNo int32) wrappers.ReadWriteCloserWithAddr {
 	return &FrameRWC{
 		conn:       conn,
 		readSeqNo:  seqNo,

telegram/dialer.go

@@ -1,7 +1,6 @@
 package telegram
 
 import (
-	"io"
 	"net"
 	"time"
 
@@ -29,7 +28,7 @@ func (t *tgDialer) dial(addr string) (net.Conn, error) {
 	return conn, nil
 }
 
-func (t *tgDialer) dialRWC(addr string) (io.ReadWriteCloser, error) {
+func (t *tgDialer) dialRWC(addr string) (wrappers.ReadWriteCloserWithAddr, error) {
 	conn, err := t.dial(addr)
 	if err != nil {
 		return nil, err

telegram/direct.go

@@ -1,7 +1,6 @@
 package telegram
 
 import (
-	"io"
 	"net"
 
 	"github.com/juju/errors"
@@ -33,7 +32,7 @@ type directTelegram struct {
 	baseTelegram
 }
 
-func (t *directTelegram) Dial(connOpts *mtproto.ConnectionOpts) (io.ReadWriteCloser, error) {
+func (t *directTelegram) Dial(connOpts *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error) {
 	dc := connOpts.DC
 	if dc < 0 {
 		dc = -dc
@@ -44,7 +43,7 @@ func (t *directTelegram) Dial(connOpts *mtproto.ConnectionOpts) (io.ReadWriteClo
 	return t.baseTelegram.dial(dc-1, connOpts.ConnectionProto)
 }
 
-func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts, conn io.ReadWriteCloser) (io.ReadWriteCloser, error) {
+func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.ReadWriteCloserWithAddr) (wrappers.ReadWriteCloserWithAddr, error) {
 	obfs2, frame := obfuscated2.MakeTelegramObfuscated2Frame(connOpts)
 
 	if n, err := conn.Write(frame); err != nil || n != obfuscated2.FrameLen {

telegram/middle_caller.go

@@ -2,7 +2,6 @@ package telegram
 
 import (
 	"bufio"
-	"io"
 	"io/ioutil"
 	"net"
 	"net/http"
@@ -16,6 +15,7 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/wrappers"
 )
 
 const (
@@ -39,7 +39,7 @@ type middleTelegramCaller struct {
 	httpClient  *http.Client
 }
 
-func (t *middleTelegramCaller) Dial(connOpts *mtproto.ConnectionOpts) (io.ReadWriteCloser, error) {
+func (t *middleTelegramCaller) Dial(connOpts *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error) {
 	dc := connOpts.DC
 	if dc == 0 {
 		dc = 1

telegram/telegram.go

@@ -1,20 +1,20 @@
 package telegram
 
 import (
-	"io"
 	"math/rand"
 
 	"github.com/juju/errors"
 
 	"github.com/9seconds/mtg/mtproto"
+	"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(*mtproto.ConnectionOpts) (io.ReadWriteCloser, error)
+	Dial(*mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error)
-	Init(*mtproto.ConnectionOpts, io.ReadWriteCloser) (io.ReadWriteCloser, error)
+	Init(*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr) (wrappers.ReadWriteCloserWithAddr, error)
 }
 
 type baseTelegram struct {
@@ -24,7 +24,7 @@ type baseTelegram struct {
 	v6Addresses map[int16][]string
 }
 
-func (b *baseTelegram) dial(dcIdx int16, proto mtproto.ConnectionProtocol) (io.ReadWriteCloser, error) {
+func (b *baseTelegram) dial(dcIdx int16, proto mtproto.ConnectionProtocol) (wrappers.ReadWriteCloserWithAddr, error) {
 	addrs := make([]string, 2)
 
 	if proto&mtproto.ConnectionProtocolIPv6 != 0 {

wrappers/blockcipherrwc.go

@@ -4,20 +4,20 @@ import (
 	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
-	"io"
+	"net"
 
 	"github.com/juju/errors"
 )
 
-type BlockCipherReadWriteCloser struct {
+type BlockCipherReadWriteCloserWithAddr struct {
 	buf *bytes.Buffer
 
-	conn      io.ReadWriteCloser
+	conn      ReadWriteCloserWithAddr
 	encryptor cipher.BlockMode
 	decryptor cipher.BlockMode
 }
 
-func (c *BlockCipherReadWriteCloser) Read(p []byte) (int, error) {
+func (c *BlockCipherReadWriteCloserWithAddr) Read(p []byte) (int, error) {
 	if c.buf.Len() > 0 {
 		return c.flush(p)
 	}
@@ -34,7 +34,7 @@ func (c *BlockCipherReadWriteCloser) Read(p []byte) (int, error) {
 	return c.flush(p)
 }
 
-func (c *BlockCipherReadWriteCloser) Write(p []byte) (int, error) {
+func (c *BlockCipherReadWriteCloserWithAddr) Write(p []byte) (int, error) {
 	if len(p)%aes.BlockSize > 0 {
 		return 0, errors.Errorf("Incorrect block size %d", len(p))
 	}
@@ -50,12 +50,16 @@ func (c *BlockCipherReadWriteCloser) Write(p []byte) (int, error) {
 	return c.conn.Write(encrypted)
 }
 
-func (c *BlockCipherReadWriteCloser) Close() error {
+func (c *BlockCipherReadWriteCloserWithAddr) Close() error {
 	defer putBuffer(c.buf)
 	return c.conn.Close()
 }
 
-func (c *BlockCipherReadWriteCloser) flush(p []byte) (int, error) {
+func (c *BlockCipherReadWriteCloserWithAddr) Addr() *net.TCPAddr {
+	return c.conn.Addr()
+}
+
+func (c *BlockCipherReadWriteCloserWithAddr) flush(p []byte) (int, error) {
 	sizeToRead := len(p)
 	if c.buf.Len() < sizeToRead {
 		sizeToRead = c.buf.Len()
@@ -76,8 +80,8 @@ func (c *BlockCipherReadWriteCloser) flush(p []byte) (int, error) {
 	return sizeToRead, nil
 }
 
-func NewBlockCipherRWC(conn io.ReadWriteCloser, encryptor, decryptor cipher.BlockMode) io.ReadWriteCloser {
+func NewBlockCipherRWC(conn ReadWriteCloserWithAddr, encryptor, decryptor cipher.BlockMode) ReadWriteCloserWithAddr {
-	return &BlockCipherReadWriteCloser{
+	return &BlockCipherReadWriteCloserWithAddr{
 		buf:       getBuffer(),
 		conn:      conn,
 		encryptor: encryptor,

wrappers/ctxrwc.go

@@ -2,21 +2,21 @@ package wrappers
 
 import (
 	"context"
-	"io"
+	"net"
 
 	"github.com/juju/errors"
 )
 
 // CtxReadWriteCloser wraps underlying connection and does management of the
 // context and its cancel function.
-type CtxReadWriteCloser struct {
+type CtxReadWriteCloserWithAddr struct {
 	ctx    context.Context
-	conn   io.ReadWriteCloser
+	conn   ReadWriteCloserWithAddr
 	cancel context.CancelFunc
 }
 
 // Read reads from connection
-func (c *CtxReadWriteCloser) Read(p []byte) (int, error) {
+func (c *CtxReadWriteCloserWithAddr) Read(p []byte) (int, error) {
 	select {
 	case <-c.ctx.Done():
 		return 0, errors.Annotate(c.ctx.Err(), "Read is failed because of closed context")
@@ -30,7 +30,7 @@ func (c *CtxReadWriteCloser) Read(p []byte) (int, error) {
 }
 
 // Write writes into connection.
-func (c *CtxReadWriteCloser) Write(p []byte) (int, error) {
+func (c *CtxReadWriteCloserWithAddr) Write(p []byte) (int, error) {
 	select {
 	case <-c.ctx.Done():
 		return 0, errors.Annotate(c.ctx.Err(), "Write is failed because of closed context")
@@ -44,14 +44,18 @@ func (c *CtxReadWriteCloser) Write(p []byte) (int, error) {
 }
 
 // Close closes underlying connection.
-func (c *CtxReadWriteCloser) Close() error {
+func (c *CtxReadWriteCloserWithAddr) Close() error {
 	return c.conn.Close()
 }
 
+func (c *CtxReadWriteCloserWithAddr) Addr() *net.TCPAddr {
+	return c.conn.Addr()
+}
+
 // NewCtxRWC returns ReadWriteCloser which respects given context,
 // cancellation etc.
-func NewCtxRWC(ctx context.Context, cancel context.CancelFunc, conn io.ReadWriteCloser) io.ReadWriteCloser {
+func NewCtxRWC(ctx context.Context, cancel context.CancelFunc, conn ReadWriteCloserWithAddr) ReadWriteCloserWithAddr {
-	return &CtxReadWriteCloser{
+	return &CtxReadWriteCloserWithAddr{
 		conn:   conn,
 		ctx:    ctx,
 		cancel: cancel,

wrappers/logrwc.go

@@ -1,44 +1,48 @@
 package wrappers
 
 import (
-	"io"
+	"net"
 
 	"go.uber.org/zap"
 )
 
 // LogReadWriteCloser adds additional logging for reading/writing. All
 // logging is performed for debug mode only.
-type LogReadWriteCloser struct {
+type LogReadWriteCloserWithAddr struct {
-	conn   io.ReadWriteCloser
+	conn   ReadWriteCloserWithAddr
 	logger *zap.SugaredLogger
 	sockid string
 	name   string
 }
 
 // Read reads from connection
-func (l *LogReadWriteCloser) Read(p []byte) (n int, err error) {
+func (l *LogReadWriteCloserWithAddr) Read(p []byte) (n int, err error) {
 	n, err = l.conn.Read(p)
 	l.logger.Debugw("Finish reading", "name", l.name, "socketid", l.sockid, "nbytes", n, "error", err)
 	return
 }
 
 // Write writes into connection.
-func (l *LogReadWriteCloser) Write(p []byte) (n int, err error) {
+func (l *LogReadWriteCloserWithAddr) Write(p []byte) (n int, err error) {
 	n, err = l.conn.Write(p)
 	l.logger.Debugw("Finish writing", "name", l.name, "socketid", l.sockid, "nbytes", n, "error", err)
 	return
 }
 
 // Close closes underlying connection.
-func (l *LogReadWriteCloser) Close() error {
+func (l *LogReadWriteCloserWithAddr) Close() error {
 	err := l.conn.Close()
 	l.logger.Debugw("Finish closing socket", "name", l.name, "socketid", l.sockid, "error", err)
 	return err
 }
 
+func (l *LogReadWriteCloserWithAddr) Addr() *net.TCPAddr {
+	return l.conn.Addr()
+}
+
 // NewLogRWC wraps ReadWriteCloser with logger calls.
-func NewLogRWC(conn io.ReadWriteCloser, logger *zap.SugaredLogger, sockid string, name string) io.ReadWriteCloser {
+func NewLogRWC(conn ReadWriteCloserWithAddr, logger *zap.SugaredLogger, sockid string, name string) ReadWriteCloserWithAddr {
-	return &LogReadWriteCloser{
+	return &LogReadWriteCloserWithAddr{
 		conn:   conn,
 		logger: logger,
 		sockid: sockid,

wrappers/rwcaddr.go

@@ -0,0 +1,12 @@
+package wrappers
+
+import (
+	"io"
+	"net"
+)
+
+type ReadWriteCloserWithAddr interface {
+	io.ReadWriteCloser
+
+	Addr() *net.TCPAddr
+}

wrappers/streamcipherrwc.go

@@ -2,26 +2,26 @@ package wrappers
 
 import (
 	"crypto/cipher"
-	"io"
+	"net"
 )
 
 // StreamCipherReadWriteCloser is a ReadWriteCloser which ciphers
 // incoming and outgoing data with givem cipher.Stream instances.
-type StreamCipherReadWriteCloser struct {
+type StreamCipherReadWriteCloserWithAddr struct {
 	encryptor cipher.Stream
 	decryptor cipher.Stream
-	conn      io.ReadWriteCloser
+	conn      ReadWriteCloserWithAddr
 }
 
 // Read reads from connection
-func (c *StreamCipherReadWriteCloser) Read(p []byte) (n int, err error) {
+func (c *StreamCipherReadWriteCloserWithAddr) Read(p []byte) (n int, err error) {
 	n, err = c.conn.Read(p)
 	c.decryptor.XORKeyStream(p, p[:n])
 	return
 }
 
 // Write writes into connection.
-func (c *StreamCipherReadWriteCloser) Write(p []byte) (int, error) {
+func (c *StreamCipherReadWriteCloserWithAddr) Write(p []byte) (int, error) {
 	// This is to decrease an amount of allocations. Unfortunately, escape
 	// analysis in (at least Golang 1.10) is absolutely not perfect. For
 	// example, it understands that we want to have a slice locally, right?
@@ -39,14 +39,18 @@ func (c *StreamCipherReadWriteCloser) Write(p []byte) (int, error) {
 }
 
 // Close closes underlying connection.
-func (c *StreamCipherReadWriteCloser) Close() error {
+func (c *StreamCipherReadWriteCloserWithAddr) Close() error {
 	return c.conn.Close()
 }
 
+func (c *StreamCipherReadWriteCloserWithAddr) Addr() *net.TCPAddr {
+	return c.conn.Addr()
+}
+
 // NewStreamCipherRWC returns wrapper which transparently
 // encrypts/decrypts traffic with obfuscated2 protocol.
-func NewStreamCipherRWC(conn io.ReadWriteCloser, encryptor, decryptor cipher.Stream) io.ReadWriteCloser {
+func NewStreamCipherRWC(conn ReadWriteCloserWithAddr, encryptor, decryptor cipher.Stream) ReadWriteCloserWithAddr {
-	return &StreamCipherReadWriteCloser{
+	return &StreamCipherReadWriteCloserWithAddr{
 		conn:      conn,
 		encryptor: encryptor,
 		decryptor: decryptor,

wrappers/timeoutrwc.go

@@ -1,31 +1,34 @@
 package wrappers
 
 import (
-	"io"
 	"net"
 	"time"
 
 	"github.com/9seconds/mtg/config"
 )
 
-type TimeoutReadWriteCloser struct {
+type TimeoutReadWriteCloserWithAddr struct {
 	conn net.Conn
 }
 
-func (t *TimeoutReadWriteCloser) Read(p []byte) (int, error) {
+func (t *TimeoutReadWriteCloserWithAddr) Read(p []byte) (int, error) {
 	t.conn.SetReadDeadline(time.Now().Add(config.TimeoutRead))
 	return t.conn.Read(p)
 }
 
-func (t *TimeoutReadWriteCloser) Write(p []byte) (int, error) {
+func (t *TimeoutReadWriteCloserWithAddr) Write(p []byte) (int, error) {
 	t.conn.SetWriteDeadline(time.Now().Add(config.TimeoutWrite))
 	return t.conn.Write(p)
 }
 
-func (t *TimeoutReadWriteCloser) Close() error {
+func (t *TimeoutReadWriteCloserWithAddr) Close() error {
 	return t.conn.Close()
 }
 
-func NewTimeoutRWC(conn net.Conn) io.ReadWriteCloser {
+func (t *TimeoutReadWriteCloserWithAddr) Addr() *net.TCPAddr {
-	return &TimeoutReadWriteCloser{conn}
+	return t.conn.RemoteAddr().(*net.TCPAddr)
+}
+
+func NewTimeoutRWC(conn net.Conn) ReadWriteCloserWithAddr {
+	return &TimeoutReadWriteCloserWithAddr{conn}
 }

wrappers/trafficrwc.go

@@ -1,37 +1,41 @@
 package wrappers
 
-import "io"
+import "net"
 
 // TrafficReadWriteCloser counts an amount of ingress/egress traffic by
 // calling given callbacks.
-type TrafficReadWriteCloser struct {
+type TrafficReadWriteCloserWithAddr struct {
-	conn          io.ReadWriteCloser
+	conn          ReadWriteCloserWithAddr
 	readCallback  func(int)
 	writeCallback func(int)
 }
 
 // Read reads from connection
-func (t *TrafficReadWriteCloser) Read(p []byte) (n int, err error) {
+func (t *TrafficReadWriteCloserWithAddr) Read(p []byte) (n int, err error) {
 	n, err = t.conn.Read(p)
 	t.readCallback(n)
 	return
 }
 
 // Write writes into connection.
-func (t *TrafficReadWriteCloser) Write(p []byte) (n int, err error) {
+func (t *TrafficReadWriteCloserWithAddr) Write(p []byte) (n int, err error) {
 	n, err = t.conn.Write(p)
 	t.writeCallback(n)
 	return
 }
 
 // Close closes underlying connection.
-func (t *TrafficReadWriteCloser) Close() error {
+func (t *TrafficReadWriteCloserWithAddr) Close() error {
 	return t.conn.Close()
 }
 
+func (t *TrafficReadWriteCloserWithAddr) Addr() *net.TCPAddr {
+	return t.conn.Addr()
+}
+
 // NewTrafficRWC wraps ReadWriteCloser to have read/write callbacks.
-func NewTrafficRWC(conn io.ReadWriteCloser, readCallback, writeCallback func(int)) io.ReadWriteCloser {
+func NewTrafficRWC(conn ReadWriteCloserWithAddr, readCallback, writeCallback func(int)) ReadWriteCloserWithAddr {
-	return &TrafficReadWriteCloser{
+	return &TrafficReadWriteCloserWithAddr{
 		conn:          conn,
 		readCallback:  readCallback,
 		writeCallback: writeCallback,