ScuroNeko/mtg
1
0
Fork 0
mirror of https://github.com/ScuroNeko/mtg.git synced 2026-08-31 15:44:01 +03:00
Code Issues Packages Projects Releases Wiki Activity

Remove mtproto pool

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
mtproto/bufferpool.go
-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
},
}
}
mtproto/rpc/rpc.go
-6
View File
@@ -1,12 +1,6 @@
package rpc
import "bytes"
const (
RPCNonceSeqNo = -2
RPCHandshakeSeqNo = -1
)
type RPC interface {
Bytes() *bytes.Buffer
}
mtproto/rpc/rpc_handshake_request.go
+2 -2
View File
@@ -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() {
mtproto/rpc/rpc_handshake_response.go
+2 -2
View File
@@ -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 {
mtproto/rpc/rpc_nonce_request.go
+2 -2
View File
@@ -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) {
mtproto/rpc/rpc_nonce_response.go
+2 -2
View File
@@ -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 {
mtproto/rpc/rpc_proxy_request.go
+3 -3
View File
@@ -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) {
mtproto/rwc.go
-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()
}
mtproto/crypt/crypt.go → mtproto/wrappers/crypt.go
+1 -1
View File
@@ -1,4 +1,4 @@
package crypt
package wrappers
import (
"bytes"
mtproto/wrappers/frame.go
+11 -20
View File
@@ -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{},
}
}