diff --git a/client/middle.go b/client/middle.go index 53bb64c..4e50895 100644 --- a/client/middle.go +++ b/client/middle.go @@ -17,9 +17,16 @@ func MiddleInit(socket net.Conn, connID string, conf *config.Config) (wrappers.W } connStream := conn.(wrappers.StreamReadWriteCloser) - newConn := wrappers.NewMTProtoAbridged(connStream, opts) - if opts.ConnectionType != mtproto.ConnectionTypeAbridged { + var newConn wrappers.PacketReadWriteCloser + switch opts.ConnectionType { + case mtproto.ConnectionTypeAbridged: + newConn = wrappers.NewMTProtoAbridged(connStream, opts) + case mtproto.ConnectionTypeIntermediate: newConn = wrappers.NewMTProtoIntermediate(connStream, opts) + case mtproto.ConnectionTypeSecure: + newConn = wrappers.NewMTProtoIntermediateSecure(connStream, opts) + default: + panic("Unknown connection type") } opts.ConnectionProto = mtproto.ConnectionProtocolIPv4 diff --git a/wrappers/mtproto_intermediate_secure.go b/wrappers/mtproto_intermediate_secure.go new file mode 100644 index 0000000..954c3d0 --- /dev/null +++ b/wrappers/mtproto_intermediate_secure.go @@ -0,0 +1,69 @@ +package wrappers + +import ( + "bytes" + "encoding/binary" + "math/rand" + + "github.com/9seconds/mtg/mtproto" +) + +type MTProtoIntermediateSecure struct { + MTProtoIntermediate +} + +func (m *MTProtoIntermediateSecure) Read() ([]byte, error) { + data, err := m.MTProtoIntermediate.Read() + if err != nil { + return nil, err + } + length := len(data) - (len(data) % 4) + + return data[:length], nil +} + +func (m *MTProtoIntermediateSecure) Write(p []byte) (int, error) { + defer func() { + m.writeCounter++ + }() + + m.logger.Debugw("Write packet", + "simple_ack", m.opts.WriteHacks.SimpleAck, + "quick_ack", m.opts.WriteHacks.QuickAck, + "counter", m.writeCounter, + ) + + if m.opts.ReadHacks.SimpleAck { + return m.conn.Write(p) + } + + buf := &bytes.Buffer{} + paddingLength := rand.Intn(4) + buf.Grow(4 + len(p) + paddingLength) + + binary.Write(buf, binary.LittleEndian, uint32(len(p)+paddingLength)) + buf.Write(p) + buf.Write(make([]byte, paddingLength)) + + m.logger.Debugw("Write packet with padding", + "simple_ack", m.opts.WriteHacks.SimpleAck, + "quick_ack", m.opts.WriteHacks.QuickAck, + "counter", m.writeCounter, + "padding_length", paddingLength, + "length", len(p), + ) + + _, err := m.conn.Write(buf.Bytes()) + + return len(p), err +} + +func NewMTProtoIntermediateSecure(conn StreamReadWriteCloser, opts *mtproto.ConnectionOpts) PacketReadWriteCloser { + return &MTProtoIntermediateSecure{ + MTProtoIntermediate: MTProtoIntermediate{ + conn: conn, + logger: conn.Logger().Named("mtproto-intermediate-secure"), + opts: opts, + }, + } +}