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

wip

Browse Source
This commit is contained in:
9seconds
2019-10-07 12:13:07 +03:00
parent 072bce2922
commit c9743b5675
28 changed files with 630 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wrappers/blockcipher.go
+10 -6
View File
@@ -10,6 +10,8 @@ import (
"time"
"go.uber.org/zap"
"github.com/9seconds/mtg/conntypes"
)
const blockCipherReadCurrentDataBufferSize = 1024 + 1 // +1 because telegram operates with blocks mod 4
@@ -17,7 +19,7 @@ const blockCipherReadCurrentDataBufferSize = 1024 + 1 // +1 because telegram ope
type wrapperBlockCipher struct {
buf bytes.Buffer
parent StreamReadWriteCloser
parent conntypes.StreamReadWriteCloser
encryptor cipher.BlockMode
decryptor cipher.BlockMode
}
@@ -47,7 +49,8 @@ func (w *wrapperBlockCipher) ReadTimeout(p []byte, timeout time.Duration) (int,
return w.read(p, readAllTimeout(timeout))
}
func (w *wrapperBlockCipher) read(p []byte, reader func(StreamReadWriteCloser) ([]byte, error)) (int, error) {
func (w *wrapperBlockCipher) read(p []byte,
reader func(conntypes.StreamReadWriteCloser) ([]byte, error)) (int, error) {
if w.buf.Len() > 0 {
return w.flush(p)
}
@@ -90,7 +93,7 @@ func (w *wrapperBlockCipher) encrypt(p []byte) ([]byte, error) {
return encrypted, nil
}
func readAll(src StreamReadWriteCloser) (rv []byte, err error) {
func readAll(src conntypes.StreamReadWriteCloser) (rv []byte, err error) {
buf := make([]byte, blockCipherReadCurrentDataBufferSize)
n := blockCipherReadCurrentDataBufferSize
@@ -105,8 +108,8 @@ func readAll(src StreamReadWriteCloser) (rv []byte, err error) {
return rv, nil
}
func readAllTimeout(timeout time.Duration) func(StreamReadWriteCloser) ([]byte, error) {
return func(src StreamReadWriteCloser) (rv []byte, err error) {
func readAllTimeout(timeout time.Duration) func(conntypes.StreamReadWriteCloser) ([]byte, error) {
return func(src conntypes.StreamReadWriteCloser) (rv []byte, err error) {
tmo := timeout
buf := make([]byte, blockCipherReadCurrentDataBufferSize)
n := blockCipherReadCurrentDataBufferSize
@@ -148,7 +151,8 @@ func (w *wrapperBlockCipher) RemoteAddr() *net.TCPAddr {
return w.parent.RemoteAddr()
}
func newBlockCipher(parent StreamReadWriteCloser, encryptor, decryptor cipher.BlockMode) StreamReadWriteCloser {
func newBlockCipher(parent conntypes.StreamReadWriteCloser,
encryptor, decryptor cipher.BlockMode) conntypes.StreamReadWriteCloser {
return &wrapperBlockCipher{
parent: parent,
encryptor: encryptor,
wrappers/conn.go
+29 -61
View File
@@ -1,7 +1,6 @@
package wrappers
import (
"context"
"fmt"
"net"
"time"
@@ -19,15 +18,8 @@ const (
connPurposeTelegram
)
const (
connTimeoutRead = 2 * time.Minute
connTimeoutWrite = 2 * time.Minute
)
type wrapperConn struct {
parent net.Conn
ctx context.Context
cancel context.CancelFunc
connID conntypes.ConnID
logger *zap.SugaredLogger
localAddr *net.TCPAddr
@@ -35,61 +27,45 @@ type wrapperConn struct {
}
func (w *wrapperConn) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
select {
case <-w.ctx.Done():
if err := w.parent.SetWriteDeadline(time.Now().Add(timeout)); err != nil {
w.Close()
return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err())
default:
if err := w.parent.SetWriteDeadline(time.Now().Add(timeout)); err != nil {
w.Close() // nolint: gosec
return 0, fmt.Errorf("cannot set write deadline to the socket: %w", err)
}
n, err := w.parent.Write(p)
w.logger.Debugw("Write to stream", "bytes", n, "error", err)
if err != nil {
w.Close() // nolint: gosec
}
return n, err
return 0, fmt.Errorf("cannot set write deadline to the socket: %w", err)
}
return w.Write(p)
}
func (w *wrapperConn) Write(p []byte) (int, error) {
return w.WriteTimeout(p, connTimeoutWrite)
n, err := w.parent.Write(p)
w.logger.Debugw("write to stream", "bytes", n, "error", err)
if err != nil {
w.Close() // nolint: gosec
}
return n, err
}
func (w *wrapperConn) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
select {
case <-w.ctx.Done():
if err := w.parent.SetReadDeadline(time.Now().Add(timeout)); err != nil {
w.Close()
return 0, fmt.Errorf("cannot read because context was closed: %w", w.ctx.Err())
default:
if err := w.parent.SetReadDeadline(time.Now().Add(timeout)); err != nil {
w.Close()
return 0, fmt.Errorf("cannot set read deadline to the socket: %w", err)
}
n, err := w.parent.Read(p)
w.logger.Debugw("Read from stream", "bytes", n, "error", err)
if err != nil {
w.Close()
}
return n, err
return 0, fmt.Errorf("cannot set read deadline to the socket: %w", err)
}
return w.Read(p)
}
func (w *wrapperConn) Read(p []byte) (int, error) {
return w.ReadTimeout(p, connTimeoutRead)
n, err := w.parent.Read(p)
w.logger.Debugw("Read from stream", "bytes", n, "error", err)
if err != nil {
w.Close()
}
return n, err
}
func (w *wrapperConn) Close() error {
w.logger.Debugw("Close connection")
w.cancel()
return w.parent.Close()
}
@@ -109,11 +85,9 @@ func (w *wrapperConn) RemoteAddr() *net.TCPAddr {
return w.remoteAddr
}
func newConn(ctx context.Context,
cancel context.CancelFunc,
parent net.Conn,
func newConn(parent net.Conn,
connID conntypes.ConnID,
purpose connPurpose) StreamReadWriteCloser {
purpose connPurpose) conntypes.StreamReadWriteCloser {
localAddr := *parent.LocalAddr().(*net.TCPAddr)
if parent.RemoteAddr().(*net.TCPAddr).IP.To4() != nil {
@@ -135,8 +109,6 @@ func newConn(ctx context.Context,
return &wrapperConn{
parent: parent,
ctx: ctx,
cancel: cancel,
connID: connID,
logger: logger,
remoteAddr: parent.RemoteAddr().(*net.TCPAddr),
@@ -144,15 +116,11 @@ func newConn(ctx context.Context,
}
}
func NewClientConn(ctx context.Context,
cancel context.CancelFunc,
parent net.Conn,
connID conntypes.ConnID) StreamReadWriteCloser {
return newConn(ctx, cancel, parent, connID, connPurposeClient)
func NewClientConn(parent net.Conn,
connID conntypes.ConnID) conntypes.StreamReadWriteCloser {
return newConn(parent, connID, connPurposeClient)
}
func NewTelegramConn(ctx context.Context,
cancel context.CancelFunc,
parent net.Conn) StreamReadWriteCloser {
return newConn(ctx, cancel, parent, conntypes.ConnID{}, connPurposeTelegram)
func NewTelegramConn(parent net.Conn) conntypes.StreamReadWriteCloser {
return newConn(parent, conntypes.ConnID{}, connPurposeTelegram)
}
wrappers/ctx.go
+89
View File
@@ -0,0 +1,89 @@
package wrappers
import (
"context"
"fmt"
"net"
"time"
"go.uber.org/zap"
"github.com/9seconds/mtg/conntypes"
)
type wrapperCtx struct {
parent conntypes.StreamReadWriteCloser
ctx context.Context
cancel context.CancelFunc
}
func (w *wrapperCtx) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
select {
case <-w.ctx.Done():
w.Close()
return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err())
default:
return w.parent.WriteTimeout(p, timeout)
}
}
func (w *wrapperCtx) Write(p []byte) (int, error) {
select {
case <-w.ctx.Done():
w.Close()
return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err())
default:
return w.parent.Write(p)
}
}
func (w *wrapperCtx) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
select {
case <-w.ctx.Done():
w.Close()
return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err())
default:
return w.parent.ReadTimeout(p, timeout)
}
}
func (w *wrapperCtx) Read(p []byte) (int, error) {
select {
case <-w.ctx.Done():
w.Close()
return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err())
default:
return w.parent.Read(p)
}
}
func (w *wrapperCtx) Close() error {
w.cancel()
return w.parent.Close()
}
func (w *wrapperCtx) Conn() net.Conn {
return w.parent.Conn()
}
func (w *wrapperCtx) Logger() *zap.SugaredLogger {
return w.parent.Logger().Named("ctx")
}
func (w *wrapperCtx) LocalAddr() *net.TCPAddr {
return w.parent.LocalAddr()
}
func (w *wrapperCtx) RemoteAddr() *net.TCPAddr {
return w.parent.RemoteAddr()
}
func NewCtx(ctx context.Context,
cancel context.CancelFunc,
parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser {
return &wrapperCtx{
parent: parent,
ctx: ctx,
cancel: cancel,
}
}
wrappers/interfaces.go
-117
View File
@@ -1,117 +0,0 @@
package wrappers
import (
"io"
"net"
"time"
"go.uber.org/zap"
)
type Packet []byte
// Wrap is a base interface for all wrappers in this package.
type Wrap interface {
Conn() net.Conn
Logger() *zap.SugaredLogger
LocalAddr() *net.TCPAddr
RemoteAddr() *net.TCPAddr
}
type BaseReaderWithTimeout interface {
ReadTimeout([]byte, time.Duration) (int, error)
}
type BaseWriterWithTimeout interface {
WriteTimeout([]byte, time.Duration) (int, error)
}
type BasePacketReader interface {
Read() (Packet, error)
}
type BasePacketWriter interface {
Write(Packet) error
}
type StreamReader interface {
Wrap
io.Reader
BaseReaderWithTimeout
}
type StreamWriter interface {
Wrap
io.Writer
BaseWriterWithTimeout
}
type StreamCloser interface {
Wrap
io.Closer
}
type StreamReadCloser interface {
Wrap
io.ReadCloser
BaseReaderWithTimeout
}
type StreamWriteCloser interface {
Wrap
io.WriteCloser
BaseWriterWithTimeout
}
type StreamReadWriter interface {
Wrap
io.ReadWriter
BaseReaderWithTimeout
}
type StreamReadWriteCloser interface {
Wrap
io.ReadWriteCloser
BaseReaderWithTimeout
BaseWriterWithTimeout
}
type PacketReader interface {
Wrap
BasePacketReader
}
type PacketWriter interface {
Wrap
BasePacketWriter
}
type PacketCloser interface {
Wrap
io.Closer
}
type PacketReadCloser interface {
Wrap
BasePacketReader
io.Closer
}
type PacketWriteCloser interface {
Wrap
BasePacketWriter
io.Closer
}
type PacketReadWriter interface {
Wrap
BasePacketWriter
BasePacketReader
}
type PacketReadWriteCloser interface {
Wrap
BasePacketWriter
BasePacketReader
io.Closer
}
wrappers/mtproto_cipher.go
+3 -2
View File
@@ -9,6 +9,7 @@ import (
"encoding/binary"
"net"
"github.com/9seconds/mtg/conntypes"
"github.com/9seconds/mtg/mtproto/rpc"
"github.com/9seconds/mtg/utils"
)
@@ -22,10 +23,10 @@ const (
var mtprotoEmptyIP = [4]byte{0x00, 0x00, 0x00, 0x00}
func NewMiddleProxyCipher(parent StreamReadWriteCloser,
func NewMiddleProxyCipher(parent conntypes.StreamReadWriteCloser,
req *rpc.NonceRequest,
resp *rpc.NonceResponse,
secret []byte) StreamReadWriteCloser {
secret []byte) conntypes.StreamReadWriteCloser {
localAddr := parent.LocalAddr()
remoteAddr := parent.RemoteAddr()
wrappers/mtproto_frame.go
+6 -4
View File
@@ -11,6 +11,8 @@ import (
"net"
"go.uber.org/zap"
"github.com/9seconds/mtg/conntypes"
)
const (
@@ -33,13 +35,13 @@ var mtprotoFramePadding = []byte{0x04, 0x00, 0x00, 0x00}
// PADDING is custom padding schema to complete frame length to such that
// len(frame) % 16 == 0
type wrapperMtprotoFrame struct {
parent StreamReadWriteCloser
parent conntypes.StreamReadWriteCloser
logger *zap.SugaredLogger
readSeqNo int32
writeSeqNo int32
}
func (w *wrapperMtprotoFrame) Read() (Packet, error) {
func (w *wrapperMtprotoFrame) Read() (conntypes.Packet, error) {
buf := &bytes.Buffer{}
sum := crc32.NewIEEE()
writer := io.MultiWriter(buf, sum)
@@ -101,7 +103,7 @@ func (w *wrapperMtprotoFrame) Read() (Packet, error) {
return data, nil
}
func (w *wrapperMtprotoFrame) Write(p Packet) error {
func (w *wrapperMtprotoFrame) Write(p conntypes.Packet) error {
messageLength := 4 + 4 + len(p) + 4
paddingLength := (aes.BlockSize - messageLength%aes.BlockSize) % aes.BlockSize
@@ -149,7 +151,7 @@ func (w *wrapperMtprotoFrame) RemoteAddr() *net.TCPAddr {
return w.parent.RemoteAddr()
}
func NewMtprotoFrame(parent StreamReadWriteCloser, seqNo int32) PacketReadWriteCloser {
func NewMtprotoFrame(parent conntypes.StreamReadWriteCloser, seqNo int32) conntypes.PacketReadWriteCloser {
return &wrapperMtprotoFrame{
parent: parent,
logger: parent.Logger().Named("mtproto-frame"),
wrappers/obfuscated2.go
+5 -2
View File
@@ -7,12 +7,14 @@ import (
"time"
"go.uber.org/zap"
"github.com/9seconds/mtg/conntypes"
)
type wrapperObfuscated2 struct {
encryptor cipher.Stream
decryptor cipher.Stream
parent StreamReadWriteCloser
parent conntypes.StreamReadWriteCloser
}
func (w *wrapperObfuscated2) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
@@ -71,7 +73,8 @@ func (w *wrapperObfuscated2) Close() error {
return w.parent.Close()
}
func NewObfuscated2(socket StreamReadWriteCloser, encryptor, decryptor cipher.Stream) StreamReadWriteCloser {
func NewObfuscated2(socket conntypes.StreamReadWriteCloser,
encryptor, decryptor cipher.Stream) conntypes.StreamReadWriteCloser {
return &wrapperObfuscated2{
parent: socket,
encryptor: encryptor,
wrappers/stats.go
+3 -2
View File
@@ -6,11 +6,12 @@ import (
"go.uber.org/zap"
"github.com/9seconds/mtg/conntypes"
"github.com/9seconds/mtg/stats"
)
type wrapperStats struct {
parent StreamReadWriteCloser
parent conntypes.StreamReadWriteCloser
}
func (w *wrapperStats) Write(p []byte) (int, error) {
@@ -61,6 +62,6 @@ func (w *wrapperStats) Close() error {
return w.parent.Close()
}
func NewTraffic(parent StreamReadWriteCloser) StreamReadWriteCloser {
func NewTraffic(parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser {
return &wrapperStats{parent}
}
wrappers/timeout.go
+61
View File
@@ -0,0 +1,61 @@
package wrappers
import (
"net"
"time"
"go.uber.org/zap"
"github.com/9seconds/mtg/conntypes"
)
const (
timeoutRead = 2 * time.Minute
timeoutWrite = 2 * time.Minute
)
type wrapperTimeout struct {
parent conntypes.StreamReadWriteCloser
}
func (w *wrapperTimeout) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
return w.parent.WriteTimeout(p, timeout)
}
func (w *wrapperTimeout) Write(p []byte) (int, error) {
return w.parent.WriteTimeout(p, timeoutWrite)
}
func (w *wrapperTimeout) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
return w.parent.ReadTimeout(p, timeout)
}
func (w *wrapperTimeout) Read(p []byte) (int, error) {
return w.parent.ReadTimeout(p, timeoutRead)
}
func (w *wrapperTimeout) Close() error {
return w.parent.Close()
}
func (w *wrapperTimeout) Conn() net.Conn {
return w.parent.Conn()
}
func (w *wrapperTimeout) Logger() *zap.SugaredLogger {
return w.parent.Logger().Named("timeout")
}
func (w *wrapperTimeout) LocalAddr() *net.TCPAddr {
return w.parent.LocalAddr()
}
func (w *wrapperTimeout) RemoteAddr() *net.TCPAddr {
return w.parent.RemoteAddr()
}
func NewTimeout(parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser {
return &wrapperTimeout{
parent: parent,
}
}