Refactored all the things!

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...)
+	}
 
-		w.decryptor.CryptBlocks(buf, buf)
-		w.Buffer.Write(buf)
+	b.decryptor.CryptBlocks(buf, buf)
+	b.buf.Write(buf)
 
-		return nil
-	})
+	return b.flush(p)
 }
 
-func (w *WrapBlockCipher) Write(p []byte) (int, error) {
+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 b.buf.Read(p)
+}
+
+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