mirror of
https://github.com/ScuroNeko/mtg.git
synced 2026-09-01 02:44:02 +03:00
Change rwc to rwc with addresses
This commit is contained in:
+2
-2
@@ -1,12 +1,12 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"io"
|
||||
"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, *config.Config) (*mtproto.ConnectionOpts, io.ReadWriteCloser, error)
|
||||
type Init func(net.Conn, *config.Config) (*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr, error)
|
||||
|
||||
+1
-2
@@ -1,7 +1,6 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net"
|
||||
"time"
|
||||
|
||||
@@ -16,7 +15,7 @@ import (
|
||||
const handshakeTimeout = 10 * time.Second
|
||||
|
||||
// DirectInit initializes client to access Telegram bypassing middleproxies.
|
||||
func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, io.ReadWriteCloser, error) {
|
||||
func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr, error) {
|
||||
if err := config.SetSocketOptions(conn); err != nil {
|
||||
return nil, nil, errors.Annotate(err, "Cannot set socket options")
|
||||
}
|
||||
|
||||
@@ -7,7 +7,6 @@ import (
|
||||
"crypto/md5"
|
||||
"crypto/sha1"
|
||||
"encoding/binary"
|
||||
"io"
|
||||
"net"
|
||||
|
||||
"github.com/9seconds/mtg/mtproto/rpc"
|
||||
@@ -23,14 +22,17 @@ const (
|
||||
|
||||
var emptyIP = [4]byte{0x00, 0x00, 0x00, 0x00}
|
||||
|
||||
func NewMiddleProxyCipherRWC(conn io.ReadWriteCloser, req *rpc.RPCNonceRequest, resp *rpc.RPCNonceResponse, client *net.TCPAddr, remote *net.TCPAddr, secret []byte) io.ReadWriteCloser {
|
||||
func NewMiddleProxyCipherRWC(conn wrappers.ReadWriteCloserWithAddr, req *rpc.RPCNonceRequest,
|
||||
resp *rpc.RPCNonceResponse, client *net.TCPAddr, remote *net.TCPAddr,
|
||||
secret []byte) wrappers.ReadWriteCloserWithAddr {
|
||||
encryptor := newCBCCipher(CipherPurposeClient, req, resp, client, remote, secret)
|
||||
decryptor := newCBCCipher(CipherPurposeServer, req, resp, client, remote, secret)
|
||||
|
||||
return wrappers.NewBlockCipherRWC(conn, encryptor, decryptor)
|
||||
}
|
||||
|
||||
func newCBCCipher(purpose CipherPurpose, req *rpc.RPCNonceRequest, resp *rpc.RPCNonceResponse, client *net.TCPAddr, remote *net.TCPAddr, secret []byte) cipher.BlockMode {
|
||||
func newCBCCipher(purpose CipherPurpose, req *rpc.RPCNonceRequest, resp *rpc.RPCNonceResponse,
|
||||
client *net.TCPAddr, remote *net.TCPAddr, secret []byte) cipher.BlockMode {
|
||||
message := bytes.Buffer{}
|
||||
message.Write(resp.Nonce[:])
|
||||
message.Write(req.Nonce[:])
|
||||
|
||||
@@ -6,8 +6,11 @@ import (
|
||||
"encoding/binary"
|
||||
"hash/crc32"
|
||||
"io"
|
||||
"net"
|
||||
|
||||
"github.com/juju/errors"
|
||||
|
||||
"github.com/9seconds/mtg/wrappers"
|
||||
)
|
||||
|
||||
// Frame: { MessageLength(4) | SequenceNumber(4) | Message(???) | CRC32(4) [| padding(4), ...] }
|
||||
@@ -19,7 +22,7 @@ const (
|
||||
var frameRWCPadding = [4]byte{0x04, 0x00, 0x00, 0x00}
|
||||
|
||||
type FrameRWC struct {
|
||||
conn io.ReadWriteCloser
|
||||
conn wrappers.ReadWriteCloserWithAddr
|
||||
|
||||
readSeqNo int32
|
||||
writeSeqNo int32
|
||||
@@ -102,6 +105,10 @@ func (f *FrameRWC) Close() error {
|
||||
return f.conn.Close()
|
||||
}
|
||||
|
||||
func (f *FrameRWC) Addr() *net.TCPAddr {
|
||||
return f.conn.Addr()
|
||||
}
|
||||
|
||||
func (f *FrameRWC) flush(p []byte) (int, error) {
|
||||
sizeToRead := len(p)
|
||||
if f.readBuf.Len() < sizeToRead {
|
||||
@@ -119,7 +126,7 @@ func (f *FrameRWC) flush(p []byte) (int, error) {
|
||||
return sizeToRead, nil
|
||||
}
|
||||
|
||||
func NewFrameRWC(conn io.ReadWriteCloser, seqNo int32) io.ReadWriteCloser {
|
||||
func NewFrameRWC(conn wrappers.ReadWriteCloserWithAddr, seqNo int32) wrappers.ReadWriteCloserWithAddr {
|
||||
return &FrameRWC{
|
||||
conn: conn,
|
||||
readSeqNo: seqNo,
|
||||
|
||||
+1
-2
@@ -1,7 +1,6 @@
|
||||
package telegram
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net"
|
||||
"time"
|
||||
|
||||
@@ -29,7 +28,7 @@ func (t *tgDialer) dial(addr string) (net.Conn, error) {
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func (t *tgDialer) dialRWC(addr string) (io.ReadWriteCloser, error) {
|
||||
func (t *tgDialer) dialRWC(addr string) (wrappers.ReadWriteCloserWithAddr, error) {
|
||||
conn, err := t.dial(addr)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
+2
-3
@@ -1,7 +1,6 @@
|
||||
package telegram
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net"
|
||||
|
||||
"github.com/juju/errors"
|
||||
@@ -33,7 +32,7 @@ type directTelegram struct {
|
||||
baseTelegram
|
||||
}
|
||||
|
||||
func (t *directTelegram) Dial(connOpts *mtproto.ConnectionOpts) (io.ReadWriteCloser, error) {
|
||||
func (t *directTelegram) Dial(connOpts *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error) {
|
||||
dc := connOpts.DC
|
||||
if dc < 0 {
|
||||
dc = -dc
|
||||
@@ -44,7 +43,7 @@ func (t *directTelegram) Dial(connOpts *mtproto.ConnectionOpts) (io.ReadWriteClo
|
||||
return t.baseTelegram.dial(dc-1, connOpts.ConnectionProto)
|
||||
}
|
||||
|
||||
func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts, conn io.ReadWriteCloser) (io.ReadWriteCloser, error) {
|
||||
func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.ReadWriteCloserWithAddr) (wrappers.ReadWriteCloserWithAddr, error) {
|
||||
obfs2, frame := obfuscated2.MakeTelegramObfuscated2Frame(connOpts)
|
||||
|
||||
if n, err := conn.Write(frame); err != nil || n != obfuscated2.FrameLen {
|
||||
|
||||
@@ -2,7 +2,6 @@ package telegram
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"net"
|
||||
"net/http"
|
||||
@@ -16,6 +15,7 @@ import (
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/9seconds/mtg/mtproto"
|
||||
"github.com/9seconds/mtg/wrappers"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -39,7 +39,7 @@ type middleTelegramCaller struct {
|
||||
httpClient *http.Client
|
||||
}
|
||||
|
||||
func (t *middleTelegramCaller) Dial(connOpts *mtproto.ConnectionOpts) (io.ReadWriteCloser, error) {
|
||||
func (t *middleTelegramCaller) Dial(connOpts *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error) {
|
||||
dc := connOpts.DC
|
||||
if dc == 0 {
|
||||
dc = 1
|
||||
|
||||
@@ -1,20 +1,20 @@
|
||||
package telegram
|
||||
|
||||
import (
|
||||
"io"
|
||||
"math/rand"
|
||||
|
||||
"github.com/juju/errors"
|
||||
|
||||
"github.com/9seconds/mtg/mtproto"
|
||||
"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(*mtproto.ConnectionOpts) (io.ReadWriteCloser, error)
|
||||
Init(*mtproto.ConnectionOpts, io.ReadWriteCloser) (io.ReadWriteCloser, error)
|
||||
Dial(*mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error)
|
||||
Init(*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr) (wrappers.ReadWriteCloserWithAddr, error)
|
||||
}
|
||||
|
||||
type baseTelegram struct {
|
||||
@@ -24,7 +24,7 @@ type baseTelegram struct {
|
||||
v6Addresses map[int16][]string
|
||||
}
|
||||
|
||||
func (b *baseTelegram) dial(dcIdx int16, proto mtproto.ConnectionProtocol) (io.ReadWriteCloser, error) {
|
||||
func (b *baseTelegram) dial(dcIdx int16, proto mtproto.ConnectionProtocol) (wrappers.ReadWriteCloserWithAddr, error) {
|
||||
addrs := make([]string, 2)
|
||||
|
||||
if proto&mtproto.ConnectionProtocolIPv6 != 0 {
|
||||
|
||||
@@ -4,20 +4,20 @@ import (
|
||||
"bytes"
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"io"
|
||||
"net"
|
||||
|
||||
"github.com/juju/errors"
|
||||
)
|
||||
|
||||
type BlockCipherReadWriteCloser struct {
|
||||
type BlockCipherReadWriteCloserWithAddr struct {
|
||||
buf *bytes.Buffer
|
||||
|
||||
conn io.ReadWriteCloser
|
||||
conn ReadWriteCloserWithAddr
|
||||
encryptor cipher.BlockMode
|
||||
decryptor cipher.BlockMode
|
||||
}
|
||||
|
||||
func (c *BlockCipherReadWriteCloser) Read(p []byte) (int, error) {
|
||||
func (c *BlockCipherReadWriteCloserWithAddr) Read(p []byte) (int, error) {
|
||||
if c.buf.Len() > 0 {
|
||||
return c.flush(p)
|
||||
}
|
||||
@@ -34,7 +34,7 @@ func (c *BlockCipherReadWriteCloser) Read(p []byte) (int, error) {
|
||||
return c.flush(p)
|
||||
}
|
||||
|
||||
func (c *BlockCipherReadWriteCloser) Write(p []byte) (int, error) {
|
||||
func (c *BlockCipherReadWriteCloserWithAddr) Write(p []byte) (int, error) {
|
||||
if len(p)%aes.BlockSize > 0 {
|
||||
return 0, errors.Errorf("Incorrect block size %d", len(p))
|
||||
}
|
||||
@@ -50,12 +50,16 @@ func (c *BlockCipherReadWriteCloser) Write(p []byte) (int, error) {
|
||||
return c.conn.Write(encrypted)
|
||||
}
|
||||
|
||||
func (c *BlockCipherReadWriteCloser) Close() error {
|
||||
func (c *BlockCipherReadWriteCloserWithAddr) Close() error {
|
||||
defer putBuffer(c.buf)
|
||||
return c.conn.Close()
|
||||
}
|
||||
|
||||
func (c *BlockCipherReadWriteCloser) flush(p []byte) (int, error) {
|
||||
func (c *BlockCipherReadWriteCloserWithAddr) Addr() *net.TCPAddr {
|
||||
return c.conn.Addr()
|
||||
}
|
||||
|
||||
func (c *BlockCipherReadWriteCloserWithAddr) flush(p []byte) (int, error) {
|
||||
sizeToRead := len(p)
|
||||
if c.buf.Len() < sizeToRead {
|
||||
sizeToRead = c.buf.Len()
|
||||
@@ -76,8 +80,8 @@ func (c *BlockCipherReadWriteCloser) flush(p []byte) (int, error) {
|
||||
return sizeToRead, nil
|
||||
}
|
||||
|
||||
func NewBlockCipherRWC(conn io.ReadWriteCloser, encryptor, decryptor cipher.BlockMode) io.ReadWriteCloser {
|
||||
return &BlockCipherReadWriteCloser{
|
||||
func NewBlockCipherRWC(conn ReadWriteCloserWithAddr, encryptor, decryptor cipher.BlockMode) ReadWriteCloserWithAddr {
|
||||
return &BlockCipherReadWriteCloserWithAddr{
|
||||
buf: getBuffer(),
|
||||
conn: conn,
|
||||
encryptor: encryptor,
|
||||
|
||||
+12
-8
@@ -2,21 +2,21 @@ package wrappers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net"
|
||||
|
||||
"github.com/juju/errors"
|
||||
)
|
||||
|
||||
// CtxReadWriteCloser wraps underlying connection and does management of the
|
||||
// context and its cancel function.
|
||||
type CtxReadWriteCloser struct {
|
||||
type CtxReadWriteCloserWithAddr struct {
|
||||
ctx context.Context
|
||||
conn io.ReadWriteCloser
|
||||
conn ReadWriteCloserWithAddr
|
||||
cancel context.CancelFunc
|
||||
}
|
||||
|
||||
// Read reads from connection
|
||||
func (c *CtxReadWriteCloser) Read(p []byte) (int, error) {
|
||||
func (c *CtxReadWriteCloserWithAddr) Read(p []byte) (int, error) {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
return 0, errors.Annotate(c.ctx.Err(), "Read is failed because of closed context")
|
||||
@@ -30,7 +30,7 @@ func (c *CtxReadWriteCloser) Read(p []byte) (int, error) {
|
||||
}
|
||||
|
||||
// Write writes into connection.
|
||||
func (c *CtxReadWriteCloser) Write(p []byte) (int, error) {
|
||||
func (c *CtxReadWriteCloserWithAddr) Write(p []byte) (int, error) {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
return 0, errors.Annotate(c.ctx.Err(), "Write is failed because of closed context")
|
||||
@@ -44,14 +44,18 @@ func (c *CtxReadWriteCloser) Write(p []byte) (int, error) {
|
||||
}
|
||||
|
||||
// Close closes underlying connection.
|
||||
func (c *CtxReadWriteCloser) Close() error {
|
||||
func (c *CtxReadWriteCloserWithAddr) Close() error {
|
||||
return c.conn.Close()
|
||||
}
|
||||
|
||||
func (c *CtxReadWriteCloserWithAddr) Addr() *net.TCPAddr {
|
||||
return c.conn.Addr()
|
||||
}
|
||||
|
||||
// NewCtxRWC returns ReadWriteCloser which respects given context,
|
||||
// cancellation etc.
|
||||
func NewCtxRWC(ctx context.Context, cancel context.CancelFunc, conn io.ReadWriteCloser) io.ReadWriteCloser {
|
||||
return &CtxReadWriteCloser{
|
||||
func NewCtxRWC(ctx context.Context, cancel context.CancelFunc, conn ReadWriteCloserWithAddr) ReadWriteCloserWithAddr {
|
||||
return &CtxReadWriteCloserWithAddr{
|
||||
conn: conn,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
|
||||
+12
-8
@@ -1,44 +1,48 @@
|
||||
package wrappers
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// LogReadWriteCloser adds additional logging for reading/writing. All
|
||||
// logging is performed for debug mode only.
|
||||
type LogReadWriteCloser struct {
|
||||
conn io.ReadWriteCloser
|
||||
type LogReadWriteCloserWithAddr struct {
|
||||
conn ReadWriteCloserWithAddr
|
||||
logger *zap.SugaredLogger
|
||||
sockid string
|
||||
name string
|
||||
}
|
||||
|
||||
// Read reads from connection
|
||||
func (l *LogReadWriteCloser) Read(p []byte) (n int, err error) {
|
||||
func (l *LogReadWriteCloserWithAddr) Read(p []byte) (n int, err error) {
|
||||
n, err = l.conn.Read(p)
|
||||
l.logger.Debugw("Finish reading", "name", l.name, "socketid", l.sockid, "nbytes", n, "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Write writes into connection.
|
||||
func (l *LogReadWriteCloser) Write(p []byte) (n int, err error) {
|
||||
func (l *LogReadWriteCloserWithAddr) Write(p []byte) (n int, err error) {
|
||||
n, err = l.conn.Write(p)
|
||||
l.logger.Debugw("Finish writing", "name", l.name, "socketid", l.sockid, "nbytes", n, "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Close closes underlying connection.
|
||||
func (l *LogReadWriteCloser) Close() error {
|
||||
func (l *LogReadWriteCloserWithAddr) Close() error {
|
||||
err := l.conn.Close()
|
||||
l.logger.Debugw("Finish closing socket", "name", l.name, "socketid", l.sockid, "error", err)
|
||||
return err
|
||||
}
|
||||
|
||||
func (l *LogReadWriteCloserWithAddr) Addr() *net.TCPAddr {
|
||||
return l.conn.Addr()
|
||||
}
|
||||
|
||||
// NewLogRWC wraps ReadWriteCloser with logger calls.
|
||||
func NewLogRWC(conn io.ReadWriteCloser, logger *zap.SugaredLogger, sockid string, name string) io.ReadWriteCloser {
|
||||
return &LogReadWriteCloser{
|
||||
func NewLogRWC(conn ReadWriteCloserWithAddr, logger *zap.SugaredLogger, sockid string, name string) ReadWriteCloserWithAddr {
|
||||
return &LogReadWriteCloserWithAddr{
|
||||
conn: conn,
|
||||
logger: logger,
|
||||
sockid: sockid,
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
package wrappers
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net"
|
||||
)
|
||||
|
||||
type ReadWriteCloserWithAddr interface {
|
||||
io.ReadWriteCloser
|
||||
|
||||
Addr() *net.TCPAddr
|
||||
}
|
||||
@@ -2,26 +2,26 @@ package wrappers
|
||||
|
||||
import (
|
||||
"crypto/cipher"
|
||||
"io"
|
||||
"net"
|
||||
)
|
||||
|
||||
// StreamCipherReadWriteCloser is a ReadWriteCloser which ciphers
|
||||
// incoming and outgoing data with givem cipher.Stream instances.
|
||||
type StreamCipherReadWriteCloser struct {
|
||||
type StreamCipherReadWriteCloserWithAddr struct {
|
||||
encryptor cipher.Stream
|
||||
decryptor cipher.Stream
|
||||
conn io.ReadWriteCloser
|
||||
conn ReadWriteCloserWithAddr
|
||||
}
|
||||
|
||||
// Read reads from connection
|
||||
func (c *StreamCipherReadWriteCloser) Read(p []byte) (n int, err error) {
|
||||
func (c *StreamCipherReadWriteCloserWithAddr) Read(p []byte) (n int, err error) {
|
||||
n, err = c.conn.Read(p)
|
||||
c.decryptor.XORKeyStream(p, p[:n])
|
||||
return
|
||||
}
|
||||
|
||||
// Write writes into connection.
|
||||
func (c *StreamCipherReadWriteCloser) Write(p []byte) (int, error) {
|
||||
func (c *StreamCipherReadWriteCloserWithAddr) Write(p []byte) (int, error) {
|
||||
// This is to decrease an amount of allocations. Unfortunately, escape
|
||||
// analysis in (at least Golang 1.10) is absolutely not perfect. For
|
||||
// example, it understands that we want to have a slice locally, right?
|
||||
@@ -39,14 +39,18 @@ func (c *StreamCipherReadWriteCloser) Write(p []byte) (int, error) {
|
||||
}
|
||||
|
||||
// Close closes underlying connection.
|
||||
func (c *StreamCipherReadWriteCloser) Close() error {
|
||||
func (c *StreamCipherReadWriteCloserWithAddr) Close() error {
|
||||
return c.conn.Close()
|
||||
}
|
||||
|
||||
func (c *StreamCipherReadWriteCloserWithAddr) Addr() *net.TCPAddr {
|
||||
return c.conn.Addr()
|
||||
}
|
||||
|
||||
// NewStreamCipherRWC returns wrapper which transparently
|
||||
// encrypts/decrypts traffic with obfuscated2 protocol.
|
||||
func NewStreamCipherRWC(conn io.ReadWriteCloser, encryptor, decryptor cipher.Stream) io.ReadWriteCloser {
|
||||
return &StreamCipherReadWriteCloser{
|
||||
func NewStreamCipherRWC(conn ReadWriteCloserWithAddr, encryptor, decryptor cipher.Stream) ReadWriteCloserWithAddr {
|
||||
return &StreamCipherReadWriteCloserWithAddr{
|
||||
conn: conn,
|
||||
encryptor: encryptor,
|
||||
decryptor: decryptor,
|
||||
|
||||
+10
-7
@@ -1,31 +1,34 @@
|
||||
package wrappers
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net"
|
||||
"time"
|
||||
|
||||
"github.com/9seconds/mtg/config"
|
||||
)
|
||||
|
||||
type TimeoutReadWriteCloser struct {
|
||||
type TimeoutReadWriteCloserWithAddr struct {
|
||||
conn net.Conn
|
||||
}
|
||||
|
||||
func (t *TimeoutReadWriteCloser) Read(p []byte) (int, error) {
|
||||
func (t *TimeoutReadWriteCloserWithAddr) Read(p []byte) (int, error) {
|
||||
t.conn.SetReadDeadline(time.Now().Add(config.TimeoutRead))
|
||||
return t.conn.Read(p)
|
||||
}
|
||||
|
||||
func (t *TimeoutReadWriteCloser) Write(p []byte) (int, error) {
|
||||
func (t *TimeoutReadWriteCloserWithAddr) Write(p []byte) (int, error) {
|
||||
t.conn.SetWriteDeadline(time.Now().Add(config.TimeoutWrite))
|
||||
return t.conn.Write(p)
|
||||
}
|
||||
|
||||
func (t *TimeoutReadWriteCloser) Close() error {
|
||||
func (t *TimeoutReadWriteCloserWithAddr) Close() error {
|
||||
return t.conn.Close()
|
||||
}
|
||||
|
||||
func NewTimeoutRWC(conn net.Conn) io.ReadWriteCloser {
|
||||
return &TimeoutReadWriteCloser{conn}
|
||||
func (t *TimeoutReadWriteCloserWithAddr) Addr() *net.TCPAddr {
|
||||
return t.conn.RemoteAddr().(*net.TCPAddr)
|
||||
}
|
||||
|
||||
func NewTimeoutRWC(conn net.Conn) ReadWriteCloserWithAddr {
|
||||
return &TimeoutReadWriteCloserWithAddr{conn}
|
||||
}
|
||||
|
||||
+12
-8
@@ -1,37 +1,41 @@
|
||||
package wrappers
|
||||
|
||||
import "io"
|
||||
import "net"
|
||||
|
||||
// TrafficReadWriteCloser counts an amount of ingress/egress traffic by
|
||||
// calling given callbacks.
|
||||
type TrafficReadWriteCloser struct {
|
||||
conn io.ReadWriteCloser
|
||||
type TrafficReadWriteCloserWithAddr struct {
|
||||
conn ReadWriteCloserWithAddr
|
||||
readCallback func(int)
|
||||
writeCallback func(int)
|
||||
}
|
||||
|
||||
// Read reads from connection
|
||||
func (t *TrafficReadWriteCloser) Read(p []byte) (n int, err error) {
|
||||
func (t *TrafficReadWriteCloserWithAddr) Read(p []byte) (n int, err error) {
|
||||
n, err = t.conn.Read(p)
|
||||
t.readCallback(n)
|
||||
return
|
||||
}
|
||||
|
||||
// Write writes into connection.
|
||||
func (t *TrafficReadWriteCloser) Write(p []byte) (n int, err error) {
|
||||
func (t *TrafficReadWriteCloserWithAddr) Write(p []byte) (n int, err error) {
|
||||
n, err = t.conn.Write(p)
|
||||
t.writeCallback(n)
|
||||
return
|
||||
}
|
||||
|
||||
// Close closes underlying connection.
|
||||
func (t *TrafficReadWriteCloser) Close() error {
|
||||
func (t *TrafficReadWriteCloserWithAddr) Close() error {
|
||||
return t.conn.Close()
|
||||
}
|
||||
|
||||
func (t *TrafficReadWriteCloserWithAddr) Addr() *net.TCPAddr {
|
||||
return t.conn.Addr()
|
||||
}
|
||||
|
||||
// NewTrafficRWC wraps ReadWriteCloser to have read/write callbacks.
|
||||
func NewTrafficRWC(conn io.ReadWriteCloser, readCallback, writeCallback func(int)) io.ReadWriteCloser {
|
||||
return &TrafficReadWriteCloser{
|
||||
func NewTrafficRWC(conn ReadWriteCloserWithAddr, readCallback, writeCallback func(int)) ReadWriteCloserWithAddr {
|
||||
return &TrafficReadWriteCloserWithAddr{
|
||||
conn: conn,
|
||||
readCallback: readCallback,
|
||||
writeCallback: writeCallback,
|
||||
|
||||
Reference in New Issue
Block a user