feat(proxy): add traffic counting wrapper for connections

server/proxy/traffic_counter.go

@@ -0,0 +1,36 @@
+package proxy
+
+import "io"
+
+type countingReadWriteCloser struct {
+	io.ReadWriteCloser
+	onRead  func(int64)
+	onWrite func(int64)
+}
+
+func wrapCountingReadWriteCloser(rwc io.ReadWriteCloser, onRead, onWrite func(int64)) io.ReadWriteCloser {
+	if onRead == nil && onWrite == nil {
+		return rwc
+	}
+	return &countingReadWriteCloser{
+		ReadWriteCloser: rwc,
+		onRead:          onRead,
+		onWrite:         onWrite,
+	}
+}
+
+func (c *countingReadWriteCloser) Read(p []byte) (n int, err error) {
+	n, err = c.ReadWriteCloser.Read(p)
+	if n > 0 && c.onRead != nil {
+		c.onRead(int64(n))
+	}
+	return
+}
+
+func (c *countingReadWriteCloser) Write(p []byte) (n int, err error) {
+	n, err = c.ReadWriteCloser.Write(p)
+	if n > 0 && c.onWrite != nil {
+		c.onWrite(int64(n))
+	}
+	return
+}