Move all wrappers

mtproto/wrappers/crypt_test.go

@@ -1,48 +0,0 @@
-package wrappers
-
-import (
-	"encoding/binary"
-	"net"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-
-	"github.com/9seconds/mtg/mtproto/rpc"
-)
-
-var proxySecret = []byte{196, 249, 250, 202, 150, 120, 230, 187, 72, 173,
-	108, 126, 44, 229, 192, 210, 68, 48, 100, 93, 85, 74, 221, 235, 85, 65,
-	158, 3, 77, 166, 39, 33, 208, 70, 234, 171, 110, 82, 171, 20, 169, 90, 68,
-	62, 207, 179, 70, 62, 121, 160, 90, 102, 97, 42, 223, 156, 174, 218, 139,
-	233, 168, 13, 166, 152, 111, 176, 166, 255, 56, 122, 248, 77, 136, 239,
-	58, 100, 19, 113, 62, 92, 51, 119, 246, 225, 163, 212, 125, 153, 245, 224,
-	197, 110, 236, 232, 240, 92, 84, 196, 144, 176, 121, 227, 27, 239, 130,
-	255, 14, 232, 242, 176, 163, 39, 86, 210, 73, 197, 242, 18, 105, 129, 108,
-	183, 6, 27, 38, 93, 178, 18}
-
-func TestMakeKeys(t *testing.T) {
-	req, err := rpc.NewNonceRequest(proxySecret)
-	assert.Nil(t, err)
-
-	copy(req.Nonce[:], []byte{24, 49, 53, 111, 198, 10, 235, 180, 230, 112, 92, 78, 1, 201, 106, 105})
-	binary.LittleEndian.PutUint32(req.CryptoTS[:], 1528396015)
-
-	resp := &rpc.NonceResponse{
-		NonceRequest: rpc.NonceRequest{
-			Nonce: []byte{247, 40, 210, 56, 65, 12, 101, 170, 216, 155, 14, 253, 250, 238, 219, 226},
-		},
-	}
-
-	cltAddr := &net.TCPAddr{
-		IP:   net.ParseIP("80.211.29.34"),
-		Port: 54208,
-	}
-	srvAddr := &net.TCPAddr{
-		IP:   net.ParseIP("149.154.162.38"),
-		Port: 80,
-	}
-
-	key, iv := makeKeys(CipherPurposeClient, req, resp, cltAddr, srvAddr, proxySecret)
-	assert.Equal(t, key, []byte{165, 158, 127, 49, 41, 232, 187, 69, 38, 29, 163, 226, 183, 146, 28, 67, 225, 224, 134, 191, 207, 152, 255, 166, 152, 66, 169, 196, 54, 135, 50, 188})
-	assert.Equal(t, iv, []byte{33, 110, 125, 221, 183, 121, 160, 116, 130, 180, 156, 249, 52, 111, 37, 178})
-}

mtproto/wrappers/proxy_request.go

@@ -1,111 +0,0 @@
-package wrappers
-
-import (
-	"bytes"
-	"io"
-	"io/ioutil"
-	"net"
-
-	"github.com/juju/errors"
-
-	"github.com/9seconds/mtg/mtproto"
-	"github.com/9seconds/mtg/mtproto/rpc"
-	"github.com/9seconds/mtg/utils"
-	"github.com/9seconds/mtg/wrappers"
-)
-
-type ProxyRequestReadWriteCloserWithAddr struct {
-	wrappers.BufferedReader
-
-	conn wrappers.ReadWriteCloserWithAddr
-	req  *rpc.ProxyRequest
-}
-
-func (p *ProxyRequestReadWriteCloserWithAddr) Read(buf []byte) (int, error) {
-	return p.BufferedRead(buf, func() error {
-		ans := make([]byte, 4)
-		if _, err := io.ReadFull(p.conn, ans); err != nil {
-			return errors.Annotate(err, "Cannot read RPC tag")
-		}
-
-		switch {
-		case bytes.Equal(ans, rpc.TagProxyAns):
-			return p.readProxyAns()
-		case bytes.Equal(ans, rpc.TagSimpleAck):
-			return p.readSimpleAck()
-		case bytes.Equal(ans, rpc.TagCloseExt):
-			return p.readCloseExt()
-		}
-
-		return errors.Errorf("Unknown RPC answer %v", ans)
-	})
-}
-
-func (p *ProxyRequestReadWriteCloserWithAddr) readCloseExt() error {
-	return errors.New("Connection has been closed remotely")
-}
-
-func (p *ProxyRequestReadWriteCloserWithAddr) readProxyAns() (err error) {
-	if _, err = io.CopyN(ioutil.Discard, p.conn, 8+4); err != nil {
-		return errors.Annotate(err, "Cannot skip flags and connid")
-	}
-
-	buf, err := utils.ReadCurrentData(p.conn)
-	if err != nil {
-		return errors.Annotate(err, "Cannot read proxy answer")
-	}
-	p.Buffer.Write(buf)
-
-	return nil
-}
-
-func (p *ProxyRequestReadWriteCloserWithAddr) readSimpleAck() error {
-	if _, err := io.CopyN(ioutil.Discard, p.conn, 8); err != nil {
-		return errors.Annotate(err, "Cannot skip connid")
-	}
-
-	ackData := make([]byte, 4)
-	if _, err := io.ReadFull(p.conn, ackData); err != nil {
-		return errors.Annotate(err, "Cannot read simple ack")
-	}
-	p.Buffer.Write(ackData)
-
-	return nil
-}
-
-func (p *ProxyRequestReadWriteCloserWithAddr) Write(raw []byte) (int, error) {
-	if _, err := p.conn.Write(p.req.Bytes(raw)); err != nil {
-		return 0, err
-	}
-
-	return len(raw), nil
-}
-
-func (p *ProxyRequestReadWriteCloserWithAddr) Close() error {
-	return p.conn.Close()
-}
-
-func (p *ProxyRequestReadWriteCloserWithAddr) LocalAddr() *net.TCPAddr {
-	return p.conn.LocalAddr()
-}
-
-func (p *ProxyRequestReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
-	return p.conn.RemoteAddr()
-}
-
-func (p *ProxyRequestReadWriteCloserWithAddr) SocketID() string {
-	return p.conn.SocketID()
-}
-
-func NewProxyRequestRWC(conn wrappers.ReadWriteCloserWithAddr, connOpts *mtproto.ConnectionOpts, adTag []byte) (wrappers.ReadWriteCloserWithAddr, 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 &ProxyRequestReadWriteCloserWithAddr{
-		BufferedReader: wrappers.NewBufferedReader(),
-		conn:           conn,
-		req:            req,
-	}, nil
-}

