Change rwc to rwc with addresses

This commit is contained in:
9seconds
2018-07-02 16:12:31 +03:00
parent 6fa5f31ca8
commit 09476cc467
15 changed files with 109 additions and 68 deletions
+2 -2
View File
@@ -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)
+1 -2
View File
@@ -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")
}
+5 -3
View File
@@ -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[:])
+9 -2
View File
@@ -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,
+1 -2
View File
@@ -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
+2 -3
View File
@@ -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 {
+2 -2
View File
@@ -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
+4 -4
View File
@@ -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 {
+13 -9
View File
@@ -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,
+12 -8
View File
@@ -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,
+12 -8
View File
@@ -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,
+12
View File
@@ -0,0 +1,12 @@
package wrappers
import (
"io"
"net"
)
type ReadWriteCloserWithAddr interface {
io.ReadWriteCloser
Addr() *net.TCPAddr
}
+12 -8
View File
@@ -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,
+10 -7
View File
@@ -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}
}
+12 -8
View File
@@ -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,