From 4264fdadf84b1cc61a68caa2c37c6fa501918317 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Thu, 5 Jul 2018 09:20:28 +0300
Subject: [PATCH] Add abridged rwc

---
 mtproto/wrappers/abridged.go | 120 +++++++++++++++++++++++++++++++++++
 telegram/middle.go           |   6 +-
 wrappers/buffered_reader.go  |  16 ++---
 3 files changed, 128 insertions(+), 14 deletions(-)
 create mode 100644 mtproto/wrappers/abridged.go

diff --git a/mtproto/wrappers/abridged.go b/mtproto/wrappers/abridged.go
new file mode 100644
index 0000000..db5a99e
--- /dev/null
+++ b/mtproto/wrappers/abridged.go
@@ -0,0 +1,120 @@
+package wrappers
+
+import (
+	"bytes"
+	"encoding/binary"
+	"io"
+	"net"
+
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+type uint24 [3]byte
+
+const (
+	abridgedSmallPacketLength = 0x7f
+	abridgedQuickAckLength    = 0x80
+	abridgedLargePacketLength = 16777216 // 256 ^ 3
+)
+
+type AbridgedReadWriteCloserAddr struct {
+	wrappers.BufferedReader
+
+	conn wrappers.ReadWriteCloserWithAddr
+	opts *mtproto.ConnectionOpts
+}
+
+func (a *AbridgedReadWriteCloserAddr) Read(p []byte) (int, error) {
+	return a.BufferedRead(p, func() error {
+		var msgLength uint8
+		if err := binary.Read(a.conn, binary.LittleEndian, &msgLength); err != nil {
+			return errors.Annotate(err, "Cannot read message length")
+		}
+
+		a.opts.QuickAck = false
+		if msgLength >= abridgedQuickAckLength {
+			a.opts.QuickAck = true
+			msgLength -= 0x80
+		}
+
+		msgLength32 := uint32(msgLength)
+		if msgLength == abridgedSmallPacketLength {
+			buf := &bytes.Buffer{}
+			buf.Grow(3)
+
+			if _, err := io.CopyN(buf, a.conn, 3); err != nil {
+				return errors.Annotate(err, "Cannot read correct message length")
+			}
+			number := uint24{}
+			copy(number[:], buf.Bytes())
+			msgLength32 = fromUint24(number)
+		}
+		msgLength32 *= 4
+
+		if _, err := io.CopyN(a.Buffer, a.conn, int64(msgLength32)); err != nil {
+			return errors.Annotate(err, "Cannot read message")
+		}
+
+		return nil
+	})
+}
+
+func (a *AbridgedReadWriteCloserAddr) Write(p []byte) (int, error) {
+	if len(p)%4 != 0 {
+		return 0, errors.Errorf("Incorrect packet length %d", len(p))
+	}
+	if a.opts.SimpleAck {
+		return a.conn.Write(reverseBytes(p))
+	}
+
+	packetLength := len(p) / 4
+	switch {
+	case packetLength < abridgedSmallPacketLength:
+		newData := append([]byte{byte(packetLength)}, p...)
+		return a.conn.Write(newData)
+
+	case packetLength < abridgedLargePacketLength:
+		length24 := toUint24(uint32(packetLength))
+
+		buf := &bytes.Buffer{}
+		buf.Grow(1 + 3 + len(p))
+		buf.WriteByte(byte(abridgedSmallPacketLength))
+		buf.Write(length24[:])
+		buf.Write(p)
+
+		return a.conn.Write(buf.Bytes())
+
+	default:
+		return 0, errors.Errorf("Packet is too big %d", len(p))
+	}
+}
+
+func (a *AbridgedReadWriteCloserAddr) Close() error {
+	return a.conn.Close()
+}
+
+func (a *AbridgedReadWriteCloserAddr) LocalAddr() *net.TCPAddr {
+	return a.conn.LocalAddr()
+}
+
+func (a *AbridgedReadWriteCloserAddr) RemoteAddr() *net.TCPAddr {
+	return a.conn.RemoteAddr()
+}
+
+func toUint24(number uint32) uint24 {
+	return uint24{byte(number), byte(number >> 8), byte(number >> 16)}
+}
+
+func fromUint24(number uint24) uint32 {
+	return uint32(number[0]) + (uint32(number[1]) << 8) + (uint32(number[2]) << 16)
+}
+
+func NewAbridgedRWC(conn wrappers.ReadWriteCloserWithAddr, connOpts *mtproto.ConnectionOpts) wrappers.ReadWriteCloserWithAddr {
+	return &AbridgedReadWriteCloserAddr{
+		conn: conn,
+		opts: connOpts,
+	}
+}
diff --git a/telegram/middle.go b/telegram/middle.go
index 91f22d5..8e1bd7b 100644
--- a/telegram/middle.go
+++ b/telegram/middle.go
@@ -20,7 +20,7 @@ import (
 type middleTelegram struct {
 	middleTelegramCaller
 
-	adtag []byte
+	conf *config.Config
 }
 
 func NewMiddleTelegram(conf *config.Config, logger *zap.SugaredLogger) Telegram {
@@ -38,7 +38,7 @@ func NewMiddleTelegram(conf *config.Config, logger *zap.SugaredLogger) Telegram
 			},
 			dialerMutex: &sync.RWMutex{},
 		},
-		adtag: conf.AdTag,
+		conf: conf,
 	}
 
 	if err := tg.update(); err != nil {
@@ -73,7 +73,7 @@ func (t *middleTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.Re
 		return nil, err
 	}
 
-	return mtwrappers.NewProxyRequestRWC(secureConn, connOpts, t.adtag)
+	return mtwrappers.NewProxyRequestRWC(secureConn, connOpts, t.conf.AdTag)
 }
 
 func (t *middleTelegram) sendRPCNonceRequest(conn io.Writer) (*rpc.RPCNonceRequest, error) {
diff --git a/wrappers/buffered_reader.go b/wrappers/buffered_reader.go
index a824805..cb48d56 100644
--- a/wrappers/buffered_reader.go
+++ b/wrappers/buffered_reader.go
@@ -17,20 +17,14 @@ func (b *BufferedReader) BufferedRead(p []byte, callback func() error) (int, err
 }
 
 func (b *BufferedReader) flush(p []byte) (int, error) {
-	sizeToRead := len(p)
-	if b.Buffer.Len() < sizeToRead {
-		sizeToRead = b.Buffer.Len()
-	}
-
-	data := b.Buffer.Bytes()
-	copy(p, data[:sizeToRead])
-	if sizeToRead == b.Buffer.Len() {
+	if b.Buffer.Len() < len(p) {
+		sizeToReturn := b.Buffer.Len()
+		copy(p, b.Buffer.Bytes())
 		b.Buffer.Reset()
-	} else {
-		b.Buffer = bytes.NewBuffer(data[sizeToRead:])
+		return sizeToReturn, nil
 	}
 
-	return sizeToRead, nil
+	return b.Buffer.Read(p)
 }
 
 func NewBufferedReader() BufferedReader {