diff --git a/client/client.go b/client/client.go index de5e32a..cb6b291 100644 --- a/client/client.go +++ b/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) diff --git a/client/direct.go b/client/direct.go index 23e9613..19d1ede 100644 --- a/client/direct.go +++ b/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") } diff --git a/mtproto/wrappers/crypt.go b/mtproto/wrappers/crypt.go index fbff067..7ae5255 100644 --- a/mtproto/wrappers/crypt.go +++ b/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[:]) diff --git a/mtproto/wrappers/frame.go b/mtproto/wrappers/frame.go index 9e7a4c7..fc3e1d6 100644 --- a/mtproto/wrappers/frame.go +++ b/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, diff --git a/telegram/dialer.go b/telegram/dialer.go index f6b2cf7..f0d9cb5 100644 --- a/telegram/dialer.go +++ b/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 diff --git a/telegram/direct.go b/telegram/direct.go index 5681415..81fb2eb 100644 --- a/telegram/direct.go +++ b/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 { diff --git a/telegram/middle_caller.go b/telegram/middle_caller.go index 2d0d86e..acc1b27 100644 --- a/telegram/middle_caller.go +++ b/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 diff --git a/telegram/telegram.go b/telegram/telegram.go index 6bbb504..5ce16ad 100644 --- a/telegram/telegram.go +++ b/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) - Init(*mtproto.ConnectionOpts, io.ReadWriteCloser) (io.ReadWriteCloser, error) + Dial(*mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, 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 { diff --git a/wrappers/blockcipherrwc.go b/wrappers/blockcipherrwc.go index 1645034..ad2c9d1 100644 --- a/wrappers/blockcipherrwc.go +++ b/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 { - return &BlockCipherReadWriteCloser{ +func NewBlockCipherRWC(conn ReadWriteCloserWithAddr, encryptor, decryptor cipher.BlockMode) ReadWriteCloserWithAddr { + return &BlockCipherReadWriteCloserWithAddr{ buf: getBuffer(), conn: conn, encryptor: encryptor, diff --git a/wrappers/ctxrwc.go b/wrappers/ctxrwc.go index 26f47be..b378a96 100644 --- a/wrappers/ctxrwc.go +++ b/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 { - return &CtxReadWriteCloser{ +func NewCtxRWC(ctx context.Context, cancel context.CancelFunc, conn ReadWriteCloserWithAddr) ReadWriteCloserWithAddr { + return &CtxReadWriteCloserWithAddr{ conn: conn, ctx: ctx, cancel: cancel, diff --git a/wrappers/logrwc.go b/wrappers/logrwc.go index 355e239..bab7af7 100644 --- a/wrappers/logrwc.go +++ b/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 { - conn io.ReadWriteCloser +type LogReadWriteCloserWithAddr struct { + 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 { - return &LogReadWriteCloser{ +func NewLogRWC(conn ReadWriteCloserWithAddr, logger *zap.SugaredLogger, sockid string, name string) ReadWriteCloserWithAddr { + return &LogReadWriteCloserWithAddr{ conn: conn, logger: logger, sockid: sockid, diff --git a/wrappers/rwcaddr.go b/wrappers/rwcaddr.go new file mode 100644 index 0000000..804bc5a --- /dev/null +++ b/wrappers/rwcaddr.go @@ -0,0 +1,12 @@ +package wrappers + +import ( + "io" + "net" +) + +type ReadWriteCloserWithAddr interface { + io.ReadWriteCloser + + Addr() *net.TCPAddr +} diff --git a/wrappers/streamcipherrwc.go b/wrappers/streamcipherrwc.go index 1d7d73c..d847c72 100644 --- a/wrappers/streamcipherrwc.go +++ b/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 { - return &StreamCipherReadWriteCloser{ +func NewStreamCipherRWC(conn ReadWriteCloserWithAddr, encryptor, decryptor cipher.Stream) ReadWriteCloserWithAddr { + return &StreamCipherReadWriteCloserWithAddr{ conn: conn, encryptor: encryptor, decryptor: decryptor, diff --git a/wrappers/timeoutrwc.go b/wrappers/timeoutrwc.go index e4ea9b0..e758600 100644 --- a/wrappers/timeoutrwc.go +++ b/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 { - return &TimeoutReadWriteCloser{conn} +func (t *TimeoutReadWriteCloserWithAddr) Addr() *net.TCPAddr { + return t.conn.RemoteAddr().(*net.TCPAddr) +} + +func NewTimeoutRWC(conn net.Conn) ReadWriteCloserWithAddr { + return &TimeoutReadWriteCloserWithAddr{conn} } diff --git a/wrappers/trafficrwc.go b/wrappers/trafficrwc.go index 485a54c..78518ad 100644 --- a/wrappers/trafficrwc.go +++ b/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 { - conn io.ReadWriteCloser +type TrafficReadWriteCloserWithAddr struct { + 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 { - return &TrafficReadWriteCloser{ +func NewTrafficRWC(conn ReadWriteCloserWithAddr, readCallback, writeCallback func(int)) ReadWriteCloserWithAddr { + return &TrafficReadWriteCloserWithAddr{ conn: conn, readCallback: readCallback, writeCallback: writeCallback,