From a3933e6ede78eb40ea347b699b7a7637d7bca18d Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Fri, 6 Jul 2018 17:29:24 +0300
Subject: [PATCH] Preliminary debug state

---
 client/client.go                  |  2 +-
 client/direct.go                  |  4 ++--
 client/middle.go                  |  4 ++--
 mtproto/wrappers/abridged.go      | 16 ++++++++++++----
 mtproto/wrappers/frame.go         |  4 ++++
 mtproto/wrappers/intermediate.go  |  4 ++++
 mtproto/wrappers/proxy_request.go | 22 +++++++++++-----------
 proxy/server.go                   | 26 ++++++++++----------------
 telegram/dialer.go                |  4 ++--
 telegram/direct.go                |  4 ++--
 telegram/middle_caller.go         |  4 ++--
 telegram/telegram.go              |  6 +++---
 utils/read_current_data.go        | 20 ++++++++++++++++++++
 wrappers/blockcipherrwc.go        | 23 ++++++++++++++++-------
 wrappers/buffered_reader.go       | 12 ++++++++++--
 wrappers/ctxrwc.go                |  4 ++++
 wrappers/logrwc.go                |  8 ++++++--
 wrappers/rwcaddr.go               |  1 +
 wrappers/streamcipherrwc.go       |  4 ++++
 wrappers/timeoutrwc.go            |  8 +++++++-
 wrappers/trafficrwc.go            |  4 ++++
 21 files changed, 127 insertions(+), 57 deletions(-)
 create mode 100644 utils/read_current_data.go

diff --git a/client/client.go b/client/client.go
index 32e23a9..318c476 100644
--- a/client/client.go
+++ b/client/client.go
@@ -9,4 +9,4 @@ import (
 )
 
 // Init has to initialize client connection based on given config.
