Move all wrappers

wrappers/mtproto_cipher.go

@@ -0,0 +1,94 @@
+package wrappers
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/md5"
+	"crypto/sha1"
+	"encoding/binary"
+	"net"
+
+	"github.com/9seconds/mtg/mtproto/rpc"
+	"github.com/9seconds/mtg/utils"
+)
+
+type CipherPurpose uint8
+
+const (
+	CipherPurposeClient CipherPurpose = iota
+	CipherPurposeServer
+)
+
+var emptyIP = [4]byte{0x00, 0x00, 0x00, 0x00}
+
+func NewMiddleProxyCipher(conn WrapStreamReadWriteCloser, req *rpc.NonceRequest, resp *rpc.NonceResponse, secret []byte) WrapStreamReadWriteCloser {
+	localAddr := conn.LocalAddr()
+	remoteAddr := conn.RemoteAddr()
+
+	encKey, encIV := deriveKeys(CipherPurposeClient, req, resp, localAddr, remoteAddr, secret)
+	decKey, decIV := deriveKeys(CipherPurposeServer, req, resp, localAddr, remoteAddr, secret)
+
+	enc, _ := makeEncrypterDecrypter(encKey, encIV)
+	_, dec := makeEncrypterDecrypter(decKey, decIV)
+
+	return NewWrapBlockCipher(conn, enc, dec)
+}
+
+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[:])
+	message.Write(req.CryptoTS[:])
+
+	clientIPv4 := emptyIP[:]
+	serverIPv4 := emptyIP[:]
+	if client.IP.To4() != nil {
+		clientIPv4 = utils.ReverseBytes(client.IP.To4())
+		serverIPv4 = utils.ReverseBytes(remote.IP.To4())
+	}
+	message.Write(serverIPv4)
+
+	var port [2]byte
+	binary.LittleEndian.PutUint16(port[:], uint16(client.Port))
+	message.Write(port[:])
+
+	switch purpose {
+	case CipherPurposeClient:
+		message.WriteString("CLIENT")
+	case CipherPurposeServer:
+		message.WriteString("SERVER")
+	default:
+		panic("Unexpected cipher purpose")
+	}
+
+	message.Write(clientIPv4)
+	binary.LittleEndian.PutUint16(port[:], uint16(remote.Port))
+	message.Write(port[:])
+	message.Write(secret)
+	message.Write(resp.Nonce[:])
+
+	if client.IP.To4() == nil {
+		message.Write(client.IP.To16())
+		message.Write(remote.IP.To16())
+	}
+	message.Write(req.Nonce[:])
+
+	data := message.Bytes()
+	md5sum := md5.Sum(data[1:])
+	sha1sum := sha1.Sum(data)
+
+	key := append(md5sum[:12], sha1sum[:]...)
+	iv := md5.Sum(data[2:])
+
+	return key, iv[:]
+}
+
+func makeEncrypterDecrypter(key, iv []byte) (cipher.BlockMode, cipher.BlockMode) {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		panic(err)
+	}
+
+	return cipher.NewCBCEncrypter(block, iv), cipher.NewCBCDecrypter(block, iv)
+}

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,
+	}
+}