Remove mtproto pool

This commit is contained in:
9seconds
2018-07-02 15:44:52 +03:00
parent 5ea0c4f7bf
commit 6fa5f31ca8
10 changed files with 23 additions and 118 deletions
-32
View File
@@ -1,32 +0,0 @@
package mtproto
import (
"bytes"
"sync"
)
const bufferPoolSize = 4 * 1024
var bufferPool sync.Pool
func GetBuffer() *bytes.Buffer {
buf := bufferPool.Get().(*bytes.Buffer)
buf.Reset()
return buf
}
func ReturnBuffer(buf *bytes.Buffer) {
bufferPool.Put(buf)
}
func init() {
bufferPool = sync.Pool{
New: func() interface{} {
buf := &bytes.Buffer{}
buf.Grow(bufferPoolSize)
return buf
},
}
}
-6
View File
@@ -1,12 +1,6 @@
package rpc package rpc
import "bytes"
const ( const (
RPCNonceSeqNo = -2 RPCNonceSeqNo = -2
RPCHandshakeSeqNo = -1 RPCHandshakeSeqNo = -1
) )
type RPC interface {
Bytes() *bytes.Buffer
}
+2 -2
View File
@@ -25,7 +25,7 @@ var (
type RPCHandshakeRequest struct { type RPCHandshakeRequest struct {
} }
func (r *RPCHandshakeRequest) Bytes() *bytes.Buffer { func (r *RPCHandshakeRequest) Bytes() []byte {
buf := &bytes.Buffer{} buf := &bytes.Buffer{}
buf.Grow(rpcHandshakeRequestLength) buf.Grow(rpcHandshakeRequestLength)
@@ -34,7 +34,7 @@ func (r *RPCHandshakeRequest) Bytes() *bytes.Buffer {
buf.Write(rpcHandshakeSenderPID[:]) buf.Write(rpcHandshakeSenderPID[:])
buf.Write(rpcHandshakePeerPID[:]) buf.Write(rpcHandshakePeerPID[:])
return buf return buf.Bytes()
} }
func init() { func init() {
+2 -2
View File
@@ -15,7 +15,7 @@ type RPCHandshakeResponse struct {
PeerPID [rpcHandshakePeerPIDLength]byte PeerPID [rpcHandshakePeerPIDLength]byte
} }
func (r *RPCHandshakeResponse) Bytes() *bytes.Buffer { func (r *RPCHandshakeResponse) Bytes() []byte {
buf := &bytes.Buffer{} buf := &bytes.Buffer{}
buf.Grow(rpcHandshakeResponseLength) buf.Grow(rpcHandshakeResponseLength)
@@ -24,7 +24,7 @@ func (r *RPCHandshakeResponse) Bytes() *bytes.Buffer {
buf.Write(r.SenderPID[:]) buf.Write(r.SenderPID[:])
buf.Write(r.PeerPID[:]) buf.Write(r.PeerPID[:])
return buf return buf.Bytes()
} }
func (r *RPCHandshakeResponse) Valid(req *RPCHandshakeRequest) error { func (r *RPCHandshakeResponse) Valid(req *RPCHandshakeRequest) error {
+2 -2
View File
@@ -31,7 +31,7 @@ type RPCNonceRequest struct {
Nonce [rpcNonceLength]byte Nonce [rpcNonceLength]byte
} }
func (r *RPCNonceRequest) Bytes() *bytes.Buffer { func (r *RPCNonceRequest) Bytes() []byte {
buf := &bytes.Buffer{} buf := &bytes.Buffer{}
buf.Grow(rpcNonceRequestLength) buf.Grow(rpcNonceRequestLength)
@@ -41,7 +41,7 @@ func (r *RPCNonceRequest) Bytes() *bytes.Buffer {
buf.Write(r.CryptoTS[:]) buf.Write(r.CryptoTS[:])
buf.Write(r.Nonce[:]) buf.Write(r.Nonce[:])
return buf return buf.Bytes()
} }
func NewRPCNonceRequest(proxySecret []byte) (*RPCNonceRequest, error) { func NewRPCNonceRequest(proxySecret []byte) (*RPCNonceRequest, error) {
+2 -2
View File
@@ -15,7 +15,7 @@ type RPCNonceResponse struct {
Crypto [rpcNonceCryptoAESLength]byte Crypto [rpcNonceCryptoAESLength]byte
} }
func (r *RPCNonceResponse) Bytes() *bytes.Buffer { func (r *RPCNonceResponse) Bytes() []byte {
buf := &bytes.Buffer{} buf := &bytes.Buffer{}
buf.Grow(rpcNonceResponseLength) buf.Grow(rpcNonceResponseLength)
@@ -25,7 +25,7 @@ func (r *RPCNonceResponse) Bytes() *bytes.Buffer {
buf.Write(r.CryptoTS[:]) buf.Write(r.CryptoTS[:])
buf.Write(r.Nonce[:]) buf.Write(r.Nonce[:])
return buf return buf.Bytes()
} }
func (r *RPCNonceResponse) Valid(req *RPCNonceRequest) error { func (r *RPCNonceResponse) Valid(req *RPCNonceRequest) error {
+3 -3
View File
@@ -32,8 +32,8 @@ type RPCProxyRequest struct {
Extras *mtproto.Extras Extras *mtproto.Extras
} }
func (r *RPCProxyRequest) Bytes() *bytes.Buffer { func (r *RPCProxyRequest) Bytes() []byte {
buf := mtproto.GetBuffer() buf := &bytes.Buffer{}
flags := r.Flags flags := r.Flags
if r.Extras.QuickAck { if r.Extras.QuickAck {
@@ -62,7 +62,7 @@ func (r *RPCProxyRequest) Bytes() *bytes.Buffer {
buf.Write(messageBytes) buf.Write(messageBytes)
} }
return buf return buf.Bytes()
} }
func NewRPCProxyRequest(connectionType mtproto.ConnectionType, local, remote *net.TCPAddr, adTag []byte, extras *mtproto.Extras) (*RPCProxyRequest, error) { func NewRPCProxyRequest(connectionType mtproto.ConnectionType, local, remote *net.TCPAddr, adTag []byte, extras *mtproto.Extras) (*RPCProxyRequest, error) {
-48
View File
@@ -1,48 +0,0 @@
package mtproto
import (
"bytes"
"io"
)
type BytesRWC interface {
Write(*bytes.Buffer) (int, error)
Read([]byte) (int, error)
Close() error
}
type StartBytesRWC struct {
conn BytesRWC
}
func (s *StartBytesRWC) Write(p []byte) (int, error) {
buf := GetBuffer()
buf.Write(p)
defer ReturnBuffer(buf)
return s.conn.Write(buf)
}
func (s *StartBytesRWC) Read(p []byte) (int, error) {
return s.conn.Read(p)
}
func (s *StartBytesRWC) Close() error {
return s.conn.Close()
}
type FinishBytesRWC struct {
conn io.ReadWriteCloser
}
func (f *FinishBytesRWC) Write(buf *bytes.Buffer) (int, error) {
return f.conn.Write(buf.Bytes())
}
func (f *FinishBytesRWC) Read(p []byte) (int, error) {
return f.conn.Read(p)
}
func (f *FinishBytesRWC) Close() error {
return f.conn.Close()
}
@@ -1,4 +1,4 @@
package crypt package wrappers
import ( import (
"bytes" "bytes"
+11 -20
View File
@@ -8,8 +8,6 @@ import (
"io" "io"
"github.com/juju/errors" "github.com/juju/errors"
"github.com/9seconds/mtg/mtproto"
) )
// Frame: { MessageLength(4) | SequenceNumber(4) | Message(???) | CRC32(4) [| padding(4), ...] } // Frame: { MessageLength(4) | SequenceNumber(4) | Message(???) | CRC32(4) [| padding(4), ...] }
@@ -21,35 +19,35 @@ const (
var frameRWCPadding = [4]byte{0x04, 0x00, 0x00, 0x00} var frameRWCPadding = [4]byte{0x04, 0x00, 0x00, 0x00}
type FrameRWC struct { type FrameRWC struct {
conn mtproto.BytesRWC conn io.ReadWriteCloser
readSeqNo int32 readSeqNo int32
writeSeqNo int32 writeSeqNo int32
readBuf *bytes.Buffer readBuf *bytes.Buffer
} }
func (f *FrameRWC) Write(buf *bytes.Buffer) (int, error) { func (f *FrameRWC) Write(buf []byte) (int, error) {
writeBuf := mtproto.GetBuffer() writeBuf := &bytes.Buffer{}
defer mtproto.ReturnBuffer(writeBuf)
// 4 - len bytes // 4 - len bytes
// 4 - seq bytes // 4 - seq bytes
// . - message // . - message
// 4 - crc32 // 4 - crc32
messageLength := 4 + 4 + buf.Len() + 4 messageLength := 4 + 4 + len(buf) + 4
paddingLength := (aes.BlockSize - messageLength%aes.BlockSize) % aes.BlockSize paddingLength := (aes.BlockSize - messageLength%aes.BlockSize) % aes.BlockSize
writeBuf.Grow(messageLength + paddingLength) writeBuf.Grow(messageLength + paddingLength)
binary.Write(writeBuf, binary.LittleEndian, uint32(messageLength)) binary.Write(writeBuf, binary.LittleEndian, uint32(messageLength))
binary.Write(writeBuf, binary.LittleEndian, f.writeSeqNo) binary.Write(writeBuf, binary.LittleEndian, f.writeSeqNo)
writeBuf.Write(buf.Bytes()) writeBuf.Write(buf)
f.writeSeqNo++ f.writeSeqNo++
checksum := crc32.ChecksumIEEE(writeBuf.Bytes()) checksum := crc32.ChecksumIEEE(writeBuf.Bytes())
binary.Write(writeBuf, binary.LittleEndian, checksum) binary.Write(writeBuf, binary.LittleEndian, checksum)
writeBuf.Write(bytes.Repeat(frameRWCPadding[:], paddingLength/4)) writeBuf.Write(bytes.Repeat(frameRWCPadding[:], paddingLength/4))
return f.conn.Write(writeBuf) _, err := f.conn.Write(writeBuf.Bytes())
return len(buf), err
} }
func (f *FrameRWC) Read(p []byte) (int, error) { func (f *FrameRWC) Read(p []byte) (int, error) {
@@ -57,9 +55,7 @@ func (f *FrameRWC) Read(p []byte) (int, error) {
return f.flush(p) return f.flush(p)
} }
buf := mtproto.GetBuffer() buf := &bytes.Buffer{}
defer mtproto.ReturnBuffer(buf)
for { for {
buf.Reset() buf.Reset()
if _, err := io.CopyN(buf, f.conn, 4); err != nil { if _, err := io.CopyN(buf, f.conn, 4); err != nil {
@@ -103,7 +99,6 @@ func (f *FrameRWC) Read(p []byte) (int, error) {
} }
func (f *FrameRWC) Close() error { func (f *FrameRWC) Close() error {
defer mtproto.ReturnBuffer(f.readBuf)
return f.conn.Close() return f.conn.Close()
} }
@@ -118,21 +113,17 @@ func (f *FrameRWC) flush(p []byte) (int, error) {
if sizeToRead == f.readBuf.Len() { if sizeToRead == f.readBuf.Len() {
f.readBuf.Reset() f.readBuf.Reset()
} else { } else {
newBuf := mtproto.GetBuffer() f.readBuf = bytes.NewBuffer(data[sizeToRead:])
newBuf.Write(data[sizeToRead:])
mtproto.ReturnBuffer(f.readBuf)
f.readBuf = newBuf
} }
return sizeToRead, nil return sizeToRead, nil
} }
func NewFrameRWC(conn mtproto.BytesRWC, seqNo int32) mtproto.BytesRWC { func NewFrameRWC(conn io.ReadWriteCloser, seqNo int32) io.ReadWriteCloser {
return &FrameRWC{ return &FrameRWC{
conn: conn, conn: conn,
readSeqNo: seqNo, readSeqNo: seqNo,
writeSeqNo: seqNo, writeSeqNo: seqNo,
readBuf: mtproto.GetBuffer(), readBuf: &bytes.Buffer{},
} }
} }