Refactor trafficrwc to wrappers

proxy/server.go

@@ -113,7 +113,7 @@ func (s *Server) makeSocketID() string {
 
 func (s *Server) getClientStream(ctx context.Context, cancel context.CancelFunc, conn net.Conn, socketID string) (io.ReadWriteCloser, int16, error) {
 	wConn := wrappers.NewTimeoutRWC(conn, s.readTimeout, s.writeTimeout)
-	wConn = newTrafficReadWriteCloser(wConn, s.stats.addIncomingTraffic, s.stats.addOutgoingTraffic)
+	wConn = wrappers.NewTrafficRWC(wConn, s.stats.addIncomingTraffic, s.stats.addOutgoingTraffic)
 	frame, err := obfuscated2.ExtractFrame(wConn)
 	if err != nil {
 		return nil, 0, errors.Annotate(err, "Cannot create client stream")
@@ -137,7 +137,7 @@ func (s *Server) getTelegramStream(ctx context.Context, cancel context.CancelFun
 		return nil, errors.Annotate(err, "Cannot dial")
 	}
 	wConn := wrappers.NewTimeoutRWC(socket, s.readTimeout, s.writeTimeout)
-	wConn = newTrafficReadWriteCloser(wConn, s.stats.addIncomingTraffic, s.stats.addOutgoingTraffic)
+	wConn = wrappers.NewTrafficRWC(wConn, s.stats.addIncomingTraffic, s.stats.addOutgoingTraffic)
 
 	obfs2, frame := obfuscated2.MakeTelegramObfuscated2Frame()
 	if n, err := socket.Write(frame); err != nil || n != len(frame) {

proxy/trafficrwc.go

@@ -1,38 +0,0 @@
-package proxy
-
-import "io"
-
-// TrafficReadWriteCloser counts an amount of ingress/egress traffic by
-// calling given callbacks.
-type TrafficReadWriteCloser struct {
-	conn          io.ReadWriteCloser
-	readCallback  func(int)
-	writeCallback func(int)
-}
-
-// Read reads from connection
-func (t *TrafficReadWriteCloser) 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) {
-	n, err = t.conn.Write(p)
-	t.writeCallback(n)
-	return
-}
-
-// Close closes underlying connection.
-func (t *TrafficReadWriteCloser) Close() error {
-	return t.conn.Close()
-}
-
-func newTrafficReadWriteCloser(conn io.ReadWriteCloser, readCallback, writeCallback func(int)) io.ReadWriteCloser {
-	return &TrafficReadWriteCloser{
-		conn:          conn,
-		readCallback:  readCallback,
-		writeCallback: writeCallback,
-	}
-}