Refactor trafficrwc to wrappers

wrappers/timeoutrwc.go

@@ -0,0 +1,40 @@
+package wrappers
+
+import (
+	"io"
+	"net"
+	"time"
+)
+
+// TimeoutReadWriteCloser sets timeouts for read/write into underlying
+// network connection.
+type TimeoutReadWriteCloser struct {
+	conn         net.Conn
+	readTimeout  time.Duration
+	writeTimeout time.Duration
+}
+
+// Read reads from connection
+func (t *TimeoutReadWriteCloser) Read(p []byte) (int, error) {
+	t.conn.SetReadDeadline(time.Now().Add(t.readTimeout)) // nolint: errcheck, gas
+	return t.conn.Read(p)
+}
+
+// Write writes into connection.
+func (t *TimeoutReadWriteCloser) Write(p []byte) (int, error) {
+	t.conn.SetWriteDeadline(time.Now().Add(t.writeTimeout)) // nolint: errcheck, gas
+	return t.conn.Write(p)
+}
+
+// Close closes underlying connection.
+func (t *TimeoutReadWriteCloser) Close() error {
+	return t.conn.Close()
+}
+
+func NewTimeoutRWC(conn net.Conn, readTimeout, writeTimeout time.Duration) io.ReadWriteCloser {
+	return &TimeoutReadWriteCloser{
+		conn:         conn,
+		readTimeout:  readTimeout,
+		writeTimeout: writeTimeout,
+	}
+}

wrappers/trafficrwc.go

@@ -0,0 +1,38 @@
+package wrappers
+
+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 NewTrafficRWC(conn io.ReadWriteCloser, readCallback, writeCallback func(int)) io.ReadWriteCloser {
+	return &TrafficReadWriteCloser{
+		conn:          conn,
+		readCallback:  readCallback,
+		writeCallback: writeCallback,
+	}
+}