wrappers/mtproto_cipher.go

@@ -10,7 +10,7 @@ import (
 	"net"
 
 	"github.com/9seconds/mtg/mtproto/rpc"
-	"github.com/9seconds/mtg/wrappers"
+	"github.com/9seconds/mtg/utils"
 )
 
 type CipherPurpose uint8
@@ -22,20 +22,20 @@ const (
 
 var emptyIP = [4]byte{0x00, 0x00, 0x00, 0x00}
 
-func NewMiddleProxyCipherRWC(conn wrappers.ReadWriteCloserWithAddr, req *rpc.NonceRequest, resp *rpc.NonceResponse, secret []byte) wrappers.ReadWriteCloserWithAddr {
+func NewMiddleProxyCipher(conn WrapStreamReadWriteCloser, req *rpc.NonceRequest, resp *rpc.NonceResponse, secret []byte) WrapStreamReadWriteCloser {
 	localAddr := conn.LocalAddr()
 	remoteAddr := conn.RemoteAddr()
 
-	encKey, encIV := makeKeys(CipherPurposeClient, req, resp, localAddr, remoteAddr, secret)
+	encKey, encIV := deriveKeys(CipherPurposeClient, req, resp, localAddr, remoteAddr, secret)
-	decKey, decIV := makeKeys(CipherPurposeServer, req, resp, localAddr, remoteAddr, secret)
+	decKey, decIV := deriveKeys(CipherPurposeServer, req, resp, localAddr, remoteAddr, secret)
 
 	enc, _ := makeEncrypterDecrypter(encKey, encIV)
 	_, dec := makeEncrypterDecrypter(decKey, decIV)
 
-	return wrappers.NewBlockCipherRWC(conn, enc, dec)
+	return NewWrapBlockCipher(conn, enc, dec)
 }
 
-func makeKeys(purpose CipherPurpose, req *rpc.NonceRequest, resp *rpc.NonceResponse, client *net.TCPAddr, remote *net.TCPAddr, secret []byte) ([]byte, []byte) {
+func deriveKeys(purpose CipherPurpose, req *rpc.NonceRequest, resp *rpc.NonceResponse, client *net.TCPAddr, remote *net.TCPAddr, secret []byte) ([]byte, []byte) {
 	message := bytes.Buffer{}
 	message.Write(resp.Nonce[:])
 	message.Write(req.Nonce[:])
@@ -44,8 +44,8 @@ func makeKeys(purpose CipherPurpose, req *rpc.NonceRequest, resp *rpc.NonceRespo
 	clientIPv4 := emptyIP[:]
 	serverIPv4 := emptyIP[:]
 	if client.IP.To4() != nil {
-		clientIPv4 = reverseBytes(client.IP.To4())
+		clientIPv4 = utils.ReverseBytes(client.IP.To4())
-		serverIPv4 = reverseBytes(remote.IP.To4())
+		serverIPv4 = utils.ReverseBytes(remote.IP.To4())
 	}
 	message.Write(serverIPv4)
 
@@ -92,16 +92,3 @@ func makeEncrypterDecrypter(key, iv []byte) (cipher.BlockMode, cipher.BlockMode)
 
 	return cipher.NewCBCEncrypter(block, iv), cipher.NewCBCDecrypter(block, iv)
 }
-
-func reverseBytes(data []byte) []byte {
-	dataLen := len(data)
-	rv := make([]byte, dataLen)
-
-	rv[dataLen/2] = data[dataLen/2]
-	for i := dataLen/2 - 1; i >= 0; i-- {
-		opp := dataLen - i - 1
-		rv[i], rv[opp] = data[opp], data[i]
-	}
-
-	return rv
-}

wrappers/mtproto_proxy.go

@@ -0,0 +1,135 @@
+package wrappers
+
+import (
+	"bytes"
+	"net"
+
+	"github.com/9seconds/mtg/mtproto/rpc"
+	"github.com/juju/errors"
+)
+
+type MTProtoProxy struct {
+	conn WrapPacketReadWriteCloser
+	req  *rpc.ProxyRequest
+
+	readCounter  uint32
+	writeCounter uint32
+}
+
+func (m *MTProtoProxy) Read() ([]byte, error) {
+	m.LogDebug("Read packet",
+		"counter", m.readCounter,
+		"simple_ack", m.req.Options.WriteHacks.SimpleAck,
+		"quick_ack", m.req.Options.WriteHacks.QuickAck,
+	)
+
+	packet, err := m.conn.Read()
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot read packet")
+	}
+	m.LogDebug("Read packet length",
+		"counter", m.readCounter,
+		"simple_ack", m.req.Options.WriteHacks.SimpleAck,
+		"quick_ack", m.req.Options.WriteHacks.QuickAck,
+		"length", len(packet),
+	)
+
+	if len(packet) < 4 {
+		return nil, errors.Annotate(err, "Incorrect packet length")
+	}
+
+	tag, packet := packet[:4], packet[4:]
+	m.LogDebug("Read RPC tag",
+		"counter", m.readCounter,
+		"simple_ack", m.req.Options.WriteHacks.SimpleAck,
+		"quick_ack", m.req.Options.WriteHacks.QuickAck,
+		"tag", tag,
+	)
+
+	m.readCounter++
+	switch {
+	case bytes.Equal(tag, rpc.TagProxyAns):
+		return m.readProxyAns(packet)
+	case bytes.Equal(tag, rpc.TagSimpleAck):
+		return m.readSimpleAck(packet)
+	case bytes.Equal(tag, rpc.TagCloseExt):
+		return m.readCloseExt(packet)
+	}
+
+	return nil, errors.Errorf("Unknown RPC answer %v", tag)
+}
+
+func (m *MTProtoProxy) readProxyAns(data []byte) ([]byte, error) {
+	if len(data) < 12 {
+		return nil, errors.Errorf("Incorrect data of proxy answer: %d", len(data))
+	}
+
+	return data[12:], nil
+}
+
+func (m *MTProtoProxy) readSimpleAck(data []byte) ([]byte, error) {
+	if len(data) != 12 {
+		return nil, errors.Errorf("Incorrect data of simple ack: %d", len(data))
+	}
+
+	return data[8:12], nil // 0:8 - connection id
+}
+
+func (m *MTProtoProxy) readCloseExt(data []byte) ([]byte, error) {
+	return nil, errors.New("Connection has been closed remotely by RPC call")
+}
+
+func (m *MTProtoProxy) Write(p []byte) (int, error) {
+	m.LogDebug("Write packet",
+		"length", len(p),
+		"counter", m.writeCounter,
+		"simple_ack", m.req.Options.ReadHacks.SimpleAck,
+		"quick_ack", m.req.Options.ReadHacks.QuickAck,
+	)
+	m.writeCounter++
+
+	if _, err := m.conn.Write(p); err != nil {
+		return 0, err
+	}
+
+	return len(p), nil
+}
+
+func (m *MTProtoProxy) LogDebug(msg string, data ...interface{}) {
+	data = append(data, []interface{}{"type", "proxy"})
+	m.conn.LogDebug(msg, data...)
+}
+
+func (m *MTProtoProxy) LogInfo(msg string, data ...interface{}) {
+	data = append(data, []interface{}{"type", "proxy"})
+	m.conn.LogInfo(msg, data...)
+}
+
+func (m *MTProtoProxy) LogWarn(msg string, data ...interface{}) {
+	data = append(data, []interface{}{"type", "proxy"})
+	m.conn.LogWarn(msg, data...)
+}
+
+func (m *MTProtoProxy) LogError(msg string, data ...interface{}) {
+	data = append(data, []interface{}{"type", "proxy"})
+	m.conn.LogError(msg, data...)
+}
+
+func (m *MTProtoProxy) LocalAddr() *net.TCPAddr {
+	return m.conn.LocalAddr()
+}
+
+func (m *MTProtoProxy) RemoteAddr() *net.TCPAddr {
+	return m.conn.RemoteAddr()
+}
+
+func (m *MTProtoProxy) Close() error {
+	return m.conn.Close()
+}
+
+func NewMTProtoProxy(conn WrapPacketReadWriteCloser, req *rpc.ProxyRequest) WrapPacketReadWriteCloser {
+	return &MTProtoProxy{
+		conn: conn,
+		req:  req,
+	}
+}