mirror of
https://github.com/ScuroNeko/mtg.git
synced 2026-08-31 16:14:02 +03:00
Add client abridged protocol
This commit is contained in:
9seconds
@@ -0,0 +1,21 @@
|
||||
package stream
|
||||
|
||||
import (
|
||||
"net"
|
||||
|
||||
"github.com/9seconds/mtg/conntypes"
|
||||
)
|
||||
|
||||
func NewClientConn(parent net.Conn, connID conntypes.ConnID) conntypes.StreamReadWriteCloser {
|
||||
conn := newConn(parent, connID, connPurposeClient)
|
||||
conn = NewTrafficStats(conn)
|
||||
|
||||
return conn
|
||||
}
|
||||
|
||||
func NewTelegramConn(dc conntypes.DC, parent net.Conn) conntypes.StreamReadWriteCloser {
|
||||
conn := newConn(parent, conntypes.ConnID{}, connPurposeTelegram)
|
||||
conn = NewTelegramStats(dc, conn)
|
||||
|
||||
return conn
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
package stream
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"fmt"
|
||||
"net"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/9seconds/mtg/conntypes"
|
||||
"github.com/9seconds/mtg/utils"
|
||||
)
|
||||
|
||||
type wrapperBlockCipher struct {
|
||||
buf bytes.Buffer
|
||||
|
||||
parent conntypes.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) {
|
||||
if w.buf.Len() > 0 {
|
||||
return w.flush(p)
|
||||
}
|
||||
|
||||
var currentBuffer []byte
|
||||
for len(currentBuffer) == 0 || len(currentBuffer)%aes.BlockSize != 0 {
|
||||
rv, err := utils.ReadFull(w.parent)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("cannot read data: %w", err)
|
||||
}
|
||||
currentBuffer = append(currentBuffer, rv...)
|
||||
}
|
||||
|
||||
w.decryptor.CryptBlocks(currentBuffer, currentBuffer)
|
||||
w.buf.Write(currentBuffer)
|
||||
|
||||
return w.flush(p)
|
||||
}
|
||||
|
||||
func (w *wrapperBlockCipher) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
|
||||
return w.Read(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 (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 conntypes.StreamReadWriteCloser,
|
||||
encryptor, decryptor cipher.BlockMode) conntypes.StreamReadWriteCloser {
|
||||
return &wrapperBlockCipher{
|
||||
parent: parent,
|
||||
encryptor: encryptor,
|
||||
decryptor: decryptor,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package stream
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/9seconds/mtg/config"
|
||||
"github.com/9seconds/mtg/conntypes"
|
||||
)
|
||||
|
||||
type connPurpose uint8
|
||||
|
||||
const (
|
||||
connPurposeClient connPurpose = 1 << iota
|
||||
connPurposeTelegram
|
||||
)
|
||||
|
||||
type wrapperConn struct {
|
||||
parent net.Conn
|
||||
connID conntypes.ConnID
|
||||
logger *zap.SugaredLogger
|
||||
localAddr *net.TCPAddr
|
||||
remoteAddr *net.TCPAddr
|
||||
}
|
||||
|
||||
func (w *wrapperConn) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
|
||||
if err := w.parent.SetWriteDeadline(time.Now().Add(timeout)); err != nil {
|
||||
w.Close()
|
||||
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) {
|
||||
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) {
|
||||
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)
|
||||
}
|
||||
|
||||
return w.Read(p)
|
||||
}
|
||||
|
||||
func (w *wrapperConn) Read(p []byte) (int, error) {
|
||||
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")
|
||||
return w.parent.Close()
|
||||
}
|
||||
|
||||
func (w *wrapperConn) Conn() net.Conn {
|
||||
return w.parent
|
||||
}
|
||||
|
||||
func (w *wrapperConn) Logger() *zap.SugaredLogger {
|
||||
return w.logger
|
||||
}
|
||||
|
||||
func (w *wrapperConn) LocalAddr() *net.TCPAddr {
|
||||
return w.localAddr
|
||||
}
|
||||
|
||||
func (w *wrapperConn) RemoteAddr() *net.TCPAddr {
|
||||
return w.remoteAddr
|
||||
}
|
||||
|
||||
func newConn(parent net.Conn,
|
||||
connID conntypes.ConnID,
|
||||
purpose connPurpose) conntypes.StreamReadWriteCloser {
|
||||
localAddr := *parent.LocalAddr().(*net.TCPAddr)
|
||||
|
||||
if parent.RemoteAddr().(*net.TCPAddr).IP.To4() != nil {
|
||||
if config.C.PublicIPv4.IP != nil {
|
||||
localAddr.IP = config.C.PublicIPv4.IP
|
||||
}
|
||||
} else if config.C.PublicIPv6.IP != nil {
|
||||
localAddr.IP = config.C.PublicIPv6.IP
|
||||
}
|
||||
|
||||
logger := zap.S().With(
|
||||
"local_address", localAddr,
|
||||
"remote_address", parent.RemoteAddr(),
|
||||
).Named("conn")
|
||||
|
||||
if purpose == connPurposeClient {
|
||||
logger = logger.With("connection_id", connID.String())
|
||||
}
|
||||
|
||||
return &wrapperConn{
|
||||
parent: parent,
|
||||
connID: connID,
|
||||
logger: logger,
|
||||
remoteAddr: parent.RemoteAddr().(*net.TCPAddr),
|
||||
localAddr: &localAddr,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
package stream
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
package stream
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"crypto/md5"
|
||||
"crypto/sha1"
|
||||
"encoding/binary"
|
||||
"net"
|
||||
|
||||
"github.com/9seconds/mtg/conntypes"
|
||||
"github.com/9seconds/mtg/mtproto/rpc"
|
||||
"github.com/9seconds/mtg/utils"
|
||||
)
|
||||
|
||||
type mtprotoCipherPurpose uint8
|
||||
|
||||
const (
|
||||
mtprotoCipherPurposeClient mtprotoCipherPurpose = iota
|
||||
mtprotoCipherPurposeServer
|
||||
)
|
||||
|
||||
var mtprotoEmptyIP = [4]byte{0x00, 0x00, 0x00, 0x00}
|
||||
|
||||
func NewMiddleProxyCipher(parent conntypes.StreamReadWriteCloser,
|
||||
req *rpc.NonceRequest,
|
||||
resp *rpc.NonceResponse,
|
||||
secret []byte) conntypes.StreamReadWriteCloser {
|
||||
localAddr := parent.LocalAddr()
|
||||
remoteAddr := parent.RemoteAddr()
|
||||
|
||||
encKey, encIV := mtprotoDeriveKeys(mtprotoCipherPurposeClient,
|
||||
req,
|
||||
resp,
|
||||
localAddr,
|
||||
remoteAddr,
|
||||
secret)
|
||||
decKey, decIV := mtprotoDeriveKeys(mtprotoCipherPurposeServer,
|
||||
req,
|
||||
resp,
|
||||
localAddr,
|
||||
remoteAddr,
|
||||
secret)
|
||||
|
||||
enc, _ := mtprotoMakeEncrypterDecrypter(encKey, encIV)
|
||||
_, dec := mtprotoMakeEncrypterDecrypter(decKey, decIV)
|
||||
|
||||
return newBlockCipher(parent, enc, dec)
|
||||
}
|
||||
|
||||
func mtprotoDeriveKeys(purpose mtprotoCipherPurpose,
|
||||
req *rpc.NonceRequest,
|
||||
resp *rpc.NonceResponse,
|
||||
client, remote *net.TCPAddr,
|
||||
secret []byte) ([]byte, []byte) {
|
||||
message := bytes.Buffer{}
|
||||
message.Write(resp.Nonce) // nolint: gosec
|
||||
message.Write(req.Nonce) // nolint: gosec
|
||||
message.Write(req.CryptoTS) // nolint: gosec
|
||||
|
||||
clientIPv4 := mtprotoEmptyIP[:]
|
||||
serverIPv4 := mtprotoEmptyIP[:]
|
||||
if client.IP.To4() != nil {
|
||||
clientIPv4 = utils.ReverseBytes(client.IP.To4())
|
||||
serverIPv4 = utils.ReverseBytes(remote.IP.To4())
|
||||
}
|
||||
message.Write(serverIPv4) // nolint: gosec
|
||||
|
||||
var port [2]byte
|
||||
binary.LittleEndian.PutUint16(port[:], uint16(client.Port))
|
||||
message.Write(port[:]) // nolint: gosec
|
||||
|
||||
switch purpose {
|
||||
case mtprotoCipherPurposeClient:
|
||||
message.WriteString("CLIENT") // nolint: gosec
|
||||
case mtprotoCipherPurposeServer:
|
||||
message.WriteString("SERVER") // nolint: gosec
|
||||
default:
|
||||
panic("Unexpected cipher purpose")
|
||||
}
|
||||
|
||||
message.Write(clientIPv4) // nolint: gosec
|
||||
binary.LittleEndian.PutUint16(port[:], uint16(remote.Port))
|
||||
message.Write(port[:]) // nolint: gosec
|
||||
message.Write(secret) // nolint: gosec
|
||||
message.Write(resp.Nonce) // nolint: gosec
|
||||
|
||||
if client.IP.To4() == nil {
|
||||
message.Write(client.IP.To16()) // nolint: gosec
|
||||
message.Write(remote.IP.To16()) // nolint: gosec
|
||||
}
|
||||
message.Write(req.Nonce) // nolint: gosec
|
||||
|
||||
data := message.Bytes()
|
||||
md5sum := md5.Sum(data[1:]) // nolint: gas
|
||||
sha1sum := sha1.Sum(data) // nolint: gosec
|
||||
|
||||
key := append(md5sum[:12], sha1sum[:]...)
|
||||
iv := md5.Sum(data[2:]) // nolint: gas
|
||||
|
||||
return key, iv[:]
|
||||
}
|
||||
|
||||
func mtprotoMakeEncrypterDecrypter(key, iv []byte) (cipher.BlockMode, cipher.BlockMode) {
|
||||
block, err := aes.NewCipher(key)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
return cipher.NewCBCEncrypter(block, iv), cipher.NewCBCDecrypter(block, iv)
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package stream
|
||||
|
||||
import (
|
||||
"crypto/cipher"
|
||||
"fmt"
|
||||
"net"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/9seconds/mtg/conntypes"
|
||||
)
|
||||
|
||||
type wrapperObfuscated2 struct {
|
||||
encryptor cipher.Stream
|
||||
decryptor cipher.Stream
|
||||
parent conntypes.StreamReadWriteCloser
|
||||
}
|
||||
|
||||
func (w *wrapperObfuscated2) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
|
||||
n, err := w.parent.ReadTimeout(p, timeout)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("cannot read stream ciphered data: %w", err)
|
||||
}
|
||||
w.decryptor.XORKeyStream(p, p[:n])
|
||||
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (w *wrapperObfuscated2) Read(p []byte) (int, error) {
|
||||
n, err := w.parent.Read(p)
|
||||
if err != nil {
|
||||
return n, err
|
||||
}
|
||||
w.decryptor.XORKeyStream(p, p[:n])
|
||||
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (w *wrapperObfuscated2) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
|
||||
buf := make([]byte, len(p))
|
||||
copy(buf, p)
|
||||
w.encryptor.XORKeyStream(buf, buf)
|
||||
|
||||
return w.parent.WriteTimeout(buf, timeout)
|
||||
}
|
||||
|
||||
func (w *wrapperObfuscated2) Write(p []byte) (int, error) {
|
||||
buf := make([]byte, len(p))
|
||||
copy(buf, p)
|
||||
w.encryptor.XORKeyStream(buf, buf)
|
||||
|
||||
return w.parent.Write(buf)
|
||||
}
|
||||
|
||||
func (w *wrapperObfuscated2) Conn() net.Conn {
|
||||
return w.parent.Conn()
|
||||
}
|
||||
|
||||
func (w *wrapperObfuscated2) Logger() *zap.SugaredLogger {
|
||||
return w.parent.Logger().Named("obfuscated2")
|
||||
}
|
||||
|
||||
func (w *wrapperObfuscated2) LocalAddr() *net.TCPAddr {
|
||||
return w.parent.LocalAddr()
|
||||
}
|
||||
|
||||
func (w *wrapperObfuscated2) RemoteAddr() *net.TCPAddr {
|
||||
return w.parent.RemoteAddr()
|
||||
}
|
||||
|
||||
func (w *wrapperObfuscated2) Close() error {
|
||||
return w.parent.Close()
|
||||
}
|
||||
|
||||
func NewObfuscated2(socket conntypes.StreamReadWriteCloser,
|
||||
encryptor, decryptor cipher.Stream) conntypes.StreamReadWriteCloser {
|
||||
return &wrapperObfuscated2{
|
||||
parent: socket,
|
||||
encryptor: encryptor,
|
||||
decryptor: decryptor,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package stream
|
||||
|
||||
import (
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/9seconds/mtg/conntypes"
|
||||
"github.com/9seconds/mtg/stats"
|
||||
)
|
||||
|
||||
type wrapperTelegramStats struct {
|
||||
parent conntypes.StreamReadWriteCloser
|
||||
dc conntypes.DC
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (w *wrapperTelegramStats) Write(p []byte) (int, error) {
|
||||
return w.parent.Write(p)
|
||||
}
|
||||
|
||||
func (w *wrapperTelegramStats) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
|
||||
return w.parent.WriteTimeout(p, timeout)
|
||||
}
|
||||
|
||||
func (w *wrapperTelegramStats) Read(p []byte) (int, error) {
|
||||
return w.parent.Read(p)
|
||||
}
|
||||
|
||||
func (w *wrapperTelegramStats) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
|
||||
return w.parent.ReadTimeout(p, timeout)
|
||||
}
|
||||
|
||||
func (w *wrapperTelegramStats) Conn() net.Conn {
|
||||
return w.parent.Conn()
|
||||
}
|
||||
|
||||
func (w *wrapperTelegramStats) Logger() *zap.SugaredLogger {
|
||||
return w.parent.Logger().Named("stats-telegram")
|
||||
}
|
||||
|
||||
func (w *wrapperTelegramStats) LocalAddr() *net.TCPAddr {
|
||||
return w.parent.LocalAddr()
|
||||
}
|
||||
|
||||
func (w *wrapperTelegramStats) RemoteAddr() *net.TCPAddr {
|
||||
return w.parent.RemoteAddr()
|
||||
}
|
||||
|
||||
func (w *wrapperTelegramStats) Close() error {
|
||||
var err error
|
||||
|
||||
w.once.Do(func() {
|
||||
err = w.parent.Close()
|
||||
stats.Stats.TelegramDisconnected(w.dc, w.RemoteAddr())
|
||||
})
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
func NewTelegramStats(dc conntypes.DC, parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser {
|
||||
conn := &wrapperTelegramStats{
|
||||
parent: parent,
|
||||
dc: dc,
|
||||
}
|
||||
stats.Stats.TelegramConnected(dc, parent.RemoteAddr())
|
||||
|
||||
return conn
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
package stream
|
||||
|
||||
import (
|
||||
"net"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/9seconds/mtg/conntypes"
|
||||
"github.com/9seconds/mtg/stats"
|
||||
)
|
||||
|
||||
type wrapperTrafficStats struct {
|
||||
parent conntypes.StreamReadWriteCloser
|
||||
}
|
||||
|
||||
func (w *wrapperTrafficStats) Write(p []byte) (int, error) {
|
||||
n, err := w.parent.Write(p)
|
||||
stats.Stats.EgressTraffic(n)
|
||||
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (w *wrapperTrafficStats) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
|
||||
n, err := w.parent.WriteTimeout(p, timeout)
|
||||
stats.Stats.EgressTraffic(n)
|
||||
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (w *wrapperTrafficStats) Read(p []byte) (int, error) {
|
||||
n, err := w.parent.Read(p)
|
||||
stats.Stats.IngressTraffic(n)
|
||||
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (w *wrapperTrafficStats) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
|
||||
n, err := w.parent.ReadTimeout(p, timeout)
|
||||
stats.Stats.IngressTraffic(n)
|
||||
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (w *wrapperTrafficStats) Conn() net.Conn {
|
||||
return w.parent.Conn()
|
||||
}
|
||||
|
||||
func (w *wrapperTrafficStats) Logger() *zap.SugaredLogger {
|
||||
return w.parent.Logger().Named("stats-traffic")
|
||||
}
|
||||
|
||||
func (w *wrapperTrafficStats) LocalAddr() *net.TCPAddr {
|
||||
return w.parent.LocalAddr()
|
||||
}
|
||||
|
||||
func (w *wrapperTrafficStats) RemoteAddr() *net.TCPAddr {
|
||||
return w.parent.RemoteAddr()
|
||||
}
|
||||
|
||||
func (w *wrapperTrafficStats) Close() error {
|
||||
return w.parent.Close()
|
||||
}
|
||||
|
||||
func NewTrafficStats(parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser {
|
||||
return &wrapperTrafficStats{parent}
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package stream
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user