mirror of
https://github.com/ScuroNeko/mtg.git
synced 2026-09-01 01:54:02 +03:00
Refactored all the things!
This commit is contained in:
9seconds
@@ -0,0 +1,66 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net"
|
||||
"sync"
|
||||
|
||||
"github.com/juju/errors"
|
||||
|
||||
"github.com/9seconds/mtg/client"
|
||||
"github.com/9seconds/mtg/config"
|
||||
"github.com/9seconds/mtg/mtproto"
|
||||
"github.com/9seconds/mtg/telegram"
|
||||
"github.com/9seconds/mtg/wrappers"
|
||||
)
|
||||
|
||||
func NewProxyDirect(conf *config.Config) *Proxy {
|
||||
tg := telegram.NewDirectTelegram(conf)
|
||||
|
||||
return &Proxy{
|
||||
conf: conf,
|
||||
acceptCallback: func(ctx context.Context, cancel context.CancelFunc, clientSocket net.Conn,
|
||||
connID string, wait *sync.WaitGroup, conf *config.Config) error {
|
||||
client, opts, err := client.DirectInit(ctx, cancel, clientSocket, connID, conf)
|
||||
if err != nil {
|
||||
return errors.Annotate(err, "Cannot initialize client connection")
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
server, err := directTelegramStream(ctx, cancel, opts, connID, tg)
|
||||
if err != nil {
|
||||
return errors.Annotate(err, "Cannot initialize telegram connection")
|
||||
}
|
||||
defer server.Close()
|
||||
|
||||
wait.Add(2)
|
||||
|
||||
go directPipe(client, server, wait)
|
||||
go directPipe(server, client, wait)
|
||||
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func directTelegramStream(ctx context.Context, cancel context.CancelFunc, opts *mtproto.ConnectionOpts,
|
||||
connID string, tg *telegram.DirectTelegram) (wrappers.WrapStreamReadWriteCloser, error) {
|
||||
streamConn, err := tg.Dial(connID, opts)
|
||||
if err != nil {
|
||||
return nil, errors.Annotate(err, "Cannot dial to Telegram")
|
||||
}
|
||||
streamConn = wrappers.NewCtx(ctx, cancel, streamConn)
|
||||
|
||||
packetConn, err := tg.Init(opts, streamConn)
|
||||
if err != nil {
|
||||
return nil, errors.Annotate(err, "Cannot handshake telegram")
|
||||
}
|
||||
|
||||
return packetConn, nil
|
||||
}
|
||||
|
||||
func directPipe(src io.Reader, dst io.Writer, wait *sync.WaitGroup) {
|
||||
defer wait.Done()
|
||||
io.Copy(dst, src)
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"sync"
|
||||
|
||||
"github.com/juju/errors"
|
||||
|
||||
"github.com/9seconds/mtg/client"
|
||||
"github.com/9seconds/mtg/config"
|
||||
"github.com/9seconds/mtg/mtproto"
|
||||
"github.com/9seconds/mtg/telegram"
|
||||
"github.com/9seconds/mtg/wrappers"
|
||||
)
|
||||
|
||||
func NewProxyMiddle(conf *config.Config) *Proxy {
|
||||
tg := telegram.NewMiddleTelegram(conf)
|
||||
|
||||
return &Proxy{
|
||||
conf: conf,
|
||||
acceptCallback: func(ctx context.Context, cancel context.CancelFunc, clientSocket net.Conn,
|
||||
connID string, wait *sync.WaitGroup, conf *config.Config) error {
|
||||
client, opts, err := client.MiddleInit(ctx, cancel, clientSocket, connID, conf)
|
||||
if err != nil {
|
||||
return errors.Annotate(err, "Cannot initialize client connection")
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
server, err := middleTelegramStream(ctx, cancel, opts, connID, tg)
|
||||
if err != nil {
|
||||
return errors.Annotate(err, "Cannot initialize telegram connection")
|
||||
}
|
||||
defer server.Close()
|
||||
|
||||
wait.Add(2)
|
||||
|
||||
go middlePipe(client, server, wait, &opts.ReadHacks)
|
||||
go middlePipe(server, client, wait, &opts.WriteHacks)
|
||||
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func middleTelegramStream(ctx context.Context, cancel context.CancelFunc, opts *mtproto.ConnectionOpts,
|
||||
connID string, tg *telegram.MiddleTelegram) (wrappers.WrapPacketReadWriteCloser, error) {
|
||||
streamConn, err := tg.Dial(connID, opts)
|
||||
if err != nil {
|
||||
return nil, errors.Annotate(err, "Cannot dial to Telegram")
|
||||
}
|
||||
streamConn = wrappers.NewCtx(ctx, cancel, streamConn)
|
||||
|
||||
packetConn, err := tg.Init(opts, streamConn)
|
||||
if err != nil {
|
||||
return nil, errors.Annotate(err, "Cannot handshake telegram")
|
||||
}
|
||||
|
||||
return packetConn, nil
|
||||
}
|
||||
|
||||
func middlePipe(src wrappers.WrapPacketReader, dst wrappers.WrapPacketWriter, wait *sync.WaitGroup, hacks *mtproto.Hacks) {
|
||||
defer wait.Done()
|
||||
|
||||
for {
|
||||
hacks.SimpleAck = false
|
||||
hacks.QuickAck = false
|
||||
|
||||
packet, err := src.Read()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if _, err = dst.Write(packet); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"sync"
|
||||
|
||||
"github.com/juju/errors"
|
||||
uuid "github.com/satori/go.uuid"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/9seconds/mtg/config"
|
||||
)
|
||||
|
||||
type proxyAcceptCallback func(context.Context, context.CancelFunc, net.Conn, string, *sync.WaitGroup, *config.Config) error
|
||||
|
||||
type Proxy struct {
|
||||
conf *config.Config
|
||||
acceptCallback proxyAcceptCallback
|
||||
}
|
||||
|
||||
func (p *Proxy) Serve() error {
|
||||
lsock, err := net.Listen("tcp", p.conf.BindAddr())
|
||||
if err != nil {
|
||||
return errors.Annotate(err, "Cannot create listen socket")
|
||||
}
|
||||
|
||||
for {
|
||||
if conn, err := lsock.Accept(); err != nil {
|
||||
zap.S().Errorw("Cannot allocate incoming connection", "error", err)
|
||||
} else {
|
||||
go p.accept(conn)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Proxy) accept(conn net.Conn) {
|
||||
connID := uuid.NewV4().String()
|
||||
log := zap.S().With("connection_id", connID)
|
||||
|
||||
defer func() {
|
||||
conn.Close()
|
||||
|
||||
if err := recover(); err != nil {
|
||||
log.Errorw("Crash of accept handler", "error", err)
|
||||
}
|
||||
}()
|
||||
|
||||
log.Infow("Client connected", "addr", conn.RemoteAddr())
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
wait := &sync.WaitGroup{}
|
||||
|
||||
if err := p.acceptCallback(ctx, cancel, conn, connID, wait, p.conf); err != nil {
|
||||
log.Errorw("Cannot initialize connection", "error", err)
|
||||
cancel()
|
||||
}
|
||||
|
||||
<-ctx.Done()
|
||||
wait.Wait()
|
||||
|
||||
log.Infow("Client disconnected", "addr", conn.RemoteAddr())
|
||||
}
|
||||
-184
@@ -1,184 +0,0 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net"
|
||||
"sync"
|
||||
|
||||
"github.com/juju/errors"
|
||||
uuid "github.com/satori/go.uuid"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/9seconds/mtg/client"
|
||||
"github.com/9seconds/mtg/config"
|
||||
"github.com/9seconds/mtg/mtproto"
|
||||
"github.com/9seconds/mtg/telegram"
|
||||
"github.com/9seconds/mtg/utils"
|
||||
"github.com/9seconds/mtg/wrappers"
|
||||
)
|
||||
|
||||
// Server is an insgtance of MTPROTO proxy.
|
||||
type Server struct {
|
||||
conf *config.Config
|
||||
logger *zap.SugaredLogger
|
||||
stats *Stats
|
||||
tg telegram.Telegram
|
||||
clientInit client.Init
|
||||
}
|
||||
|
||||
// Serve does MTPROTO proxying.
|
||||
func (s *Server) Serve() error {
|
||||
lsock, err := net.Listen("tcp", s.conf.BindAddr())
|
||||
if err != nil {
|
||||
return errors.Annotate(err, "Cannot create listen socket")
|
||||
}
|
||||
|
||||
for {
|
||||
if conn, err := lsock.Accept(); err != nil {
|
||||
s.logger.Warn("Cannot allocate incoming connection", "error", err)
|
||||
} else {
|
||||
go s.accept(conn)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) accept(conn net.Conn) {
|
||||
defer func() {
|
||||
s.stats.closeConnection()
|
||||
conn.Close() // nolint: errcheck
|
||||
|
||||
if r := recover(); r != nil {
|
||||
s.logger.Errorw("Crash of accept handler", "error", r)
|
||||
}
|
||||
}()
|
||||
|
||||
s.stats.newConnection()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
socketID := uuid.NewV4().String()
|
||||
|
||||
s.logger.Debugw("Client connected",
|
||||
"addr", conn.RemoteAddr().String(),
|
||||
"socketid", socketID,
|
||||
)
|
||||
|
||||
connOpts, clientConn, err := s.getClientStream(ctx, cancel, conn, socketID)
|
||||
if err != nil {
|
||||
s.logger.Warnw("Cannot initialize client connection",
|
||||
"addr", conn.RemoteAddr().String(),
|
||||
"socketid", socketID,
|
||||
"error", err,
|
||||
)
|
||||
return
|
||||
}
|
||||
defer clientConn.Close() // nolint: errcheck
|
||||
|
||||
tgConn, err := s.getTelegramStream(ctx, cancel, connOpts, socketID)
|
||||
if err != nil {
|
||||
s.logger.Warnw("Cannot initialize Telegram connection",
|
||||
"socketid", socketID,
|
||||
"error", err,
|
||||
)
|
||||
return
|
||||
}
|
||||
defer tgConn.Close() // nolint: errcheck
|
||||
|
||||
wait := &sync.WaitGroup{}
|
||||
wait.Add(2)
|
||||
|
||||
go func() {
|
||||
defer wait.Done()
|
||||
|
||||
for {
|
||||
connOpts.ReadHacks.QuickAck = false
|
||||
connOpts.ReadHacks.SimpleAck = false
|
||||
if err := s.pump(clientConn, tgConn, socketID, "client"); err != nil {
|
||||
s.logger.Infow("Client stream is aborted",
|
||||
"socketid", socketID, "error", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
go func() {
|
||||
defer wait.Done()
|
||||
|
||||
for {
|
||||
connOpts.WriteHacks.QuickAck = false
|
||||
connOpts.WriteHacks.SimpleAck = false
|
||||
if err := s.pump(tgConn, clientConn, socketID, "telegram"); err != nil {
|
||||
s.logger.Infow("Telegram stream is aborted",
|
||||
"socketid", socketID, "error", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
<-ctx.Done()
|
||||
wait.Wait()
|
||||
|
||||
s.logger.Debugw("Client disconnected",
|
||||
"addr", conn.RemoteAddr().String(),
|
||||
"socketid", socketID,
|
||||
)
|
||||
}
|
||||
|
||||
func (s *Server) getClientStream(ctx context.Context, cancel context.CancelFunc, conn net.Conn, socketID string) (*mtproto.ConnectionOpts, io.ReadWriteCloser, error) {
|
||||
socket, connOpts, err := s.clientInit(conn, socketID, s.conf)
|
||||
if err != nil {
|
||||
return nil, nil, errors.Annotate(err, "Cannot init client connection")
|
||||
}
|
||||
|
||||
socket = wrappers.NewTrafficRWC(socket, s.stats.addIncomingTraffic, s.stats.addOutgoingTraffic)
|
||||
socket = wrappers.NewLogRWC(socket, s.logger, socketID, "client")
|
||||
socket = wrappers.NewCtxRWC(ctx, cancel, socket)
|
||||
|
||||
return connOpts, socket, nil
|
||||
}
|
||||
|
||||
func (s *Server) getTelegramStream(ctx context.Context, cancel context.CancelFunc, connOpts *mtproto.ConnectionOpts, socketID string) (io.ReadWriteCloser, error) {
|
||||
conn, err := s.tg.Dial(socketID, connOpts)
|
||||
if err != nil {
|
||||
return nil, errors.Annotate(err, "Cannot connect to Telegram")
|
||||
}
|
||||
|
||||
conn = wrappers.NewTrafficRWC(conn, s.stats.addIncomingTraffic, s.stats.addOutgoingTraffic)
|
||||
conn, err = s.tg.Init(connOpts, conn)
|
||||
if err != nil {
|
||||
return nil, errors.Annotate(err, "Cannot handshake Telegram")
|
||||
}
|
||||
|
||||
conn = wrappers.NewLogRWC(conn, s.logger, socketID, "telegram")
|
||||
conn = wrappers.NewCtxRWC(ctx, cancel, conn)
|
||||
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func (s *Server) pump(src io.Reader, dst io.Writer, socketID, name string) error {
|
||||
buf, err := utils.ReadCurrentData(src)
|
||||
if err != nil {
|
||||
return errors.Annotate(err, "Cannot pump the socket")
|
||||
}
|
||||
|
||||
_, err = dst.Write(buf)
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
// NewServer creates new instance of MTPROTO proxy.
|
||||
func NewServer(conf *config.Config, logger *zap.SugaredLogger, stat *Stats) *Server {
|
||||
clientInit := client.DirectInit
|
||||
tg := telegram.NewDirectTelegram
|
||||
|
||||
if len(conf.AdTag) > 0 {
|
||||
clientInit = client.MiddleInit
|
||||
tg = telegram.NewMiddleTelegram
|
||||
}
|
||||
|
||||
return &Server{
|
||||
conf: conf,
|
||||
logger: logger,
|
||||
stats: stat,
|
||||
tg: tg(conf, logger),
|
||||
clientInit: clientInit,
|
||||
}
|
||||
}
|
||||
@@ -1,74 +0,0 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/9seconds/mtg/config"
|
||||
)
|
||||
|
||||
type statsUptime time.Time
|
||||
|
||||
func (s statsUptime) MarshalJSON() ([]byte, error) {
|
||||
uptime := int(time.Since(time.Time(s)).Seconds())
|
||||
return []byte(strconv.Itoa(uptime)), nil
|
||||
}
|
||||
|
||||
// Stats is a datastructure for statistics on work of this proxy.
|
||||
type Stats struct {
|
||||
AllConnections uint64 `json:"all_connections"`
|
||||
ActiveConnections uint32 `json:"active_connections"`
|
||||
Traffic struct {
|
||||
Incoming uint64 `json:"incoming"`
|
||||
Outgoing uint64 `json:"outgoing"`
|
||||
} `json:"traffic"`
|
||||
URLs config.IPURLs `json:"urls"`
|
||||
Uptime statsUptime `json:"uptime"`
|
||||
|
||||
conf *config.Config
|
||||
}
|
||||
|
||||
func (s *Stats) newConnection() {
|
||||
atomic.AddUint64(&s.AllConnections, 1)
|
||||
atomic.AddUint32(&s.ActiveConnections, 1)
|
||||
}
|
||||
|
||||
func (s *Stats) closeConnection() {
|
||||
atomic.AddUint32(&s.ActiveConnections, ^uint32(0))
|
||||
}
|
||||
|
||||
func (s *Stats) addIncomingTraffic(n int) {
|
||||
atomic.AddUint64(&s.Traffic.Incoming, uint64(n))
|
||||
}
|
||||
|
||||
func (s *Stats) addOutgoingTraffic(n int) {
|
||||
atomic.AddUint64(&s.Traffic.Outgoing, uint64(n))
|
||||
}
|
||||
|
||||
// Serve runs statistics HTTP server.
|
||||
func (s *Stats) Serve() {
|
||||
http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
|
||||
encoder := json.NewEncoder(w)
|
||||
encoder.SetEscapeHTML(false)
|
||||
encoder.SetIndent("", " ")
|
||||
encoder.Encode(s) // nolint: errcheck, gas
|
||||
})
|
||||
|
||||
http.ListenAndServe(s.conf.StatAddr(), nil) // nolint: errcheck, gas
|
||||
}
|
||||
|
||||
// NewStats returns new instance of statistics datastructure.
|
||||
func NewStats(conf *config.Config) *Stats {
|
||||
stat := &Stats{
|
||||
Uptime: statsUptime(time.Now()),
|
||||
conf: conf,
|
||||
}
|
||||
stat.URLs = conf.GetURLs()
|
||||
|
||||
return stat
|
||||
}
|
||||
Reference in New Issue
Block a user