mirror of
https://github.com/ScuroNeko/mtg.git
synced 2026-08-31 16:44:03 +03:00
Add base wrappers for mtproto
This commit is contained in:
9seconds
@@ -0,0 +1,157 @@
|
||||
package wrappers
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const blockCipherReadCurrentDataBufferSize = 1024 + 1 // +1 because telegram operates with blocks mod 4
|
||||
|
||||
type wrapperBlockCipher struct {
|
||||
buf bytes.Buffer
|
||||
|
||||
parent StreamReadWriteCloser
|
||||
encryptor cipher.BlockMode
|
||||
decryptor cipher.BlockMode
|
||||
}
|
||||
|
||||
func (w *wrapperBlockCipher) Write(p []byte) (int, error) {
|
||||
encrypted, err := w.encrypt(p)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return w.parent.Write(encrypted)
|
||||
}
|
||||
|
||||
func (w *wrapperBlockCipher) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
|
||||
encrypted, err := w.encrypt(p)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return w.parent.WriteTimeout(encrypted, timeout)
|
||||
}
|
||||
|
||||
func (w *wrapperBlockCipher) Read(p []byte) (int, error) {
|
||||
return w.read(p, readAll)
|
||||
|
||||
}
|
||||
|
||||
func (w *wrapperBlockCipher) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
|
||||
return w.read(p, readAllTimeout(timeout))
|
||||
}
|
||||
|
||||
func (w *wrapperBlockCipher) read(p []byte, reader func(StreamReadWriteCloser) ([]byte, error)) (int, error) {
|
||||
if w.buf.Len() > 0 {
|
||||
return w.flush(p)
|
||||
}
|
||||
|
||||
var buf []byte
|
||||
for len(buf) == 0 || len(buf)%aes.BlockSize != 0 {
|
||||
rv, err := reader(w.parent)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("cannot read from socket: %w", err)
|
||||
}
|
||||
buf = append(buf, rv...)
|
||||
}
|
||||
|
||||
w.decryptor.CryptBlocks(buf, buf)
|
||||
w.buf.Write(buf)
|
||||
|
||||
return w.flush(p)
|
||||
}
|
||||
|
||||
func (w *wrapperBlockCipher) flush(p []byte) (int, error) {
|
||||
if w.buf.Len() > len(p) {
|
||||
return w.buf.Read(p)
|
||||
}
|
||||
|
||||
sizeToReturn := w.buf.Len()
|
||||
copy(p, w.buf.Bytes())
|
||||
w.buf.Reset()
|
||||
|
||||
return sizeToReturn, nil
|
||||
}
|
||||
|
||||
func (w *wrapperBlockCipher) encrypt(p []byte) ([]byte, error) {
|
||||
if len(p)%aes.BlockSize > 0 {
|
||||
return nil, fmt.Errorf("incorrect block size %d", len(p))
|
||||
}
|
||||
|
||||
encrypted := make([]byte, len(p))
|
||||
w.encryptor.CryptBlocks(encrypted, p)
|
||||
|
||||
return encrypted, nil
|
||||
}
|
||||
|
||||
func readAll(src StreamReadWriteCloser) (rv []byte, err error) {
|
||||
buf := make([]byte, blockCipherReadCurrentDataBufferSize)
|
||||
n := blockCipherReadCurrentDataBufferSize
|
||||
|
||||
for n == len(buf) {
|
||||
n, err = src.Read(buf)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rv = append(rv, buf[:n]...)
|
||||
}
|
||||
|
||||
return rv, nil
|
||||
}
|
||||
|
||||
func readAllTimeout(timeout time.Duration) func(StreamReadWriteCloser) ([]byte, error) {
|
||||
return func(src StreamReadWriteCloser) (rv []byte, err error) {
|
||||
tmo := timeout
|
||||
buf := make([]byte, blockCipherReadCurrentDataBufferSize)
|
||||
n := blockCipherReadCurrentDataBufferSize
|
||||
|
||||
for n == len(buf) {
|
||||
if tmo <= 0 {
|
||||
return nil, errors.New("timeout")
|
||||
}
|
||||
startTime := time.Now()
|
||||
n, err = src.ReadTimeout(buf, tmo)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rv = append(rv, buf[:n]...)
|
||||
tmo -= time.Since(startTime)
|
||||
}
|
||||
|
||||
return rv, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (w *wrapperBlockCipher) Close() error {
|
||||
return w.parent.Close()
|
||||
}
|
||||
|
||||
func (w *wrapperBlockCipher) Conn() net.Conn {
|
||||
return w.parent.Conn()
|
||||
}
|
||||
|
||||
func (w *wrapperBlockCipher) Logger() *zap.SugaredLogger {
|
||||
return w.parent.Logger().Named("block-cipher")
|
||||
}
|
||||
|
||||
func (w *wrapperBlockCipher) LocalAddr() *net.TCPAddr {
|
||||
return w.parent.LocalAddr()
|
||||
}
|
||||
|
||||
func (w *wrapperBlockCipher) RemoteAddr() *net.TCPAddr {
|
||||
return w.parent.RemoteAddr()
|
||||
}
|
||||
|
||||
func NewBlockCipher(parent StreamReadWriteCloser, encryptor, decryptor cipher.BlockMode) StreamReadWriteCloser {
|
||||
return &wrapperBlockCipher{
|
||||
parent: parent,
|
||||
encryptor: encryptor,
|
||||
decryptor: decryptor,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
package wrappers
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/aes"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"hash/crc32"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"net"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
mtprotoFrameMinMessageLength = 12
|
||||
mtprotoFrameMaxMessageLength = 16777216
|
||||
)
|
||||
|
||||
var mtprotoFramePadding = []byte{0x04, 0x00, 0x00, 0x00}
|
||||
|
||||
// MTProtoFrame is a wrapper which converts written data to the MTProtoFrame.
|
||||
// The format of the frame:
|
||||
//
|
||||
// [ MSGLEN(4) | SEQNO(4) | MSG(...) | CRC32(4) | PADDING(4*x) ]
|
||||
//
|
||||
// MSGLEN is the length of the message + len of seqno and msglen.
|
||||
// SEQNO is the number of frame in the receive/send sequence. If client
|
||||
// sends a message with SeqNo 18, it has to receive message with SeqNo 18.
|
||||
// MSG is the data which has to be written
|
||||
// CRC32 is the CRC32 checksum of MSGLEN + SEQNO + MSG
|
||||
// PADDING is custom padding schema to complete frame length to such that
|
||||
// len(frame) % 16 == 0
|
||||
type wrapperMtprotoFrame struct {
|
||||
parent StreamReadWriteCloser
|
||||
logger *zap.SugaredLogger
|
||||
readSeqNo int32
|
||||
writeSeqNo int32
|
||||
}
|
||||
|
||||
func (w *wrapperMtprotoFrame) Read() (Packet, error) {
|
||||
buf := &bytes.Buffer{}
|
||||
sum := crc32.NewIEEE()
|
||||
writer := io.MultiWriter(buf, sum)
|
||||
|
||||
for {
|
||||
buf.Reset()
|
||||
sum.Reset()
|
||||
if _, err := io.CopyN(writer, w.parent, 4); err != nil {
|
||||
return nil, fmt.Errorf("cannot read frame padding: %w", err)
|
||||
}
|
||||
if !bytes.Equal(buf.Bytes(), mtprotoFramePadding) {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
messageLength := binary.LittleEndian.Uint32(buf.Bytes())
|
||||
w.logger.Debugw("Read MTProto frame",
|
||||
"messageLength", messageLength,
|
||||
"sequence_number", w.readSeqNo,
|
||||
)
|
||||
if messageLength%4 != 0 || messageLength < mtprotoFrameMinMessageLength ||
|
||||
messageLength > mtprotoFrameMaxMessageLength {
|
||||
return nil, fmt.Errorf("Incorrect frame message length %d", messageLength)
|
||||
}
|
||||
|
||||
buf.Reset()
|
||||
buf.Grow(int(messageLength) - 4 - 4)
|
||||
if _, err := io.CopyN(writer, w.parent, int64(messageLength)-4-4); err != nil {
|
||||
return nil, fmt.Errorf("cannot read the message frame: %w", err)
|
||||
}
|
||||
|
||||
var seqNo int32
|
||||
binary.Read(buf, binary.LittleEndian, &seqNo) // nolint: errcheck, gosec
|
||||
if seqNo != w.readSeqNo {
|
||||
return nil, fmt.Errorf("unexpected sequence number %d (wait for %d)", seqNo, w.readSeqNo)
|
||||
}
|
||||
|
||||
data, _ := ioutil.ReadAll(buf) // nolint: gosec
|
||||
buf.Reset()
|
||||
// write to buf, not to writer. This is because we are going to fetch
|
||||
// crc32 checksum.
|
||||
if _, err := io.CopyN(buf, w.parent, 4); err != nil {
|
||||
return nil, fmt.Errorf("cannot read checksum: %w", err)
|
||||
}
|
||||
|
||||
checksum := binary.LittleEndian.Uint32(buf.Bytes())
|
||||
if checksum != sum.Sum32() {
|
||||
return nil, fmt.Errorf("CRC32 checksum mismatch. wait for %d, got %d", sum.Sum32(), checksum)
|
||||
}
|
||||
|
||||
w.logger.Debugw("Read MTProto frame",
|
||||
"messageLength", messageLength,
|
||||
"sequence_number", w.readSeqNo,
|
||||
"dataLength", len(data),
|
||||
"checksum", checksum,
|
||||
)
|
||||
w.readSeqNo++
|
||||
|
||||
return data, nil
|
||||
}
|
||||
|
||||
func (w *wrapperMtprotoFrame) Write(p Packet) error {
|
||||
messageLength := 4 + 4 + len(p) + 4
|
||||
paddingLength := (aes.BlockSize - messageLength%aes.BlockSize) % aes.BlockSize
|
||||
|
||||
buf := &bytes.Buffer{}
|
||||
buf.Grow(messageLength + paddingLength)
|
||||
|
||||
binary.Write(buf, binary.LittleEndian, uint32(messageLength))
|
||||
binary.Write(buf, binary.LittleEndian, w.writeSeqNo)
|
||||
buf.Write(p)
|
||||
|
||||
checksum := crc32.ChecksumIEEE(buf.Bytes())
|
||||
binary.Write(buf, binary.LittleEndian, checksum)
|
||||
buf.Write(bytes.Repeat(mtprotoFramePadding, paddingLength/4))
|
||||
|
||||
w.logger.Debugw("Write MTProto frame",
|
||||
"length", len(p),
|
||||
"sequence_number", w.writeSeqNo,
|
||||
"crc32", checksum,
|
||||
"frame_length", buf.Len(),
|
||||
)
|
||||
w.writeSeqNo++
|
||||
|
||||
_, err := w.parent.Write(buf.Bytes())
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
func (w *wrapperMtprotoFrame) Close() error {
|
||||
return w.parent.Close()
|
||||
}
|
||||
|
||||
func (w *wrapperMtprotoFrame) Conn() net.Conn {
|
||||
return w.parent.Conn()
|
||||
}
|
||||
|
||||
func (w *wrapperMtprotoFrame) Logger() *zap.SugaredLogger {
|
||||
return w.logger
|
||||
}
|
||||
|
||||
func (w *wrapperMtprotoFrame) LocalAddr() *net.TCPAddr {
|
||||
return w.parent.LocalAddr()
|
||||
}
|
||||
|
||||
func (w *wrapperMtprotoFrame) RemoteAddr() *net.TCPAddr {
|
||||
return w.parent.RemoteAddr()
|
||||
}
|
||||
|
||||
func NewMtprotoFrame(parent StreamReadWriteCloser, seqNo int32) PacketReadWriteCloser {
|
||||
return &wrapperMtprotoFrame{
|
||||
parent: parent,
|
||||
logger: parent.Logger().Named("mtproto-frame"),
|
||||
readSeqNo: seqNo,
|
||||
writeSeqNo: seqNo,
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user