From 96bb6fbe980fb64d82ae941161c5079a75f9d054 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Mon, 2 Jul 2018 09:15:03 +0300
Subject: [PATCH] Add rpc proxy request

---
 mtproto/bufferpool/bufferpool.go | 32 +++++++++++
 mtproto/extras.go                |  6 ++
 mtproto/rpc/rpc_proxy_flags.go   | 29 +---------
 mtproto/rpc/rpc_proxy_request.go | 98 ++++++++++++++++++++++++++++++++
 4 files changed, 137 insertions(+), 28 deletions(-)
 create mode 100644 mtproto/bufferpool/bufferpool.go
 create mode 100644 mtproto/extras.go
 create mode 100644 mtproto/rpc/rpc_proxy_request.go

diff --git a/mtproto/bufferpool/bufferpool.go b/mtproto/bufferpool/bufferpool.go
new file mode 100644
index 0000000..f58a88e
--- /dev/null
+++ b/mtproto/bufferpool/bufferpool.go
@@ -0,0 +1,32 @@
+package bufferpool
+
+import (
+	"bytes"
+	"sync"
+)
+
+const bufferPoolSize = 4 * 1024
+
+var bufferPool sync.Pool
+
+func Get() *bytes.Buffer {
+	buf := bufferPool.Get().(*bytes.Buffer)
+	buf.Reset()
+
+	return buf
+}
+
+func Return(buf *bytes.Buffer) {
+	bufferPool.Put(buf)
+}
+
+func init() {
+	bufferPool = sync.Pool{
+		New: func() interface{} {
+			buf := &bytes.Buffer{}
+			buf.Grow(bufferPoolSize)
+
+			return buf
+		},
+	}
+}
diff --git a/mtproto/extras.go b/mtproto/extras.go
new file mode 100644
index 0000000..5c0ad19
--- /dev/null
+++ b/mtproto/extras.go
@@ -0,0 +1,6 @@
+package mtproto
+
+type Extras struct {
+	QuickAck  bool
+	SimpleAck bool
+}
diff --git a/mtproto/rpc/rpc_proxy_flags.go b/mtproto/rpc/rpc_proxy_flags.go
index 8e0f202..4edd74d 100644
--- a/mtproto/rpc/rpc_proxy_flags.go
+++ b/mtproto/rpc/rpc_proxy_flags.go
@@ -1,11 +1,6 @@
 package rpc
 
-import (
-	"bytes"
-	"encoding/binary"
-
-	"github.com/9seconds/mtg/mtproto"
-)
+import "encoding/binary"
 
 type RPCProxyRequestFlags uint32
 
@@ -27,25 +22,3 @@ func (r RPCProxyRequestFlags) Bytes() []byte {
 
 	return converted
 }
-
-func NewRPCRproxyRequestFlags(connectionType mtproto.ConnectionType, quickAck bool, message []byte) RPCProxyRequestFlags {
-	flags := RPCProxyRequestFlagsHasAdTag
-	flags |= RPCProxyRequestFlagsMagic
-	flags |= RPCProxyRequestFlagsExtMode2
-
-	switch connectionType {
-	case mtproto.ConnectionTypeAbridged:
-		flags |= RPCProxyRequestFlagsAbdridged
-	case mtproto.ConnectionTypeIntermediate:
-		flags |= RPCProxyRequestFlagsIntermediate
-	}
-
-	if quickAck {
-		flags |= RPCProxyRequestFlagsQuickAck
-	}
-	if bytes.HasPrefix(message, rpcProxyRequestFlagsEncryptedPrefix[:]) {
-		flags |= RPCProxyRequestFlagsEncrypted
-	}
-
-	return flags
-}
diff --git a/mtproto/rpc/rpc_proxy_request.go b/mtproto/rpc/rpc_proxy_request.go
new file mode 100644
index 0000000..f9df232
--- /dev/null
+++ b/mtproto/rpc/rpc_proxy_request.go
@@ -0,0 +1,98 @@
+package rpc
+
+import (
+	"bytes"
+	"crypto/rand"
+	"encoding/binary"
+	"net"
+
+	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/mtproto/bufferpool"
+	"github.com/juju/errors"
+)
+
+const (
+	rpcProxyRequestConnectionIDLength = 8
+	rpcProxyRequestIPPortLength       = 16 + 4
+)
+
+var (
+	rpcProxyRequestTag       = []byte{0xee, 0xf1, 0xce, 0x36}
+	rpcProxyRequestExtraSize = []byte{0x18, 0x00, 0x00, 0x00}
+	rpcProxyRequestProxyTag  = []byte{0xae, 0x26, 0x1e, 0xdb}
+)
+
+type RPCProxyRequest struct {
+	Flags        RPCProxyRequestFlags
+	ConnectionID [rpcProxyRequestConnectionIDLength]byte
+	RemoteIPPort [rpcProxyRequestIPPortLength]byte
+	LocalIPPort  [rpcProxyRequestIPPortLength]byte
+	ADTag        []byte
+	Message      *bytes.Buffer
+	Extras       *mtproto.Extras
+}
+
+func (r *RPCProxyRequest) Bytes() *bytes.Buffer {
+	buf := bufferpool.Get()
+
+	flags := r.Flags
+	if r.Extras.QuickAck {
+		flags |= RPCProxyRequestFlagsQuickAck
+	}
+
+	messageBytes := r.Message.Bytes()
+	if bytes.HasPrefix(messageBytes, rpcProxyRequestFlagsEncryptedPrefix[:]) {
+		flags |= RPCProxyRequestFlagsEncrypted
+	}
+
+	buf.Write(rpcProxyRequestTag)
+	buf.Write(flags.Bytes())
+	buf.Write(r.ConnectionID[:])
+	buf.Write(r.RemoteIPPort[:])
+	buf.Write(r.LocalIPPort[:])
+	buf.Write(rpcProxyRequestExtraSize)
+	buf.Write(rpcProxyRequestProxyTag)
+	buf.WriteByte(byte(len(r.ADTag)))
+	buf.Write(r.ADTag)
+
+	for i := 0; i < (buf.Len() % 4); i++ {
+		buf.WriteByte(0x00)
+	}
+	if r.Message != nil {
+		buf.Write(messageBytes)
+	}
+
+	return buf
+}
+
+func NewRPCProxyRequest(connectionType mtproto.ConnectionType, local, remote *net.TCPAddr, adTag []byte, extras *mtproto.Extras) (*RPCProxyRequest, error) {
+	flags := RPCProxyRequestFlagsHasAdTag | RPCProxyRequestFlagsMagic | RPCProxyRequestFlagsExtMode2
+
+	switch connectionType {
+	case mtproto.ConnectionTypeAbridged:
+		flags |= RPCProxyRequestFlagsAbdridged
+	case mtproto.ConnectionTypeIntermediate:
+		flags |= RPCProxyRequestFlagsIntermediate
+	}
+
+	request := RPCProxyRequest{
+		Flags:  flags,
+		ADTag:  adTag,
+		Extras: extras,
+	}
+
+	if _, err := rand.Read(request.ConnectionID[:]); err != nil {
+		return nil, errors.Annotate(err, "Cannot generate connection ID")
+	}
+
+	port := make([]byte, 4)
+	copy(request.LocalIPPort[:], local.IP.To16())
+	binary.LittleEndian.PutUint32(port, uint32(local.Port))
+	copy(request.LocalIPPort[16:], port)
+
+	copy(request.RemoteIPPort[:], remote.IP.To16())
+	binary.LittleEndian.PutUint32(port, uint32(remote.Port))
+	copy(request.RemoteIPPort[16:], port)
+
+	return &request, nil
+}