diff --git a/client/client.go b/client/client.go deleted file mode 100644 index 318c476..0000000 --- a/client/client.go +++ /dev/null @@ -1,12 +0,0 @@ -package client - -import ( - "net" - - "github.com/9seconds/mtg/config" - "github.com/9seconds/mtg/mtproto" - "github.com/9seconds/mtg/wrappers" -) - -// Init has to initialize client connection based on given config. -type Init func(net.Conn, string, *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) diff --git a/client/direct.go b/client/direct.go index 7ee7e75..5abe56c 100644 --- a/client/direct.go +++ b/client/direct.go @@ -1,6 +1,7 @@ package client import ( + "context" "net" "time" @@ -14,28 +15,29 @@ import ( const handshakeTimeout = 10 * time.Second -// DirectInit initializes client to access Telegram bypassing middleproxies. -func DirectInit(conn net.Conn, socketID string, conf *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) { - if err := config.SetSocketOptions(conn); err != nil { +func DirectInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn, connID string, + conf *config.Config) (wrappers.WrapStreamReadWriteCloser, *mtproto.ConnectionOpts, error) { + if err := config.SetSocketOptions(socket); err != nil { return nil, nil, errors.Annotate(err, "Cannot set socket options") } - conn.SetReadDeadline(time.Now().Add(handshakeTimeout)) // nolint: errcheck - frame, err := obfuscated2.ExtractFrame(conn) - conn.SetReadDeadline(time.Time{}) // nolint: errcheck + socket.SetReadDeadline(time.Now().Add(handshakeTimeout)) + frame, err := obfuscated2.ExtractFrame(socket) if err != nil { return nil, nil, errors.Annotate(err, "Cannot extract frame") } + socket.SetReadDeadline(time.Time{}) + conn := wrappers.NewConn(socket, connID, wrappers.ConnPurposeClient, conf.PublicIPv4, conf.PublicIPv6) obfs2, connOpts, err := obfuscated2.ParseObfuscated2ClientFrame(conf.Secret, frame) if err != nil { return nil, nil, errors.Annotate(err, "Cannot parse obfuscated frame") } connOpts.ConnectionProto = mtproto.ConnectionProtocolAny - connOpts.ClientAddr = conn.RemoteAddr().(*net.TCPAddr) + connOpts.ClientAddr = conn.RemoteAddr() - socket := wrappers.NewTimeoutRWC(conn, socketID, conf.PublicIPv4, conf.PublicIPv6) - socket = wrappers.NewStreamCipherRWC(socket, obfs2.Encryptor, obfs2.Decryptor) + conn = wrappers.NewCtx(ctx, cancel, conn) + conn = wrappers.NewStreamCipher(conn, obfs2.Encryptor, obfs2.Decryptor) - return socket, connOpts, nil + return conn, connOpts, nil } diff --git a/client/middle.go b/client/middle.go index 2c978ec..635277d 100644 --- a/client/middle.go +++ b/client/middle.go @@ -1,30 +1,30 @@ package client import ( + "context" "net" "github.com/9seconds/mtg/config" "github.com/9seconds/mtg/mtproto" - mtwrappers "github.com/9seconds/mtg/mtproto/wrappers" "github.com/9seconds/mtg/wrappers" ) -func MiddleInit(conn net.Conn, socketID string, conf *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) { - newConn, opts, err := DirectInit(conn, socketID, conf) +func MiddleInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn, connID string, + conf *config.Config) (wrappers.WrapPacketReadWriteCloser, *mtproto.ConnectionOpts, error) { + conn, opts, err := DirectInit(ctx, cancel, socket, connID, conf) if err != nil { return nil, nil, err } - if opts.ConnectionType == mtproto.ConnectionTypeAbridged { - newConn = mtwrappers.NewAbridgedRWC(newConn, opts) - } else { - newConn = mtwrappers.NewIntermediateRWC(newConn, opts) + newConn := wrappers.NewMTProtoAbridged(conn, opts) + if opts.ConnectionType != mtproto.ConnectionTypeAbridged { + newConn = wrappers.NewMTProtoIntermediate(conn, opts) } opts.ConnectionProto = mtproto.ConnectionProtocolIPv4 - if conn.LocalAddr().(*net.TCPAddr).IP.To4() == nil { + if socket.LocalAddr().(*net.TCPAddr).IP.To4() == nil { opts.ConnectionProto = mtproto.ConnectionProtocolIPv6 } - return newConn, opts, nil + return newConn, opts, err } diff --git a/config/config.go b/config/config.go index b5cfe9e..1a8f95e 100644 --- a/config/config.go +++ b/config/config.go @@ -5,7 +5,6 @@ import ( "fmt" "net" "strconv" - "time" "github.com/juju/errors" ) @@ -14,10 +13,6 @@ import ( const ( BufferWriteSize = 32 * 1024 BufferReadSize = 32 * 1024 - BufferSizeCopy = 32 * 1024 - - TimeoutRead = time.Minute - TimeoutWrite = time.Minute ) // Config represents common configuration of mtg. diff --git a/main.go b/main.go index 14cf85a..5cc71be 100644 --- a/main.go +++ b/main.go @@ -111,16 +111,21 @@ func main() { zapcore.NewJSONEncoder(encoderCfg), zapcore.Lock(os.Stderr), atom, - )).Sugar() + )) + zap.ReplaceGlobals(logger) + defer logger.Sync() - stat := proxy.NewStats(conf) - go stat.Serve() + var server *proxy.Proxy + if len(conf.AdTag) == 0 { + server = proxy.NewProxyDirect(conf) + } else { + server = proxy.NewProxyMiddle(conf) + } - srv := proxy.NewServer(conf, logger, stat) printURLs(conf.GetURLs()) - if err := srv.Serve(); err != nil { - logger.Fatal(err.Error()) + if err := server.Serve(); err != nil { + zap.S().Fatalw("Server stopped", "error", err) } } diff --git a/obfuscated2/frame.go b/obfuscated2/frame.go index 6c6f6ab..edcb8b4 100644 --- a/obfuscated2/frame.go +++ b/obfuscated2/frame.go @@ -9,6 +9,7 @@ import ( "github.com/juju/errors" "github.com/9seconds/mtg/mtproto" + "github.com/9seconds/mtg/utils" ) // [frameOffsetFirst:frameOffsetKey:frameOffsetIV:frameOffsetMagic:frameOffsetDC:frameOffsetEnd] @@ -67,14 +68,7 @@ func (f Frame) ConnectionType() (mtproto.ConnectionType, error) { // Invert inverts frame for extracting encryption keys. Pkease check that link: // https://blog.susanka.eu/how-telegram-obfuscates-its-mtproto-traffic/ func (f Frame) Invert() Frame { - reversed := make(Frame, FrameLen) - copy(reversed, f) - - for i := 0; i < frameLenKey+frameLenIV; i++ { - reversed[frameOffsetFirst+i] = f[frameOffsetIV-1-i] - } - - return reversed + return Frame(utils.ReverseBytes([]byte(f))) } // ExtractFrame extracts exact obfuscated2 handshake frame from given reader. diff --git a/proxy/direct.go b/proxy/direct.go new file mode 100644 index 0000000..56b6ce2 --- /dev/null +++ b/proxy/direct.go @@ -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) +} diff --git a/proxy/middle.go b/proxy/middle.go new file mode 100644 index 0000000..4987a13 --- /dev/null +++ b/proxy/middle.go @@ -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 + } + } +} diff --git a/proxy/proxy.go b/proxy/proxy.go new file mode 100644 index 0000000..d1baf93 --- /dev/null +++ b/proxy/proxy.go @@ -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()) +} diff --git a/proxy/server.go b/proxy/server.go deleted file mode 100644 index 18bb98e..0000000 --- a/proxy/server.go +++ /dev/null @@ -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, - } -} diff --git a/proxy/stats.go b/proxy/stats.go deleted file mode 100644 index 9469c2f..0000000 --- a/proxy/stats.go +++ /dev/null @@ -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 -} diff --git a/telegram/dialer.go b/telegram/dialer.go index 5f3e20c..d710f38 100644 --- a/telegram/dialer.go +++ b/telegram/dialer.go @@ -30,11 +30,14 @@ func (t *tgDialer) dial(addr string) (net.Conn, error) { return conn, nil } -func (t *tgDialer) dialRWC(addr, sock string) (wrappers.ReadWriteCloserWithAddr, error) { +func (t *tgDialer) dialRWC(addr, connID string) (wrappers.WrapStreamReadWriteCloser, error) { conn, err := t.dial(addr) if err != nil { return nil, err } - return wrappers.NewTimeoutRWC(conn, sock, t.conf.PublicIPv4, t.conf.PublicIPv6), nil + tgConn := wrappers.NewConn(conn, connID, wrappers.ConnPurposeTelegram, + t.conf.PublicIPv4, t.conf.PublicIPv6) + + return tgConn, nil } diff --git a/telegram/direct.go b/telegram/direct.go index 601fa52..0d1dd28 100644 --- a/telegram/direct.go +++ b/telegram/direct.go @@ -4,7 +4,6 @@ import ( "net" "github.com/juju/errors" - "go.uber.org/zap" "github.com/9seconds/mtg/config" "github.com/9seconds/mtg/mtproto" @@ -29,11 +28,11 @@ var ( } ) -type directTelegram struct { +type DirectTelegram struct { baseTelegram } -func (t *directTelegram) Dial(sock string, connOpts *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error) { +func (t *DirectTelegram) Dial(connID string, connOpts *mtproto.ConnectionOpts) (wrappers.WrapStreamReadWriteCloser, error) { dc := connOpts.DC if dc < 0 { dc = -dc @@ -41,23 +40,23 @@ func (t *directTelegram) Dial(sock string, connOpts *mtproto.ConnectionOpts) (wr dc = 1 } - return t.baseTelegram.dial(dc-1, sock, connOpts.ConnectionProto) + return t.baseTelegram.dial(dc-1, connID, connOpts.ConnectionProto) } -func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.ReadWriteCloserWithAddr) (wrappers.ReadWriteCloserWithAddr, error) { +func (t *DirectTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.WrapStreamReadWriteCloser) (wrappers.WrapStreamReadWriteCloser, error) { obfs2, frame := obfuscated2.MakeTelegramObfuscated2Frame(connOpts) - if n, err := conn.Write(frame); err != nil || n != obfuscated2.FrameLen { + if _, err := conn.Write(frame); err != nil { return nil, errors.Annotate(err, "Cannot write hadnshake frame") } - return wrappers.NewStreamCipherRWC(conn, obfs2.Encryptor, obfs2.Decryptor), nil + return wrappers.NewStreamCipher(conn, obfs2.Encryptor, obfs2.Decryptor), nil } // NewDirectTelegram returns Telegram instance which connects directly // to Telegram bypassing middleproxies. -func NewDirectTelegram(conf *config.Config, _ *zap.SugaredLogger) Telegram { - return &directTelegram{baseTelegram{ +func NewDirectTelegram(conf *config.Config) *DirectTelegram { + return &DirectTelegram{baseTelegram{ dialer: tgDialer{ Dialer: net.Dialer{Timeout: telegramDialTimeout}, conf: conf, diff --git a/telegram/middle.go b/telegram/middle.go index 62a9c65..8a25996 100644 --- a/telegram/middle.go +++ b/telegram/middle.go @@ -1,29 +1,26 @@ package telegram import ( - "io" "net" "net/http" "sync" "github.com/juju/errors" - "go.uber.org/zap" "github.com/9seconds/mtg/config" "github.com/9seconds/mtg/mtproto" "github.com/9seconds/mtg/mtproto/rpc" - mtwrappers "github.com/9seconds/mtg/mtproto/wrappers" "github.com/9seconds/mtg/wrappers" ) -type middleTelegram struct { +type MiddleTelegram struct { middleTelegramCaller conf *config.Config } -func NewMiddleTelegram(conf *config.Config, logger *zap.SugaredLogger) Telegram { - tg := &middleTelegram{ +func NewMiddleTelegram(conf *config.Config) *MiddleTelegram { + tg := &MiddleTelegram{ middleTelegramCaller: middleTelegramCaller{ baseTelegram: baseTelegram{ dialer: tgDialer{ @@ -31,7 +28,6 @@ func NewMiddleTelegram(conf *config.Config, logger *zap.SugaredLogger) Telegram conf: conf, }, }, - logger: logger, httpClient: &http.Client{ Timeout: middleTelegramHTTPClientTimeout, }, @@ -48,8 +44,8 @@ func NewMiddleTelegram(conf *config.Config, logger *zap.SugaredLogger) Telegram return tg } -func (t *middleTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.ReadWriteCloserWithAddr) (wrappers.ReadWriteCloserWithAddr, error) { - rpcNonceConn := mtwrappers.NewFrameRWC(conn, rpc.SeqNoNonce) +func (t *MiddleTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.WrapStreamReadWriteCloser) (wrappers.WrapPacketReadWriteCloser, error) { + rpcNonceConn := wrappers.NewMTProtoFrame(conn, rpc.SeqNoNonce) rpcNonceReq, err := t.sendRPCNonceRequest(rpcNonceConn) if err != nil { @@ -60,22 +56,22 @@ func (t *middleTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.Re return nil, err } - secureConn := mtwrappers.NewMiddleProxyCipherRWC(conn, rpcNonceReq, rpcNonceResp, t.proxySecret) - secureConn = mtwrappers.NewFrameRWC(secureConn, rpc.SeqNoHandshake) + secureConn := wrappers.NewMiddleProxyCipher(conn, rpcNonceReq, rpcNonceResp, t.proxySecret) + frameConn := wrappers.NewMTProtoFrame(secureConn, rpc.SeqNoHandshake) - rpcHandshakeReq, err := t.sendRPCHandshakeRequest(secureConn) + rpcHandshakeReq, err := t.sendRPCHandshakeRequest(frameConn) if err != nil { return nil, err } - _, err = t.receiveRPCHandshakeResponse(secureConn, rpcHandshakeReq) + _, err = t.receiveRPCHandshakeResponse(frameConn, rpcHandshakeReq) if err != nil { return nil, err } - return mtwrappers.NewProxyRequestRWC(secureConn, connOpts, t.conf.AdTag) + return wrappers.NewMTProtoProxy(frameConn, connOpts, t.conf.AdTag) } -func (t *middleTelegram) sendRPCNonceRequest(conn io.Writer) (*rpc.NonceRequest, error) { +func (t *MiddleTelegram) sendRPCNonceRequest(conn wrappers.WrapPacketWriter) (*rpc.NonceRequest, error) { rpcNonceReq, err := rpc.NewNonceRequest(t.proxySecret) if err != nil { return nil, errors.Annotate(err, "Cannot create RPC nonce request") @@ -87,15 +83,13 @@ func (t *middleTelegram) sendRPCNonceRequest(conn io.Writer) (*rpc.NonceRequest, return rpcNonceReq, nil } -func (t *middleTelegram) receiveRPCNonceResponse(conn io.Reader, req *rpc.NonceRequest) (*rpc.NonceResponse, error) { - var ans [128]byte - - n, err := conn.Read(ans[:]) +func (t *MiddleTelegram) receiveRPCNonceResponse(conn wrappers.WrapPacketReader, req *rpc.NonceRequest) (*rpc.NonceResponse, error) { + packet, err := conn.Read() if err != nil { return nil, errors.Annotate(err, "Cannot read RPC nonce response") } - rpcNonceResp, err := rpc.NewNonceResponse(ans[:n]) + rpcNonceResp, err := rpc.NewNonceResponse(packet) if err != nil { return nil, errors.Annotate(err, "Cannot initialize RPC nonce response") } @@ -106,7 +100,7 @@ func (t *middleTelegram) receiveRPCNonceResponse(conn io.Reader, req *rpc.NonceR return rpcNonceResp, nil } -func (t *middleTelegram) sendRPCHandshakeRequest(conn io.Writer) (*rpc.HandshakeRequest, error) { +func (t *MiddleTelegram) sendRPCHandshakeRequest(conn wrappers.WrapPacketWriter) (*rpc.HandshakeRequest, error) { req := rpc.NewHandshakeRequest() if _, err := conn.Write(req.Bytes()); err != nil { return nil, errors.Annotate(err, "Cannot send RPC handshake request") @@ -115,15 +109,13 @@ func (t *middleTelegram) sendRPCHandshakeRequest(conn io.Writer) (*rpc.Handshake return req, nil } -func (t *middleTelegram) receiveRPCHandshakeResponse(conn io.Reader, req *rpc.HandshakeRequest) (*rpc.HandshakeResponse, error) { - var ans [128]byte - - n, err := conn.Read(ans[:]) +func (t *MiddleTelegram) receiveRPCHandshakeResponse(conn wrappers.WrapPacketReader, req *rpc.HandshakeRequest) (*rpc.HandshakeResponse, error) { + packet, err := conn.Read() if err != nil { return nil, errors.Annotate(err, "Cannot read RPC handshake response") } - rpcHandshakeResp, err := rpc.NewHandshakeResponse(ans[:n]) + rpcHandshakeResp, err := rpc.NewHandshakeResponse(packet) if err != nil { return nil, errors.Annotate(err, "Cannot initialize RPC handshake response") } diff --git a/telegram/middle_caller.go b/telegram/middle_caller.go index 490ed26..c0d40a2 100644 --- a/telegram/middle_caller.go +++ b/telegram/middle_caller.go @@ -35,11 +35,10 @@ type middleTelegramCaller struct { proxySecret []byte dialerMutex *sync.RWMutex - logger *zap.SugaredLogger httpClient *http.Client } -func (t *middleTelegramCaller) Dial(sock string, connOpts *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error) { +func (t *middleTelegramCaller) Dial(connID string, connOpts *mtproto.ConnectionOpts) (wrappers.WrapStreamReadWriteCloser, error) { dc := connOpts.DC if dc == 0 { dc = 1 @@ -47,13 +46,13 @@ func (t *middleTelegramCaller) Dial(sock string, connOpts *mtproto.ConnectionOpt t.dialerMutex.RLock() defer t.dialerMutex.RUnlock() - return t.baseTelegram.dial(dc, sock, connOpts.ConnectionProto) + return t.baseTelegram.dial(dc, connID, connOpts.ConnectionProto) } func (t *middleTelegramCaller) autoUpdate() { for range time.Tick(middleTelegramAutoUpdateInterval) { if err := t.update(); err != nil { - t.logger.Warnw("Cannot update from Telegram", "error", err) + zap.S().Warnw("Cannot update from Telegram", "error", err) } } } @@ -80,7 +79,7 @@ func (t *middleTelegramCaller) update() error { t.v6Addresses = v6Addresses t.dialerMutex.Unlock() - t.logger.Infow("Telegram middle proxy data has been updated") + zap.S().Infow("Telegram middle proxy data has been updated") return nil } diff --git a/telegram/telegram.go b/telegram/telegram.go index 9413882..da8e27a 100644 --- a/telegram/telegram.go +++ b/telegram/telegram.go @@ -9,14 +9,6 @@ import ( "github.com/9seconds/mtg/wrappers" ) -// Telegram defines an interface to connect to Telegram. This -// encapsulates logic of working with middleproxies or direct -// connections. -type Telegram interface { - Dial(string, *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error) - Init(*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr) (wrappers.ReadWriteCloserWithAddr, error) -} - type baseTelegram struct { dialer tgDialer @@ -24,7 +16,7 @@ type baseTelegram struct { v6Addresses map[int16][]string } -func (b *baseTelegram) dial(dcIdx int16, sock string, proto mtproto.ConnectionProtocol) (wrappers.ReadWriteCloserWithAddr, error) { +func (b *baseTelegram) dial(dcIdx int16, connID string, proto mtproto.ConnectionProtocol) (wrappers.WrapStreamReadWriteCloser, error) { addrs := make([]string, 2) if proto&mtproto.ConnectionProtocolIPv6 != 0 { @@ -39,7 +31,7 @@ func (b *baseTelegram) dial(dcIdx int16, sock string, proto mtproto.ConnectionPr } for _, addr := range addrs { - if conn, err := b.dialer.dialRWC(addr, sock); err == nil { + if conn, err := b.dialer.dialRWC(addr, connID); err == nil { return conn, err } } diff --git a/wrappers/blockcipher.go b/wrappers/blockcipher.go index 55b90e1..5283f90 100644 --- a/wrappers/blockcipher.go +++ b/wrappers/blockcipher.go @@ -1,6 +1,7 @@ package wrappers import ( + "bytes" "crypto/aes" "crypto/cipher" "net" @@ -9,77 +10,89 @@ import ( "github.com/juju/errors" ) -type WrapBlockCipher struct { - BufferedReader +type BlockCipher struct { + buf *bytes.Buffer conn WrapStreamReadWriteCloser encryptor cipher.BlockMode decryptor cipher.BlockMode } -func (w *WrapBlockCipher) Read(p []byte) (int, error) { - return w.BufferedRead(p, func() error { - var buf []byte +func (b *BlockCipher) Read(p []byte) (int, error) { + if b.buf.Len() > 0 { + return b.flush(p) + } - for len(buf) == 0 || len(buf)%aes.BlockSize != 0 { - rv, err := utils.ReadCurrentData(w.conn) - if err != nil { - return errors.Annotate(err, "Cannot read from socket") - } - buf = append(buf, rv...) + buf := []byte{} + for len(buf) == 0 || len(buf)%aes.BlockSize != 0 { + rv, err := utils.ReadCurrentData(b.conn) + if err != nil { + return 0, errors.Annotate(err, "Cannot read from socket") } + buf = append(buf, rv...) + } - w.decryptor.CryptBlocks(buf, buf) - w.Buffer.Write(buf) + b.decryptor.CryptBlocks(buf, buf) + b.buf.Write(buf) - return nil - }) + return b.flush(p) } -func (w *WrapBlockCipher) Write(p []byte) (int, error) { +func (b *BlockCipher) flush(p []byte) (int, error) { + if b.buf.Len() <= len(p) { + sizeToReturn := b.buf.Len() + copy(p, b.buf.Bytes()) + b.buf.Reset() + return sizeToReturn, nil + } + + return b.buf.Read(p) +} + +func (b *BlockCipher) Write(p []byte) (int, error) { if len(p)%aes.BlockSize > 0 { return 0, errors.Errorf("Incorrect block size %d", len(p)) } encrypted := make([]byte, len(p)) - w.encryptor.CryptBlocks(encrypted, p) + b.encryptor.CryptBlocks(encrypted, p) - return w.conn.Write(encrypted) + return b.conn.Write(encrypted) } -func (w *WrapBlockCipher) LogDebug(msg string, data ...interface{}) { - w.conn.LogDebug(msg, data...) +func (b *BlockCipher) LogDebug(msg string, data ...interface{}) { + b.conn.LogDebug(msg, data...) } -func (w *WrapBlockCipher) LogInfo(msg string, data ...interface{}) { - w.conn.LogInfo(msg, data...) +func (b *BlockCipher) LogInfo(msg string, data ...interface{}) { + b.conn.LogInfo(msg, data...) } -func (w *WrapBlockCipher) LogWarn(msg string, data ...interface{}) { - w.conn.LogWarn(msg, data...) +func (b *BlockCipher) LogWarn(msg string, data ...interface{}) { + b.conn.LogWarn(msg, data...) } -func (w *WrapBlockCipher) LogError(msg string, data ...interface{}) { - w.conn.LogError(msg, data...) +func (b *BlockCipher) LogError(msg string, data ...interface{}) { + b.conn.LogError(msg, data...) } -func (w *WrapBlockCipher) LocalAddr() *net.TCPAddr { - return w.conn.LocalAddr() +func (b *BlockCipher) LocalAddr() *net.TCPAddr { + return b.conn.LocalAddr() } -func (w *WrapBlockCipher) RemoteAddr() *net.TCPAddr { - return w.conn.RemoteAddr() +func (b *BlockCipher) RemoteAddr() *net.TCPAddr { + return b.conn.RemoteAddr() } -func (w *WrapBlockCipher) Close() error { - return w.conn.Close() +func (b *BlockCipher) Close() error { + return b.conn.Close() } -func NewWrapBlockCipher(conn WrapStreamReadWriteCloser, encryptor, decryptor cipher.BlockMode) WrapStreamReadWriteCloser { - return &WrapBlockCipher{ - BufferedReader: NewBufferedReader(), - conn: conn, - encryptor: encryptor, - decryptor: decryptor, +func NewBlockCipher(conn WrapStreamReadWriteCloser, encryptor, decryptor cipher.BlockMode) WrapStreamReadWriteCloser { + return &BlockCipher{ + buf: &bytes.Buffer{}, + conn: conn, + encryptor: encryptor, + decryptor: decryptor, } } diff --git a/wrappers/buffered_reader.go b/wrappers/buffered_reader.go deleted file mode 100644 index 379534a..0000000 --- a/wrappers/buffered_reader.go +++ /dev/null @@ -1,32 +0,0 @@ -package wrappers - -import "bytes" - -type BufferedReader struct { - Buffer *bytes.Buffer -} - -func (b *BufferedReader) BufferedRead(p []byte, callback func() error) (int, error) { - if b.Buffer.Len() > 0 { - return b.flush(p) - } - if err := callback(); err != nil { - return 0, err - } - return b.flush(p) -} - -func (b *BufferedReader) flush(p []byte) (int, error) { - if b.Buffer.Len() <= len(p) { - sizeToReturn := b.Buffer.Len() - copy(p, b.Buffer.Bytes()) - b.Buffer.Reset() - return sizeToReturn, nil - } - - return b.Buffer.Read(p) -} - -func NewBufferedReader() BufferedReader { - return BufferedReader{Buffer: &bytes.Buffer{}} -} diff --git a/wrappers/conn.go b/wrappers/conn.go index 468820c..2a84941 100644 --- a/wrappers/conn.go +++ b/wrappers/conn.go @@ -30,8 +30,7 @@ const ( connTimeoutWrite = 5 * time.Minute ) -type WrapConn struct { - purpose ConnPurpose +type Conn struct { connID string conn net.Conn logger *zap.SugaredLogger @@ -39,77 +38,80 @@ type WrapConn struct { publicIPv6 net.IP } -func (w *WrapConn) Write(p []byte) (int, error) { - w.conn.SetWriteDeadline(time.Now().Add(connTimeoutWrite)) - n, err := w.conn.Write(p) +func (c *Conn) Write(p []byte) (int, error) { + c.conn.SetWriteDeadline(time.Now().Add(connTimeoutWrite)) + n, err := c.conn.Write(p) - w.logger.Debugw("Write to stream", "bytes", n, "error", err) + c.logger.Debugw("Write to stream", "bytes", n, "error", err) return n, err } -func (w *WrapConn) Read(p []byte) (int, error) { - w.conn.SetReadDeadline(time.Now().Add(connTimeoutRead)) - n, err := w.conn.Read(p) +func (c *Conn) Read(p []byte) (int, error) { + c.conn.SetReadDeadline(time.Now().Add(connTimeoutRead)) + n, err := c.conn.Read(p) - w.logger.Debugw("Read from stream", "bytes", n, "error", err) + c.logger.Debugw("Read from stream", "bytes", n, "error", err) return n, err } -func (w *WrapConn) Close() error { - defer w.LogDebug("Closed connection") - return w.conn.Close() +func (c *Conn) Close() error { + defer c.LogDebug("Closed connection") + return c.conn.Close() } -func (w *WrapConn) LocalAddr() *net.TCPAddr { - addr := w.conn.LocalAddr().(*net.TCPAddr) +func (c *Conn) LocalAddr() *net.TCPAddr { + addr := c.conn.LocalAddr().(*net.TCPAddr) newAddr := *addr - if w.RemoteAddr().IP.To4() != nil { - if w.publicIPv4 != nil { - newAddr.IP = w.publicIPv4 + if c.RemoteAddr().IP.To4() != nil { + if c.publicIPv4 != nil { + newAddr.IP = c.publicIPv4 } - } else if w.publicIPv6 != nil { - newAddr.IP = w.publicIPv6 + } else if c.publicIPv6 != nil { + newAddr.IP = c.publicIPv6 } return &newAddr } -func (w *WrapConn) RemoteAddr() *net.TCPAddr { - return w.conn.RemoteAddr().(*net.TCPAddr) +func (c *Conn) RemoteAddr() *net.TCPAddr { + return c.conn.RemoteAddr().(*net.TCPAddr) } -func (w *WrapConn) LogDebug(msg string, data ...interface{}) { - w.logger.Debugw(msg, data...) +func (c *Conn) LogDebug(msg string, data ...interface{}) { + c.logger.Debugw(msg, data...) } -func (w *WrapConn) LogInfo(msg string, data ...interface{}) { - w.logger.Infow(msg, data...) +func (c *Conn) LogInfo(msg string, data ...interface{}) { + c.logger.Infow(msg, data...) } -func (w *WrapConn) LogWarn(msg string, data ...interface{}) { - w.logger.Warnw(msg, data...) +func (c *Conn) LogWarn(msg string, data ...interface{}) { + c.logger.Warnw(msg, data...) } -func (w *WrapConn) LogError(msg string, data ...interface{}) { - w.logger.Errorw(msg, data...) +func (c *Conn) LogError(msg string, data ...interface{}) { + c.logger.Errorw(msg, data...) } -func NewConn(connID string, purpose ConnPurpose, conn net.Conn, publicIPv4, publicIPv6 net.IP) WrapStreamReadWriteCloser { +func NewConn(conn net.Conn, connID string, purpose ConnPurpose, publicIPv4, publicIPv6 net.IP) WrapStreamReadWriteCloser { logger := zap.S().With( "connection_id", connID, "local_address", conn.LocalAddr(), "remote_address", conn.RemoteAddr(), + "purpose", purpose, ) - return &WrapConn{ + wrapper := Conn{ logger: logger, - purpose: purpose, connID: connID, conn: conn, publicIPv4: publicIPv4, publicIPv6: publicIPv6, } + wrapper.logger = logger.With("faked_local_addr", wrapper.LocalAddr()) + + return &wrapper } diff --git a/wrappers/ctx.go b/wrappers/ctx.go index ee766f8..bc38211 100644 --- a/wrappers/ctx.go +++ b/wrappers/ctx.go @@ -7,68 +7,68 @@ import ( "github.com/juju/errors" ) -type WrapCtx struct { +type Ctx struct { cancel context.CancelFunc conn WrapStreamReadWriteCloser ctx context.Context } -func (w *WrapCtx) Read(p []byte) (int, error) { +func (c *Ctx) Read(p []byte) (int, error) { select { - case <-w.ctx.Done(): - return 0, errors.Annotate(w.ctx.Err(), "Read is failed because of closed context") + case <-c.ctx.Done(): + return 0, errors.Annotate(c.ctx.Err(), "Read is failed because of closed context") default: - n, err := w.conn.Read(p) + n, err := c.conn.Read(p) if err != nil { - w.cancel() + c.cancel() } return n, err } } -func (w *WrapCtx) Write(p []byte) (int, error) { +func (c *Ctx) Write(p []byte) (int, error) { select { - case <-w.ctx.Done(): - return 0, errors.Annotate(w.ctx.Err(), "Write is failed because of closed context") + case <-c.ctx.Done(): + return 0, errors.Annotate(c.ctx.Err(), "Write is failed because of closed context") default: - n, err := w.conn.Write(p) + n, err := c.conn.Write(p) if err != nil { - w.cancel() + c.cancel() } return n, err } } -func (w *WrapCtx) LogDebug(msg string, data ...interface{}) { - w.conn.LogDebug(msg, data...) +func (c *Ctx) LogDebug(msg string, data ...interface{}) { + c.conn.LogDebug(msg, data...) } -func (w *WrapCtx) LogInfo(msg string, data ...interface{}) { - w.conn.LogInfo(msg, data...) +func (c *Ctx) LogInfo(msg string, data ...interface{}) { + c.conn.LogInfo(msg, data...) } -func (w *WrapCtx) LogWarn(msg string, data ...interface{}) { - w.conn.LogWarn(msg, data...) +func (c *Ctx) LogWarn(msg string, data ...interface{}) { + c.conn.LogWarn(msg, data...) } -func (w *WrapCtx) LogError(msg string, data ...interface{}) { - w.conn.LogError(msg, data...) +func (c *Ctx) LogError(msg string, data ...interface{}) { + c.conn.LogError(msg, data...) } -func (w *WrapCtx) LocalAddr() *net.TCPAddr { - return w.conn.LocalAddr() +func (c *Ctx) LocalAddr() *net.TCPAddr { + return c.conn.LocalAddr() } -func (w *WrapCtx) RemoteAddr() *net.TCPAddr { - return w.conn.RemoteAddr() +func (c *Ctx) RemoteAddr() *net.TCPAddr { + return c.conn.RemoteAddr() } -func (w *WrapCtx) Close() error { - return w.conn.Close() +func (c *Ctx) Close() error { + return c.conn.Close() } func NewCtx(ctx context.Context, cancel context.CancelFunc, conn WrapStreamReadWriteCloser) WrapStreamReadWriteCloser { - return &WrapCtx{ + return &Ctx{ ctx: ctx, cancel: cancel, conn: conn, diff --git a/wrappers/mtproto_cipher.go b/wrappers/mtproto_cipher.go index 146016e..71170dc 100644 --- a/wrappers/mtproto_cipher.go +++ b/wrappers/mtproto_cipher.go @@ -32,7 +32,7 @@ func NewMiddleProxyCipher(conn WrapStreamReadWriteCloser, req *rpc.NonceRequest, enc, _ := makeEncrypterDecrypter(encKey, encIV) _, dec := makeEncrypterDecrypter(decKey, decIV) - return NewWrapBlockCipher(conn, enc, dec) + return NewBlockCipher(conn, enc, dec) } func deriveKeys(purpose CipherPurpose, req *rpc.NonceRequest, resp *rpc.NonceResponse, client *net.TCPAddr, remote *net.TCPAddr, secret []byte) ([]byte, []byte) { diff --git a/wrappers/mtproto_proxy.go b/wrappers/mtproto_proxy.go index 0a6d1a8..c8f1796 100644 --- a/wrappers/mtproto_proxy.go +++ b/wrappers/mtproto_proxy.go @@ -4,8 +4,10 @@ import ( "bytes" "net" - "github.com/9seconds/mtg/mtproto/rpc" "github.com/juju/errors" + + "github.com/9seconds/mtg/mtproto" + "github.com/9seconds/mtg/mtproto/rpc" ) type MTProtoProxy struct { @@ -127,9 +129,14 @@ func (m *MTProtoProxy) Close() error { return m.conn.Close() } -func NewMTProtoProxy(conn WrapPacketReadWriteCloser, req *rpc.ProxyRequest) WrapPacketReadWriteCloser { +func NewMTProtoProxy(conn WrapPacketReadWriteCloser, connOpts *mtproto.ConnectionOpts, adTag []byte) (WrapPacketReadWriteCloser, error) { + req, err := rpc.NewProxyRequest(connOpts.ClientAddr, conn.LocalAddr(), connOpts, adTag) + if err != nil { + return nil, errors.Annotate(err, "Cannot create new RPC proxy request") + } + return &MTProtoProxy{ conn: conn, req: req, - } + }, nil } diff --git a/wrappers/streamcipher.go b/wrappers/streamcipher.go index 1ad764f..1057d33 100644 --- a/wrappers/streamcipher.go +++ b/wrappers/streamcipher.go @@ -7,59 +7,59 @@ import ( "github.com/juju/errors" ) -type WrapStreamCipher struct { +type StreamCipher struct { encryptor cipher.Stream decryptor cipher.Stream conn WrapStreamReadWriteCloser } -func (w *WrapStreamCipher) Read(p []byte) (int, error) { - n, err := w.conn.Read(p) +func (s *StreamCipher) Read(p []byte) (int, error) { + n, err := s.conn.Read(p) if err != nil { return 0, errors.Annotate(err, "Cannot read stream ciphered data") } - w.decryptor.XORKeyStream(p, p[:n]) + s.decryptor.XORKeyStream(p, p[:n]) return n, nil } -func (w *WrapStreamCipher) Write(p []byte) (int, error) { +func (s *StreamCipher) Write(p []byte) (int, error) { encrypted := make([]byte, len(p)) - w.encryptor.XORKeyStream(encrypted, p) + s.encryptor.XORKeyStream(encrypted, p) - return w.conn.Write(encrypted) + return s.conn.Write(encrypted) } -func (w *WrapStreamCipher) LogDebug(msg string, data ...interface{}) { - w.conn.LogDebug(msg, data...) +func (s *StreamCipher) LogDebug(msg string, data ...interface{}) { + s.conn.LogDebug(msg, data...) } -func (w *WrapStreamCipher) LogInfo(msg string, data ...interface{}) { - w.conn.LogInfo(msg, data...) +func (s *StreamCipher) LogInfo(msg string, data ...interface{}) { + s.conn.LogInfo(msg, data...) } -func (w *WrapStreamCipher) LogWarn(msg string, data ...interface{}) { - w.conn.LogWarn(msg, data...) +func (s *StreamCipher) LogWarn(msg string, data ...interface{}) { + s.conn.LogWarn(msg, data...) } -func (w *WrapStreamCipher) LogError(msg string, data ...interface{}) { - w.conn.LogError(msg, data...) +func (s *StreamCipher) LogError(msg string, data ...interface{}) { + s.conn.LogError(msg, data...) } -func (w *WrapStreamCipher) LocalAddr() *net.TCPAddr { - return w.conn.LocalAddr() +func (s *StreamCipher) LocalAddr() *net.TCPAddr { + return s.conn.LocalAddr() } -func (w *WrapStreamCipher) RemoteAddr() *net.TCPAddr { - return w.conn.RemoteAddr() +func (s *StreamCipher) RemoteAddr() *net.TCPAddr { + return s.conn.RemoteAddr() } -func (w *WrapStreamCipher) Close() error { - return w.conn.Close() +func (s *StreamCipher) Close() error { + return s.conn.Close() } func NewStreamCipher(conn WrapStreamReadWriteCloser, encryptor, decryptor cipher.Stream) WrapStreamReadWriteCloser { - return &WrapStreamCipher{ + return &StreamCipher{ conn: conn, encryptor: encryptor, decryptor: decryptor, diff --git a/wrappers/wrap.go b/wrappers/wrap.go index d4bc7d4..3d54bc1 100644 --- a/wrappers/wrap.go +++ b/wrappers/wrap.go @@ -56,6 +56,11 @@ type WrapPacketReader interface { Wrap } +type WrapPacketWriter interface { + io.Writer + Wrap +} + type WrapPacketReadWriter interface { io.Writer WrapPacketReader