From 69c4997fd4657c829c3bf375812cee7fa878a1dc Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Wed, 4 Jul 2018 13:28:23 +0300
Subject: [PATCH] Add proxy request wrapper

---
 client/direct.go                  |   1 +
 mtproto/connection_options.go     |   2 +
 mtproto/wrappers/proxy_request.go | 102 ++++++++++++++++++++++++++++++
 3 files changed, 105 insertions(+)
 create mode 100644 mtproto/wrappers/proxy_request.go

diff --git a/client/direct.go b/client/direct.go
index 1e38e5e..98b518d 100644
--- a/client/direct.go
+++ b/client/direct.go
@@ -32,6 +32,7 @@ func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, wr
 		return nil, nil, errors.Annotate(err, "Cannot parse obfuscated frame")
 	}
 	connOpts.ConnectionProto = mtproto.ConnectionProtocolAny
+	connOpts.ClientAddr = conn.RemoteAddr().(*net.TCPAddr)
 
 	socket := wrappers.NewTimeoutRWC(conn, conf.PublicIPv4, conf.PublicIPv6)
 	socket = wrappers.NewStreamCipherRWC(socket, obfs2.Encryptor, obfs2.Decryptor)
diff --git a/mtproto/connection_options.go b/mtproto/connection_options.go
index 823b81b..585e0f7 100644
--- a/mtproto/connection_options.go
+++ b/mtproto/connection_options.go
@@ -2,6 +2,7 @@ package mtproto
 
 import (
 	"bytes"
+	"net"
 
 	"github.com/juju/errors"
 )
@@ -20,6 +21,7 @@ type ConnectionOpts struct {
 	ConnectionProto ConnectionProtocol
 	QuickAck        bool
 	SimpleAck       bool
+	ClientAddr      *net.TCPAddr
 }
 
 // Different connection types which user requests from Telegram.
diff --git a/mtproto/wrappers/proxy_request.go b/mtproto/wrappers/proxy_request.go
new file mode 100644
index 0000000..d048a85
--- /dev/null
+++ b/mtproto/wrappers/proxy_request.go
@@ -0,0 +1,102 @@
+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/wrappers"
+)
+
+var (
+	rpcCloseExtTag  = [4]byte{0xa2, 0x34, 0xb6, 0x5e}
+	rpcProxyAnsTag  = [4]byte{0x0d, 0xda, 0x03, 0x44}
+	rpcSimpleAckTag = [4]byte{0x9b, 0x40, 0xac, 0x3b}
+)
+
+type ProxyRequestReadWriteCloserWithAddr struct {
+	wrappers.BufferedReader
+
+	conn wrappers.ReadWriteCloserWithAddr
+	req  *rpc.RPCProxyRequest
+}
+
+func (p *ProxyRequestReadWriteCloserWithAddr) Read(buf []byte) (int, error) {
+	return p.BufferedRead(buf, func() error {
+		ansBuf := &bytes.Buffer{}
+		ansBuf.Grow(4)
+
+		if _, err := io.CopyN(ansBuf, p.conn, 4); err != nil {
+			return errors.Annotate(err, "Cannot read RPC tag")
+		}
+
+		if bytes.Equal(ansBuf.Bytes(), rpcCloseExtTag[:]) {
+			return errors.New("Connection has been closed remotely")
+		} else if bytes.Equal(ansBuf.Bytes(), rpcProxyAnsTag[:]) {
+			if _, err := io.CopyN(ioutil.Discard, p.conn, 8+4); err != nil {
+				return errors.Annotate(err, "Cannot skip flags and connid")
+			}
+			for {
+				n, err := p.conn.Read(buf)
+				if err != nil {
+					return errors.Annotate(err, "Cannot read proxy answer")
+				}
+				if n == 0 {
+					break
+				}
+				p.Buffer.Write(buf[:n])
+			}
+			return nil
+		} else if bytes.Equal(ansBuf.Bytes(), rpcSimpleAckTag[:]) {
+			if _, err := io.CopyN(ioutil.Discard, p.conn, 8); err != nil {
+				return errors.Annotate(err, "Cannot skip connid")
+			}
+			if _, err := io.CopyN(p.Buffer, p.conn, 4); err != nil {
+				return errors.Annotate(err, "Cannot read simple ack")
+			}
+			p.req.Options.SimpleAck = true
+			return nil
+		}
+
+		return nil
+	})
+}
+
+func (p *ProxyRequestReadWriteCloserWithAddr) Write(raw []byte) (int, error) {
+	if _, err := p.conn.Write(p.req.Bytes(raw)); err != nil {
+		return 0, err
+	}
+	p.req.Options.SimpleAck = false
+	p.req.Options.QuickAck = false
+
+	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 NewProxyRequestRWC(conn wrappers.ReadWriteCloserWithAddr, connOpts *mtproto.ConnectionOpts, adTag []byte) (wrappers.ReadWriteCloserWithAddr, error) {
+	req, err := rpc.NewRPCProxyRequest(connOpts.ClientAddr, conn.LocalAddr(), connOpts, adTag)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot create new RPC proxy request")
+	}
+
+	return &ProxyRequestReadWriteCloserWithAddr{
+		conn: conn,
+		req:  req,
+	}, nil
+}