From 5f54a235d28eca7a89065954dfb1dc58305efb07 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Mon, 2 Jul 2018 17:51:19 +0300
Subject: [PATCH] Add base of rpc handshake

---
 mtproto/wrappers/crypt.go |  4 +-
 telegram/middle.go        | 81 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+), 2 deletions(-)

diff --git a/mtproto/wrappers/crypt.go b/mtproto/wrappers/crypt.go
index 7ae5255..4150ec2 100644
--- a/mtproto/wrappers/crypt.go
+++ b/mtproto/wrappers/crypt.go
@@ -23,8 +23,8 @@ const (
 var emptyIP = [4]byte{0x00, 0x00, 0x00, 0x00}
 
 func NewMiddleProxyCipherRWC(conn wrappers.ReadWriteCloserWithAddr, req *rpc.RPCNonceRequest,
-	resp *rpc.RPCNonceResponse, client *net.TCPAddr, remote *net.TCPAddr,
-	secret []byte) wrappers.ReadWriteCloserWithAddr {
+	resp *rpc.RPCNonceResponse, client *net.TCPAddr, secret []byte) wrappers.ReadWriteCloserWithAddr {
+	remote := conn.Addr()
 	encryptor := newCBCCipher(CipherPurposeClient, req, resp, client, remote, secret)
 	decryptor := newCBCCipher(CipherPurposeServer, req, resp, client, remote, secret)
 
diff --git a/telegram/middle.go b/telegram/middle.go
index 4f38b8f..65e4af2 100644
--- a/telegram/middle.go
+++ b/telegram/middle.go
@@ -1,6 +1,7 @@
 package telegram
 
 import (
+	"io/ioutil"
 	"net"
 	"net/http"
 	"sync"
@@ -9,7 +10,10 @@ import (
 
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/mtproto/rpc"
+	mtwrappers "github.com/9seconds/mtg/mtproto/wrappers"
 	"github.com/9seconds/mtg/wrappers"
+	"github.com/juju/errors"
 )
 
 type middleTelegram struct {
@@ -39,5 +43,82 @@ func NewMiddleTelegram(conf *config.Config, logger *zap.SugaredLogger) Telegram
 }
 
 func (t *middleTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.ReadWriteCloserWithAddr) (wrappers.ReadWriteCloserWithAddr, error) {
+	rpcNonceConn := mtwrappers.NewFrameRWC(conn, rpc.RPCNonceSeqNo)
+
+	rpcNonceReq, err := t.sendRPCNonceRequest(rpcNonceConn)
+	if err != nil {
+		return nil, err
+	}
+	rpcNonceResp, err := t.receiveRPCNonceResponse(rpcNonceConn, rpcNonceReq)
+	if err != nil {
+		return nil, err
+	}
+
+	secureConn := mtwrappers.NewFrameRWC(conn, rpc.RPCHandshakeSeqNo)
+	secureConn = mtwrappers.NewMiddleProxyCipherRWC(conn, rpcNonceReq,
+		rpcNonceResp, connOpts.ClientAddr, t.proxySecret)
+
+	rpcHandshakeReq, err := t.sendRPCHandshakeRequest(secureConn)
+	if err != nil {
+		return nil, err
+	}
+	_, err = t.receiveRPCHandshakeResponse(secureConn, rpcHandshakeReq)
+	if err != nil {
+		return nil, err
+	}
+
 	return nil, nil
 }
+
+func (t *middleTelegram) sendRPCNonceRequest(conn wrappers.ReadWriteCloserWithAddr) (*rpc.RPCNonceRequest, error) {
+	rpcNonceReq, err := rpc.NewRPCNonceRequest(t.proxySecret)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot create RPC nonce request")
+	}
+	if _, err = conn.Write(rpcNonceReq.Bytes()); err != nil {
+		return nil, errors.Annotate(err, "Cannot send RPC nonce request")
+	}
+
+	return rpcNonceReq, nil
+}
+
+func (t *middleTelegram) receiveRPCNonceResponse(conn wrappers.ReadWriteCloserWithAddr, req *rpc.RPCNonceRequest) (*rpc.RPCNonceResponse, error) {
+	ans, err := ioutil.ReadAll(conn)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot read RPC nonce response")
+	}
+	rpcNonceResp, err := rpc.NewRPCNonceResponse(ans)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot initialize RPC nonce response")
+	}
+	if err = rpcNonceResp.Valid(req); err != nil {
+		return nil, errors.Annotate(err, "Invalid RPC nonce response")
+	}
+
+	return rpcNonceResp, nil
+}
+
+func (t *middleTelegram) sendRPCHandshakeRequest(conn wrappers.ReadWriteCloserWithAddr) (*rpc.RPCHandshakeRequest, error) {
+	req := rpc.NewRPCHandshakeRequest()
+	if _, err := conn.Write(req.Bytes()); err != nil {
+		return nil, errors.Annotate(err, "Cannot send RPC handshake request")
+	}
+
+	return req, nil
+}
+
+func (t *middleTelegram) receiveRPCHandshakeResponse(conn wrappers.ReadWriteCloserWithAddr, req *rpc.RPCHandshakeRequest) (*rpc.RPCHandshakeResponse, error) {
+	ans, err := ioutil.ReadAll(conn)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot read RPC handshake response")
+	}
+	rpcHandshakeResp, err := rpc.NewRPCHandshakeResponse(ans)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot initialize RPC handshake response")
+	}
+	if err = rpcHandshakeResp.Valid(req); err != nil {
+		return nil, errors.Annotate(err, "Invalid RPC handshake response")
+	}
+
+	return rpcHandshakeResp, nil
+}