From 76e88edac342ffdfe41b53b5f11400c7b9308f5e Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Sun, 8 Jul 2018 13:19:16 +0300
Subject: [PATCH] Send correct proxy header

---
 client/direct.go             |  2 ++
 mtproto/rpc/proxy_request.go | 17 ++++++++++++++---
 telegram/middle.go           |  8 +++++++-
 wrappers/mtproto_proxy.go    | 15 ++++++++++++++-
 4 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/client/direct.go b/client/direct.go
index 2c16436..9f0df10 100644
--- a/client/direct.go
+++ b/client/direct.go
@@ -36,5 +36,7 @@ func DirectInit(socket net.Conn, connID string, conf *config.Config) (wrappers.W
 
 	conn = wrappers.NewStreamCipher(conn, obfs2.Encryptor, obfs2.Decryptor)
 
+	conn.Logger().Infow("Client connection initialized")
+
 	return conn, connOpts, nil
 }
diff --git a/mtproto/rpc/proxy_request.go b/mtproto/rpc/proxy_request.go
index e36a2ec..0895f1a 100644
--- a/mtproto/rpc/proxy_request.go
+++ b/mtproto/rpc/proxy_request.go
@@ -20,8 +20,20 @@ type ProxyRequest struct {
 	Options      *mtproto.ConnectionOpts
 }
 
-func (r *ProxyRequest) Bytes(message []byte) []byte {
+func (r *ProxyRequest) MakeHeader(message []byte) *bytes.Buffer {
+	bufferLength := len(TagProxyRequest) +
+		4 + // len(flags)
+		len(r.ConnectionID) +
+		len(r.ClientIPPort) +
+		len(r.OurIPPort) +
+		len(ProxyRequestExtraSize) +
+		len(ProxyRequestProxyTag) +
+		1 + // len(AdTag)
+		len(r.ADTag)
+	bufferLength += bufferLength % 4
+
 	buf := &bytes.Buffer{}
+	buf.Grow(bufferLength)
 
 	flags := r.Flags
 	if r.Options.ReadHacks.QuickAck {
@@ -42,9 +54,8 @@ func (r *ProxyRequest) Bytes(message []byte) []byte {
 	buf.WriteByte(byte(len(r.ADTag)))
 	buf.Write(r.ADTag)
 	buf.Write(make([]byte, (4-buf.Len()%4)%4))
-	buf.Write(message)
 
-	return buf.Bytes()
+	return buf
 }
 
 func NewProxyRequest(clientAddr, ownAddr *net.TCPAddr, opts *mtproto.ConnectionOpts, adTag []byte) (*ProxyRequest, error) {
diff --git a/telegram/middle.go b/telegram/middle.go
index b800860..7404e3d 100644
--- a/telegram/middle.go
+++ b/telegram/middle.go
@@ -43,7 +43,13 @@ func (t *MiddleTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.St
 		return nil, err
 	}
 
-	return wrappers.NewMTProtoProxy(frameConn, connOpts, t.conf.AdTag)
+	proxyConn, err := wrappers.NewMTProtoProxy(frameConn, connOpts, t.conf.AdTag)
+	if err != nil {
+		return nil, err
+	}
+	proxyConn.Logger().Infow("Telegram connection initialized")
+
+	return proxyConn, nil
 }
 
 func (t *MiddleTelegram) sendRPCNonceRequest(conn wrappers.PacketWriter) (*rpc.NonceRequest, error) {
diff --git a/wrappers/mtproto_proxy.go b/wrappers/mtproto_proxy.go
index 0a808a3..1a82eb4 100644
--- a/wrappers/mtproto_proxy.go
+++ b/wrappers/mtproto_proxy.go
@@ -2,6 +2,7 @@ package wrappers
 
 import (
 	"bytes"
+	"fmt"
 	"net"
 
 	"github.com/juju/errors"
@@ -100,7 +101,19 @@ func (m *MTProtoProxy) Write(p []byte) (int, error) {
 	)
 	m.writeCounter++
 
-	if _, err := m.conn.Write(p); err != nil {
+	header := m.req.MakeHeader(p)
+	if ce := m.logger.Desugar().Check(zap.DebugLevel, "RPC_PROXY_REQ header"); ce != nil {
+		ce.Write(
+			zap.Int("length", len(p)),
+			zap.Uint32("counter", m.writeCounter),
+			zap.Bool("simple_ack", m.req.Options.ReadHacks.QuickAck),
+			zap.Bool("quick_ack", m.req.Options.ReadHacks.SimpleAck),
+			zap.String("header", fmt.Sprintf("%v", header.Bytes())),
+		)
+	}
+	header.Write(p)
+
+	if _, err := m.conn.Write(header.Bytes()); err != nil {
 		return 0, err
 	}