diff --git a/mtproto/bufferpool.go b/mtproto/bufferpool.go deleted file mode 100644 index 2c8c246..0000000 --- a/mtproto/bufferpool.go +++ /dev/null @@ -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 - }, - } -} diff --git a/mtproto/rpc/rpc.go b/mtproto/rpc/rpc.go index f72bfdd..8f2ae99 100644 --- a/mtproto/rpc/rpc.go +++ b/mtproto/rpc/rpc.go @@ -1,12 +1,6 @@ package rpc -import "bytes" - const ( RPCNonceSeqNo = -2 RPCHandshakeSeqNo = -1 ) - -type RPC interface { - Bytes() *bytes.Buffer -} diff --git a/mtproto/rpc/rpc_handshake_request.go b/mtproto/rpc/rpc_handshake_request.go index 662d7a7..12e79e6 100644 --- a/mtproto/rpc/rpc_handshake_request.go +++ b/mtproto/rpc/rpc_handshake_request.go @@ -25,7 +25,7 @@ var ( type RPCHandshakeRequest struct { } -func (r *RPCHandshakeRequest) Bytes() *bytes.Buffer { +func (r *RPCHandshakeRequest) Bytes() []byte { buf := &bytes.Buffer{} buf.Grow(rpcHandshakeRequestLength) @@ -34,7 +34,7 @@ func (r *RPCHandshakeRequest) Bytes() *bytes.Buffer { buf.Write(rpcHandshakeSenderPID[:]) buf.Write(rpcHandshakePeerPID[:]) - return buf + return buf.Bytes() } func init() { diff --git a/mtproto/rpc/rpc_handshake_response.go b/mtproto/rpc/rpc_handshake_response.go index 9af0984..f415840 100644 --- a/mtproto/rpc/rpc_handshake_response.go +++ b/mtproto/rpc/rpc_handshake_response.go @@ -15,7 +15,7 @@ type RPCHandshakeResponse struct { PeerPID [rpcHandshakePeerPIDLength]byte } -func (r *RPCHandshakeResponse) Bytes() *bytes.Buffer { +func (r *RPCHandshakeResponse) Bytes() []byte { buf := &bytes.Buffer{} buf.Grow(rpcHandshakeResponseLength) @@ -24,7 +24,7 @@ func (r *RPCHandshakeResponse) Bytes() *bytes.Buffer { buf.Write(r.SenderPID[:]) buf.Write(r.PeerPID[:]) - return buf + return buf.Bytes() } func (r *RPCHandshakeResponse) Valid(req *RPCHandshakeRequest) error { diff --git a/mtproto/rpc/rpc_nonce_request.go b/mtproto/rpc/rpc_nonce_request.go index b94d814..3fcfe15 100644 --- a/mtproto/rpc/rpc_nonce_request.go +++ b/mtproto/rpc/rpc_nonce_request.go @@ -31,7 +31,7 @@ type RPCNonceRequest struct { Nonce [rpcNonceLength]byte } -func (r *RPCNonceRequest) Bytes() *bytes.Buffer { +func (r *RPCNonceRequest) Bytes() []byte { buf := &bytes.Buffer{} buf.Grow(rpcNonceRequestLength) @@ -41,7 +41,7 @@ func (r *RPCNonceRequest) Bytes() *bytes.Buffer { buf.Write(r.CryptoTS[:]) buf.Write(r.Nonce[:]) - return buf + return buf.Bytes() } func NewRPCNonceRequest(proxySecret []byte) (*RPCNonceRequest, error) { diff --git a/mtproto/rpc/rpc_nonce_response.go b/mtproto/rpc/rpc_nonce_response.go index cad02c3..7cf90fc 100644 --- a/mtproto/rpc/rpc_nonce_response.go +++ b/mtproto/rpc/rpc_nonce_response.go @@ -15,7 +15,7 @@ type RPCNonceResponse struct { Crypto [rpcNonceCryptoAESLength]byte } -func (r *RPCNonceResponse) Bytes() *bytes.Buffer { +func (r *RPCNonceResponse) Bytes() []byte { buf := &bytes.Buffer{} buf.Grow(rpcNonceResponseLength) @@ -25,7 +25,7 @@ func (r *RPCNonceResponse) Bytes() *bytes.Buffer { buf.Write(r.CryptoTS[:]) buf.Write(r.Nonce[:]) - return buf + return buf.Bytes() } func (r *RPCNonceResponse) Valid(req *RPCNonceRequest) error { diff --git a/mtproto/rpc/rpc_proxy_request.go b/mtproto/rpc/rpc_proxy_request.go index e7d0a4a..84edafb 100644 --- a/mtproto/rpc/rpc_proxy_request.go +++ b/mtproto/rpc/rpc_proxy_request.go @@ -32,8 +32,8 @@ type RPCProxyRequest struct { Extras *mtproto.Extras } -func (r *RPCProxyRequest) Bytes() *bytes.Buffer { - buf := mtproto.GetBuffer() +func (r *RPCProxyRequest) Bytes() []byte { + buf := &bytes.Buffer{} flags := r.Flags if r.Extras.QuickAck { @@ -62,7 +62,7 @@ func (r *RPCProxyRequest) Bytes() *bytes.Buffer { buf.Write(messageBytes) } - return buf + return buf.Bytes() } func NewRPCProxyRequest(connectionType mtproto.ConnectionType, local, remote *net.TCPAddr, adTag []byte, extras *mtproto.Extras) (*RPCProxyRequest, error) { diff --git a/mtproto/rwc.go b/mtproto/rwc.go deleted file mode 100644 index ad14b8c..0000000 --- a/mtproto/rwc.go +++ /dev/null @@ -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() -} diff --git a/mtproto/crypt/crypt.go b/mtproto/wrappers/crypt.go similarity index 99% rename from mtproto/crypt/crypt.go rename to mtproto/wrappers/crypt.go index 6530f2b..fbff067 100644 --- a/mtproto/crypt/crypt.go +++ b/mtproto/wrappers/crypt.go @@ -1,4 +1,4 @@ -package crypt +package wrappers import ( "bytes" diff --git a/mtproto/wrappers/frame.go b/mtproto/wrappers/frame.go index b03484b..9e7a4c7 100644 --- a/mtproto/wrappers/frame.go +++ b/mtproto/wrappers/frame.go @@ -8,8 +8,6 @@ import ( "io" "github.com/juju/errors" - - "github.com/9seconds/mtg/mtproto" ) // Frame: { MessageLength(4) | SequenceNumber(4) | Message(???) | CRC32(4) [| padding(4), ...] } @@ -21,35 +19,35 @@ const ( var frameRWCPadding = [4]byte{0x04, 0x00, 0x00, 0x00} type FrameRWC struct { - conn mtproto.BytesRWC + conn io.ReadWriteCloser readSeqNo int32 writeSeqNo int32 readBuf *bytes.Buffer } -func (f *FrameRWC) Write(buf *bytes.Buffer) (int, error) { - writeBuf := mtproto.GetBuffer() - defer mtproto.ReturnBuffer(writeBuf) +func (f *FrameRWC) Write(buf []byte) (int, error) { + writeBuf := &bytes.Buffer{} // 4 - len bytes // 4 - seq bytes // . - message // 4 - crc32 - messageLength := 4 + 4 + buf.Len() + 4 + messageLength := 4 + 4 + len(buf) + 4 paddingLength := (aes.BlockSize - messageLength%aes.BlockSize) % aes.BlockSize writeBuf.Grow(messageLength + paddingLength) binary.Write(writeBuf, binary.LittleEndian, uint32(messageLength)) binary.Write(writeBuf, binary.LittleEndian, f.writeSeqNo) - writeBuf.Write(buf.Bytes()) + writeBuf.Write(buf) f.writeSeqNo++ checksum := crc32.ChecksumIEEE(writeBuf.Bytes()) binary.Write(writeBuf, binary.LittleEndian, checksum) 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) { @@ -57,9 +55,7 @@ func (f *FrameRWC) Read(p []byte) (int, error) { return f.flush(p) } - buf := mtproto.GetBuffer() - defer mtproto.ReturnBuffer(buf) - + buf := &bytes.Buffer{} for { buf.Reset() 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 { - defer mtproto.ReturnBuffer(f.readBuf) return f.conn.Close() } @@ -118,21 +113,17 @@ func (f *FrameRWC) flush(p []byte) (int, error) { if sizeToRead == f.readBuf.Len() { f.readBuf.Reset() } else { - newBuf := mtproto.GetBuffer() - newBuf.Write(data[sizeToRead:]) - - mtproto.ReturnBuffer(f.readBuf) - f.readBuf = newBuf + f.readBuf = bytes.NewBuffer(data[sizeToRead:]) } return sizeToRead, nil } -func NewFrameRWC(conn mtproto.BytesRWC, seqNo int32) mtproto.BytesRWC { +func NewFrameRWC(conn io.ReadWriteCloser, seqNo int32) io.ReadWriteCloser { return &FrameRWC{ conn: conn, readSeqNo: seqNo, writeSeqNo: seqNo, - readBuf: mtproto.GetBuffer(), + readBuf: &bytes.Buffer{}, } }