-type Init func(net.Conn, *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error)
+type Init func(net.Conn, string, *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error)
diff --git a/client/direct.go b/client/direct.go
index 3ce1eed..7ee7e75 100644
--- a/client/direct.go
+++ b/client/direct.go
@@ -15,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) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) {
+func DirectInit(conn net.Conn, socketID string, conf *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) {
 	if err := config.SetSocketOptions(conn); err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot set socket options")
 	}
@@ -34,7 +34,7 @@ func DirectInit(conn net.Conn, conf *config.Config) (wrappers.ReadWriteCloserWit
 	connOpts.ConnectionProto = mtproto.ConnectionProtocolAny
 	connOpts.ClientAddr = conn.RemoteAddr().(*net.TCPAddr)
 
-	socket := wrappers.NewTimeoutRWC(conn, conf.PublicIPv4, conf.PublicIPv6)
+	socket := wrappers.NewTimeoutRWC(conn, socketID, conf.PublicIPv4, conf.PublicIPv6)
 	socket = wrappers.NewStreamCipherRWC(socket, obfs2.Encryptor, obfs2.Decryptor)
 
 	return socket, connOpts, nil
diff --git a/client/middle.go b/client/middle.go
index f8877fc..2c978ec 100644
--- a/client/middle.go
+++ b/client/middle.go
@@ -9,8 +9,8 @@ import (
 	"github.com/9seconds/mtg/wrappers"
 )
 
-func MiddleInit(conn net.Conn, conf *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) {
-	newConn, opts, err := DirectInit(conn, conf)
+func MiddleInit(conn net.Conn, socketID string, conf *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) {
+	newConn, opts, err := DirectInit(conn, socketID, conf)
 	if err != nil {
 		return nil, nil, err
 	}
diff --git a/mtproto/wrappers/abridged.go b/mtproto/wrappers/abridged.go
index 1850220..e4bfc13 100644
--- a/mtproto/wrappers/abridged.go
+++ b/mtproto/wrappers/abridged.go
@@ -31,15 +31,18 @@ func (a *AbridgedReadWriteCloserWithAddr) Read(p []byte) (int, error) {
 		buf := &bytes.Buffer{}
 		buf.Grow(3)
 
+		q := make([]byte, 1)
+
 		if _, err := io.CopyN(buf, a.conn, 1); err != nil {
 			return errors.Annotate(err, "Cannot read message length")
 		}
 		msgLength := uint8(buf.Bytes()[0])
+		q[0] = msgLength
 		buf.Reset()
 
 		if msgLength >= abridgedQuickAckLength {
 			a.opts.ReadHacks.QuickAck = true
-			msgLength -= 0x80
+			msgLength -= abridgedQuickAckLength
 		}
 
 		msgLength32 := uint32(msgLength)
@@ -49,6 +52,7 @@ func (a *AbridgedReadWriteCloserWithAddr) Read(p []byte) (int, error) {
 			}
 			number := uint24{}
 			copy(number[:], buf.Bytes())
+			q = append(q, buf.Bytes()...)
 			msgLength32 = fromUint24(number)
 		}
 		msgLength32 *= 4
@@ -59,6 +63,7 @@ func (a *AbridgedReadWriteCloserWithAddr) Read(p []byte) (int, error) {
 		if _, err := io.CopyN(buf, a.conn, int64(msgLength32)); err != nil {
 			return errors.Annotate(err, "Cannot read message")
 		}
+		q = append(q, buf.Bytes()...)
 		a.Buffer.Write(buf.Bytes())
 
 		return nil
@@ -87,10 +92,9 @@ func (a *AbridgedReadWriteCloserWithAddr) Write(p []byte) (int, error) {
 		buf.Write(length24[:])
 		buf.Write(p)
 		return a.conn.Write(buf.Bytes())
-
-	default:
-		return 0, errors.Errorf("Packet is too big %d", len(p))
 	}
+
+	return 0, errors.Errorf("Packet is too big %d", len(p))
 }
 
 func (a *AbridgedReadWriteCloserWithAddr) Close() error {
@@ -105,6 +109,10 @@ func (a *AbridgedReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
 	return a.conn.RemoteAddr()
 }
 
+func (a *AbridgedReadWriteCloserWithAddr) SocketID() string {
+	return a.conn.SocketID()
+}
+
 func toUint24(number uint32) uint24 {
 	return uint24{byte(number), byte(number >> 8), byte(number >> 16)}
 }
diff --git a/mtproto/wrappers/frame.go b/mtproto/wrappers/frame.go
index 2c5cf8a..633c4ad 100644
--- a/mtproto/wrappers/frame.go
+++ b/mtproto/wrappers/frame.go
@@ -120,6 +120,10 @@ func (f *FrameRWC) RemoteAddr() *net.TCPAddr {
 	return f.conn.RemoteAddr()
 }
 
+func (f *FrameRWC) SocketID() string {
+	return f.conn.SocketID()
+}
+
 func NewFrameRWC(conn wrappers.ReadWriteCloserWithAddr, seqNo int32) wrappers.ReadWriteCloserWithAddr {
 	return &FrameRWC{
 		BufferedReader: wrappers.NewBufferedReader(),
diff --git a/mtproto/wrappers/intermediate.go b/mtproto/wrappers/intermediate.go
index a22e56d..0468d47 100644
--- a/mtproto/wrappers/intermediate.go
+++ b/mtproto/wrappers/intermediate.go
@@ -77,6 +77,10 @@ func (i *IntermediateReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
 	return i.conn.RemoteAddr()
 }
 
+func (i *IntermediateReadWriteCloserWithAddr) SocketID() string {
+	return i.conn.SocketID()
+}
+
 func NewIntermediateRWC(conn wrappers.ReadWriteCloserWithAddr, connOpts *mtproto.ConnectionOpts) wrappers.ReadWriteCloserWithAddr {
 	return &IntermediateReadWriteCloserWithAddr{
 		BufferedReader: wrappers.NewBufferedReader(),
diff --git a/mtproto/wrappers/proxy_request.go b/mtproto/wrappers/proxy_request.go
index 2641715..088565d 100644
--- a/mtproto/wrappers/proxy_request.go
+++ b/mtproto/wrappers/proxy_request.go
@@ -10,6 +10,7 @@ import (
 
 	"github.com/9seconds/mtg/mtproto"
 	"github.com/9seconds/mtg/mtproto/rpc"
+	"github.com/9seconds/mtg/utils"
 	"github.com/9seconds/mtg/wrappers"
 )
 
@@ -29,7 +30,7 @@ func (p *ProxyRequestReadWriteCloserWithAddr) Read(buf []byte) (int, error) {
 
 		switch {
 		case bytes.Equal(ans, rpc.TagProxyAns):
-			return p.readProxyAns(buf)
+			return p.readProxyAns()
 		case bytes.Equal(ans, rpc.TagSimpleAck):
 			return p.readSimpleAck()
 		case bytes.Equal(ans, rpc.TagCloseExt):
@@ -44,21 +45,16 @@ func (p *ProxyRequestReadWriteCloserWithAddr) readCloseExt() error {
 	return errors.New("Connection has been closed remotely")
 }
 
-func (p *ProxyRequestReadWriteCloserWithAddr) readProxyAns(buf []byte) (err error) {
+func (p *ProxyRequestReadWriteCloserWithAddr) readProxyAns() (err error) {
 	if _, err = io.CopyN(ioutil.Discard, p.conn, 8+4); err != nil {
 		return errors.Annotate(err, "Cannot skip flags and connid")
 	}
 
-	n := len(buf)
-	preBuffer := &bytes.Buffer{}
-	for n == len(buf) {
-		n, err = p.conn.Read(buf)
-		if err != nil {
-			return errors.Annotate(err, "Cannot read proxy answer")
-		}
-		preBuffer.Write(buf[:n])
+	buf, err := utils.ReadCurrentData(p.conn)
+	if err != nil {
+		return errors.Annotate(err, "Cannot read proxy answer")
 	}
-	p.Buffer.Write(preBuffer.Bytes())
+	p.Buffer.Write(buf)
 
 	return nil
 }
@@ -97,6 +93,10 @@ func (p *ProxyRequestReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
 	return p.conn.RemoteAddr()
 }
 
+func (p *ProxyRequestReadWriteCloserWithAddr) SocketID() string {
+	return p.conn.SocketID()
+}
+
 func NewProxyRequestRWC(conn wrappers.ReadWriteCloserWithAddr, connOpts *mtproto.ConnectionOpts, adTag []byte) (wrappers.ReadWriteCloserWithAddr, error) {
 	req, err := rpc.NewProxyRequest(connOpts.ClientAddr, conn.LocalAddr(), connOpts, adTag)
 	if err != nil {
diff --git a/proxy/server.go b/proxy/server.go
index 9b43613..18bb98e 100644
--- a/proxy/server.go
+++ b/proxy/server.go
@@ -14,6 +14,7 @@ import (
 	"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"
 )
 
@@ -122,7 +123,7 @@ func (s *Server) accept(conn net.Conn) {
 }
 
 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, s.conf)
+	socket, connOpts, err := s.clientInit(conn, socketID, s.conf)
 	if err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot init client connection")
 	}
@@ -135,7 +136,7 @@ func (s *Server) getClientStream(ctx context.Context, cancel context.CancelFunc,
 }
 
 func (s *Server) getTelegramStream(ctx context.Context, cancel context.CancelFunc, connOpts *mtproto.ConnectionOpts, socketID string) (io.ReadWriteCloser, error) {
-	conn, err := s.tg.Dial(connOpts)
+	conn, err := s.tg.Dial(socketID, connOpts)
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot connect to Telegram")
 	}
@@ -152,22 +153,15 @@ func (s *Server) getTelegramStream(ctx context.Context, cancel context.CancelFun
 	return conn, nil
 }
 
-func (s *Server) pump(src io.Reader, dst io.Writer, socketID, name string) (err error) {
-	copyBuf := make([]byte, 1024*1024*2)
-
-	n := config.BufferSizeCopy
-	for n == config.BufferSizeCopy {
-		n, err = src.Read(copyBuf)
-		if err != nil {
-			break
-		}
-		_, err = dst.Write(copyBuf[:n])
-		if err != nil {
-			break
-		}
+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")
 	}
 
-	return
+	_, err = dst.Write(buf)
+
+	return err
 }
 
 // NewServer creates new instance of MTPROTO proxy.
diff --git a/telegram/dialer.go b/telegram/dialer.go
index 9cf048b..5f3e20c 100644
--- a/telegram/dialer.go
+++ b/telegram/dialer.go
@@ -30,11 +30,11 @@ func (t *tgDialer) dial(addr string) (net.Conn, error) {
 	return conn, nil
 }
 
-func (t *tgDialer) dialRWC(addr string) (wrappers.ReadWriteCloserWithAddr, error) {
+func (t *tgDialer) dialRWC(addr, sock string) (wrappers.ReadWriteCloserWithAddr, error) {
 	conn, err := t.dial(addr)
 	if err != nil {
 		return nil, err
 	}
 
-	return wrappers.NewTimeoutRWC(conn, t.conf.PublicIPv4, t.conf.PublicIPv6), nil
+	return wrappers.NewTimeoutRWC(conn, sock, t.conf.PublicIPv4, t.conf.PublicIPv6), nil
 }
diff --git a/telegram/direct.go b/telegram/direct.go
index ae89276..601fa52 100644
--- a/telegram/direct.go
+++ b/telegram/direct.go
@@ -33,7 +33,7 @@ type directTelegram struct {
 	baseTelegram
 }
 
-func (t *directTelegram) Dial(connOpts *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error) {
+func (t *directTelegram) Dial(sock string, connOpts *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error) {
 	dc := connOpts.DC
 	if dc < 0 {
 		dc = -dc
@@ -41,7 +41,7 @@ func (t *directTelegram) Dial(connOpts *mtproto.ConnectionOpts) (wrappers.ReadWr
 		dc = 1
 	}
 
-	return t.baseTelegram.dial(dc-1, connOpts.ConnectionProto)
+	return t.baseTelegram.dial(dc-1, sock, connOpts.ConnectionProto)
 }
 
 func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.ReadWriteCloserWithAddr) (wrappers.ReadWriteCloserWithAddr, error) {
diff --git a/telegram/middle_caller.go b/telegram/middle_caller.go
index acc1b27..490ed26 100644
--- a/telegram/middle_caller.go
+++ b/telegram/middle_caller.go
@@ -39,7 +39,7 @@ type middleTelegramCaller struct {
 	httpClient  *http.Client
 }
 
-func (t *middleTelegramCaller) Dial(connOpts *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error) {
+func (t *middleTelegramCaller) Dial(sock string, connOpts *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error) {
 	dc := connOpts.DC
 	if dc == 0 {
 		dc = 1
@@ -47,7 +47,7 @@ func (t *middleTelegramCaller) Dial(connOpts *mtproto.ConnectionOpts) (wrappers.
 	t.dialerMutex.RLock()
 	defer t.dialerMutex.RUnlock()
 
-	return t.baseTelegram.dial(dc, connOpts.ConnectionProto)
+	return t.baseTelegram.dial(dc, sock, connOpts.ConnectionProto)
 }
 
 func (t *middleTelegramCaller) autoUpdate() {
diff --git a/telegram/telegram.go b/telegram/telegram.go
index 5ce16ad..9413882 100644
--- a/telegram/telegram.go
+++ b/telegram/telegram.go
@@ -13,7 +13,7 @@ import (
 // encapsulates logic of working with middleproxies or direct
 // connections.
 type Telegram interface {
-	Dial(*mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error)
+	Dial(string, *mtproto.ConnectionOpts) (wrappers.ReadWriteCloserWithAddr, error)
 	Init(*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr) (wrappers.ReadWriteCloserWithAddr, error)
 }
 
@@ -24,7 +24,7 @@ type baseTelegram struct {
 	v6Addresses map[int16][]string
 }
 
-func (b *baseTelegram) dial(dcIdx int16, proto mtproto.ConnectionProtocol) (wrappers.ReadWriteCloserWithAddr, error) {
+func (b *baseTelegram) dial(dcIdx int16, sock string, proto mtproto.ConnectionProtocol) (wrappers.ReadWriteCloserWithAddr, error) {
 	addrs := make([]string, 2)
 
 	if proto&mtproto.ConnectionProtocolIPv6 != 0 {
@@ -39,7 +39,7 @@ func (b *baseTelegram) dial(dcIdx int16, proto mtproto.ConnectionProtocol) (wrap
 	}
 
 	for _, addr := range addrs {
-		if conn, err := b.dialer.dialRWC(addr); err == nil {
+		if conn, err := b.dialer.dialRWC(addr, sock); err == nil {
 			return conn, err
 		}
 	}
diff --git a/utils/read_current_data.go b/utils/read_current_data.go
new file mode 100644
index 0000000..c8669e2
--- /dev/null
+++ b/utils/read_current_data.go
@@ -0,0 +1,20 @@
+package utils
+
+import "io"
+
+const readCurrentDataBufferSize = 1024 + 1
+
+func ReadCurrentData(src io.Reader) (rv []byte, err error) {
+	buf := make([]byte, readCurrentDataBufferSize)
+	n := readCurrentDataBufferSize
+
+	for n == len(buf) {
+		n, err = src.Read(buf)
+		if err != nil {
+			return nil, err
+		}
+		rv = append(rv, buf[:n]...)
+	}
+
+	return rv, nil
+}
diff --git a/wrappers/blockcipherrwc.go b/wrappers/blockcipherrwc.go
index 8d2a06c..af5ccf0 100644
--- a/wrappers/blockcipherrwc.go
+++ b/wrappers/blockcipherrwc.go
@@ -1,12 +1,14 @@
 package wrappers
 
 import (
-	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
+	"fmt"
 	"net"
 
 	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/utils"
 )
 
 type BlockCipherReadWriteCloserWithAddr struct {
@@ -19,16 +21,18 @@ type BlockCipherReadWriteCloserWithAddr struct {
 
 func (c *BlockCipherReadWriteCloserWithAddr) Read(p []byte) (int, error) {
 	return c.BufferedRead(p, func() error {
-		buf := &bytes.Buffer{}
-		for buf.Len()%aes.BlockSize != 0 || buf.Len() == 0 {
-			n, err := c.conn.Read(p)
+		var buf []byte
+
+		for len(buf) == 0 || len(buf)%aes.BlockSize != 0 {
+			rv, err := utils.ReadCurrentData(c.conn)
 			if err != nil {
 				return errors.Annotate(err, "Cannot read from socket")
 			}
-			buf.Write(p[:n])
+			buf = append(buf, rv...)
 		}
-		c.decryptor.CryptBlocks(buf.Bytes(), buf.Bytes())
-		c.Buffer.Write(buf.Bytes())
+
+		c.decryptor.CryptBlocks(buf, buf)
+		c.Buffer.Write(buf)
 
 		return nil
 	})
@@ -46,6 +50,7 @@ func (c *BlockCipherReadWriteCloserWithAddr) Write(p []byte) (int, error) {
 }
 
 func (c *BlockCipherReadWriteCloserWithAddr) Close() error {
+	fmt.Println("BlockCipherReadWriteCloserWithAddr closes", "sockid", c.SocketID(), "bufsize", c.Buffer.Len())
 	return c.conn.Close()
 }
 
@@ -57,6 +62,10 @@ func (c *BlockCipherReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
 	return c.conn.RemoteAddr()
 }
 
+func (c *BlockCipherReadWriteCloserWithAddr) SocketID() string {
+	return c.conn.SocketID()
+}
+
 func NewBlockCipherRWC(conn ReadWriteCloserWithAddr, encryptor, decryptor cipher.BlockMode) ReadWriteCloserWithAddr {
 	return &BlockCipherReadWriteCloserWithAddr{
 		BufferedReader: NewBufferedReader(),
diff --git a/wrappers/buffered_reader.go b/wrappers/buffered_reader.go
index cb48d56..09d3bca 100644
--- a/wrappers/buffered_reader.go
+++ b/wrappers/buffered_reader.go
@@ -1,11 +1,19 @@
 package wrappers
 
-import "bytes"
+import (
+	"bytes"
+
+	"github.com/juju/errors"
+)
 
 type BufferedReader struct {
 	Buffer *bytes.Buffer
 }
 
+var (
+	BufferedReaderContinue = errors.New("Please continue reading")
+)
+
 func (b *BufferedReader) BufferedRead(p []byte, callback func() error) (int, error) {
 	if b.Buffer.Len() > 0 {
 		return b.flush(p)
@@ -17,7 +25,7 @@ func (b *BufferedReader) BufferedRead(p []byte, callback func() error) (int, err
 }
 
 func (b *BufferedReader) flush(p []byte) (int, error) {
-	if b.Buffer.Len() < len(p) {
+	if b.Buffer.Len() <= len(p) {
 		sizeToReturn := b.Buffer.Len()
 		copy(p, b.Buffer.Bytes())
 		b.Buffer.Reset()
diff --git a/wrappers/ctxrwc.go b/wrappers/ctxrwc.go
index 64c1006..fe6f515 100644
--- a/wrappers/ctxrwc.go
+++ b/wrappers/ctxrwc.go
@@ -56,6 +56,10 @@ func (c *CtxReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
 	return c.conn.RemoteAddr()
 }
 
+func (c *CtxReadWriteCloserWithAddr) SocketID() string {
+	return c.conn.SocketID()
+}
+
 // NewCtxRWC returns ReadWriteCloser which respects given context,
 // cancellation etc.
 func NewCtxRWC(ctx context.Context, cancel context.CancelFunc, conn ReadWriteCloserWithAddr) ReadWriteCloserWithAddr {
diff --git a/wrappers/logrwc.go b/wrappers/logrwc.go
index fd80ad4..74d3cbf 100644
--- a/wrappers/logrwc.go
+++ b/wrappers/logrwc.go
@@ -18,14 +18,14 @@ type LogReadWriteCloserWithAddr struct {
 // Read reads from connection
 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)
+	l.logger.Debugw("Finish reading", "name", l.name, "socketid", l.sockid, "nbytes", n, "error", err, "localAddr", l.LocalAddr())
 	return
 }
 
 // Write writes into connection.
 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)
+	l.logger.Debugw("Finish writing", "name", l.name, "socketid", l.sockid, "nbytes", n, "error", err, "localAddr", l.LocalAddr())
 	return
 }
 
@@ -44,6 +44,10 @@ func (l *LogReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
 	return l.conn.RemoteAddr()
 }
 
+func (l *LogReadWriteCloserWithAddr) SocketID() string {
+	return l.sockid
+}
+
 // NewLogRWC wraps ReadWriteCloser with logger calls.
 func NewLogRWC(conn ReadWriteCloserWithAddr, logger *zap.SugaredLogger, sockid string, name string) ReadWriteCloserWithAddr {
 	return &LogReadWriteCloserWithAddr{
diff --git a/wrappers/rwcaddr.go b/wrappers/rwcaddr.go
index 6ca4299..fee186c 100644
--- a/wrappers/rwcaddr.go
+++ b/wrappers/rwcaddr.go
@@ -10,4 +10,5 @@ type ReadWriteCloserWithAddr interface {
 
 	LocalAddr() *net.TCPAddr
 	RemoteAddr() *net.TCPAddr
+	SocketID() string
 }
diff --git a/wrappers/streamcipherrwc.go b/wrappers/streamcipherrwc.go
index a5f15d1..f9ad598 100644
--- a/wrappers/streamcipherrwc.go
+++ b/wrappers/streamcipherrwc.go
@@ -51,6 +51,10 @@ func (c *StreamCipherReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
 	return c.conn.RemoteAddr()
 }
 
+func (c *StreamCipherReadWriteCloserWithAddr) SocketID() string {
+	return c.conn.SocketID()
+}
+
 // NewStreamCipherRWC returns wrapper which transparently
 // encrypts/decrypts traffic with obfuscated2 protocol.
 func NewStreamCipherRWC(conn ReadWriteCloserWithAddr, encryptor, decryptor cipher.Stream) ReadWriteCloserWithAddr {
diff --git a/wrappers/timeoutrwc.go b/wrappers/timeoutrwc.go
index 952ade1..3c48546 100644
--- a/wrappers/timeoutrwc.go
+++ b/wrappers/timeoutrwc.go
@@ -9,6 +9,7 @@ import (
 
 type TimeoutReadWriteCloserWithAddr struct {
 	conn       net.Conn
+	sock       string
 	publicIPv4 net.IP
 	publicIPv6 net.IP
 }
@@ -46,10 +47,15 @@ func (t *TimeoutReadWriteCloserWithAddr) LocalAddr() *net.TCPAddr {
 	return &newAddr
 }
 
-func NewTimeoutRWC(conn net.Conn, ipv4, ipv6 net.IP) ReadWriteCloserWithAddr {
+func (t *TimeoutReadWriteCloserWithAddr) SocketID() string {
+	return t.sock
+}
+
+func NewTimeoutRWC(conn net.Conn, sock string, ipv4, ipv6 net.IP) ReadWriteCloserWithAddr {
 	return &TimeoutReadWriteCloserWithAddr{
 		conn:       conn,
 		publicIPv4: ipv4,
 		publicIPv6: ipv6,
+		sock:       sock,
 	}
 }
diff --git a/wrappers/trafficrwc.go b/wrappers/trafficrwc.go
index 60dd8c2..344a451 100644
--- a/wrappers/trafficrwc.go
+++ b/wrappers/trafficrwc.go
@@ -37,6 +37,10 @@ func (t *TrafficReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
 	return t.conn.RemoteAddr()
 }
 
+func (t *TrafficReadWriteCloserWithAddr) SocketID() string {
+	return t.conn.SocketID()
+}
+
 // NewTrafficRWC wraps ReadWriteCloser to have read/write callbacks.
 func NewTrafficRWC(conn ReadWriteCloserWithAddr, readCallback, writeCallback func(int)) ReadWriteCloserWithAddr {
 	return &TrafficReadWriteCloserWithAddr{