From 1894aa798916718c0e98b76844c28642576b5fbc Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Wed, 4 Jul 2018 18:25:15 +0300
Subject: [PATCH] Refactorings of rpc

---
 mtproto/rpc/rpc.go                |  6 +++
 mtproto/wrappers/proxy_request.go | 73 +++++++++++++++++--------------
 2 files changed, 47 insertions(+), 32 deletions(-)

diff --git a/mtproto/rpc/rpc.go b/mtproto/rpc/rpc.go
index 8f2ae99..575b9b7 100644
--- a/mtproto/rpc/rpc.go
+++ b/mtproto/rpc/rpc.go
@@ -4,3 +4,9 @@ const (
 	RPCNonceSeqNo     = -2
 	RPCHandshakeSeqNo = -1
 )
+
+var (
+	RPCTagCloseExt  = []byte{0xa2, 0x34, 0xb6, 0x5e}
+	RPCTagProxyAns  = []byte{0x0d, 0xda, 0x03, 0x44}
+	RPCTagSimpleAck = []byte{0x9b, 0x40, 0xac, 0x3b}
+)
diff --git a/mtproto/wrappers/proxy_request.go b/mtproto/wrappers/proxy_request.go
index d048a85..5b1b269 100644
--- a/mtproto/wrappers/proxy_request.go
+++ b/mtproto/wrappers/proxy_request.go
@@ -13,12 +13,6 @@ import (
 	"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
 
@@ -35,38 +29,53 @@ func (p *ProxyRequestReadWriteCloserWithAddr) Read(buf []byte) (int, error) {
 			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
+		if bytes.Equal(ansBuf.Bytes(), rpc.RPCTagCloseExt) {
+			return p.readCloseExt()
+		} else if bytes.Equal(ansBuf.Bytes(), rpc.RPCTagProxyAns) {
+			return p.readProxyAns(buf)
+		} else if bytes.Equal(ansBuf.Bytes(), rpc.RPCTagSimpleAck) {
+			return p.readSimpleAck()
 		}
 
 		return nil
 	})
 }
 
+func (p *ProxyRequestReadWriteCloserWithAddr) readCloseExt() error {
+	return errors.New("Connection has been closed remotely")
+}
+
+func (p *ProxyRequestReadWriteCloserWithAddr) readProxyAns(buf []byte) error {
+	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
+}
+
+func (p *ProxyRequestReadWriteCloserWithAddr) readSimpleAck() error {
+	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
+}
+
 func (p *ProxyRequestReadWriteCloserWithAddr) Write(raw []byte) (int, error) {
 	if _, err := p.conn.Write(p.req.Bytes(raw)); err != nil {
 		return 0, err