Reworked base

2026-08-31 20:14:01 +03:00 · 2019-08-29 13:16:08 +03:00
parent 09c7ce45d2
commit 07985cf418
26 changed files with 1481 additions and 271 deletions
@@ -0,0 +1,117 @@
+package newwrappers
+
+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
+}
@@ -0,0 +1,178 @@
+package newwrappers
+
+import (
+	"context"
+	"crypto/rand"
+	"encoding/hex"
+	"net"
+	"time"
+
+	"github.com/juju/errors"
+	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/newconfig"
+)
+
+const ConnIDLength = 8
+
+type ConnID [ConnIDLength]byte
+
+func (c ConnID) String() string {
+	return hex.EncodeToString(c[:])
+}
+
+type connPurpose uint8
+
+const (
+	connPurposeClient connPurpose = 1 << iota
+	connPurposeTelegram
+)
+
+const (
+	connTimeoutRead  = 2 * time.Minute
+	connTimeoutWrite = 2 * time.Minute
+)
+
+type wrapperConn struct {
+	parent     net.Conn
+	ctx        context.Context
+	cancel     context.CancelFunc
+	connID     ConnID
+	logger     *zap.SugaredLogger
+	localAddr  *net.TCPAddr
+	remoteAddr *net.TCPAddr
+}
+
+func (w *wrapperConn) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
+	select {
+	case <-w.ctx.Done():
+		w.Close()
+		return 0, errors.Annotate(w.ctx.Err(), "Cannot write because context was closed")
+
+	default:
+		if err := w.parent.SetWriteDeadline(time.Now().Add(timeout)); err != nil {
+			w.Close() // nolint: gosec
+			return 0, errors.Annotate(err, "Cannot set write deadline to the socket")
+		}
+
+		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) Write(p []byte) (int, error) {
+	return w.WriteTimeout(p, connTimeoutWrite)
+}
+
+func (w *wrapperConn) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
+	select {
+	case <-w.ctx.Done():
+		w.Close()
+		return 0, errors.Annotate(w.ctx.Err(), "Cannot read because context was closed")
+
+	default:
+		if err := w.parent.SetReadDeadline(time.Now().Add(timeout)); err != nil {
+			w.Close()
+			return 0, errors.Annotate(err, "Cannot set read deadline to the socket")
+		}
+
+		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) Read(p []byte) (int, error) {
+	return w.ReadTimeout(p, connTimeoutRead)
+}
+
+func (w *wrapperConn) Close() error {
+	w.logger.Debugw("Close connection")
+	w.cancel()
+
+	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(ctx context.Context,
+	cancel context.CancelFunc,
+	parent net.Conn,
+	connID ConnID,
+	purpose connPurpose) StreamReadWriteCloser {
+	localAddr := *parent.LocalAddr().(*net.TCPAddr)
+
+	if parent.RemoteAddr().(*net.TCPAddr).IP.To4() != nil {
+		if newconfig.C.PublicIPv4Addr.IP != nil {
+			localAddr.IP = newconfig.C.PublicIPv4Addr.IP
+		}
+	} else if newconfig.C.PublicIPv6Addr.IP != nil {
+		localAddr.IP = newconfig.C.PublicIPv6Addr.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,
+		ctx:        ctx,
+		cancel:     cancel,
+		connID:     connID,
+		logger:     logger,
+		remoteAddr: parent.RemoteAddr().(*net.TCPAddr),
+		localAddr:  &localAddr,
+	}
+}
+
+func NewClientConn(ctx context.Context,
+	cancel context.CancelFunc,
+	parent net.Conn,
+	connID ConnID) StreamReadWriteCloser {
+	return newConn(ctx, cancel, parent, connID, connPurposeClient)
+}
+
+func NewTelegramConn(ctx context.Context,
+	cancel context.CancelFunc,
+	parent net.Conn,
+	connID ConnID) StreamReadWriteCloser {
+	return newConn(ctx, cancel, parent, connID, connPurposeTelegram)
+}
+
+func NewConnID() ConnID {
+	var id ConnID
+
+	if _, err := rand.Read(id[:]); err != nil {
+		panic(err)
+	}
+
+	return id
+}
@@ -0,0 +1,80 @@
+package newwrappers
+
+import (
+	"crypto/cipher"
+	"net"
+	"time"
+
+	"github.com/juju/errors"
+	"go.uber.org/zap"
+)
+
+type wrapperObfuscated2 struct {
+	encryptor cipher.Stream
+	decryptor cipher.Stream
+	parent    StreamReadWriteCloser
+}
+
+func (w *wrapperObfuscated2) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
+	n, err := w.parent.ReadTimeout(p, timeout)
+	if err != nil {
+		return 0, errors.Annotate(err, "Cannot read stream ciphered data")
+	}
+	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 0, errors.Annotate(err, "Cannot read stream ciphered data")
+	}
+	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 StreamReadWriteCloser, encryptor, decryptor cipher.Stream) StreamReadWriteCloser {
+	return &wrapperObfuscated2{
+		parent:    socket,
+		encryptor: encryptor,
+		decryptor: decryptor,
+	}
+}
@@ -0,0 +1,66 @@
+package newwrappers
+
+import (
+	"net"
+	"time"
+
+	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/newstats"
+)
+
+type wrapperStats struct {
+	parent StreamReadWriteCloser
+}
+
+func (w *wrapperStats) Write(p []byte) (int, error) {
+	n, err := w.parent.Write(p)
+	newstats.S.EgressTraffic(n)
+
+	return n, err
+}
+
+func (w *wrapperStats) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
+	n, err := w.parent.WriteTimeout(p, timeout)
+	newstats.S.EgressTraffic(n)
+
+	return n, err
+}
+
+func (w *wrapperStats) Read(p []byte) (int, error) {
+	n, err := w.parent.Read(p)
+	newstats.S.IngressTraffic(n)
+
+	return n, err
+}
+
+func (w *wrapperStats) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
+	n, err := w.parent.ReadTimeout(p, timeout)
+	newstats.S.IngressTraffic(n)
+
+	return n, err
+}
+
+func (w *wrapperStats) Conn() net.Conn {
+	return w.parent.Conn()
+}
+
+func (w *wrapperStats) Logger() *zap.SugaredLogger {
+	return w.parent.Logger().Named("traffic")
+}
+
+func (w *wrapperStats) LocalAddr() *net.TCPAddr {
+	return w.parent.LocalAddr()
+}
+
+func (w *wrapperStats) RemoteAddr() *net.TCPAddr {
+	return w.parent.RemoteAddr()
+}
+
+func (w *wrapperStats) Close() error {
+	return w.parent.Close()
+}
+
+func NewTraffic(parent StreamReadWriteCloser) StreamReadWriteCloser {
+	return &wrapperStats{parent}
+}