From 10e940126d0cd52994e6c1f12a5cedc007a6267d Mon Sep 17 00:00:00 2001
From: maride <git@maride.cc>
Date: Tue, 20 Apr 2021 10:28:03 +0200
Subject: [PATCH] Add statistics

---
 go.mod          |  2 ++
 go.sum          |  2 ++
 main.go         |  4 ++++
 net/listener.go |  4 ++++
 net/peer.go     |  6 +++++-
 stats/stats.go  | 38 ++++++++++++++++++++++++++++++++++++++
 6 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 go.sum
 create mode 100644 stats/stats.go

diff --git a/go.mod b/go.mod
index 057686f..0618dc8 100644
--- a/go.mod
+++ b/go.mod
@@ -1,3 +1,5 @@
 module github.com/maride/afl-transmit
 
 go 1.13
+
+require github.com/dustin/go-humanize v1.0.0
diff --git a/go.sum b/go.sum
new file mode 100644
index 0000000..4d39dd1
--- /dev/null
+++ b/go.sum
@@ -0,0 +1,2 @@
+github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo=
+github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
diff --git a/main.go b/main.go
index 5f4ca84..e781f9b 100644
--- a/main.go
+++ b/main.go
@@ -4,6 +4,7 @@ import (
 	"flag"
 	"fmt"
 	"github.com/maride/afl-transmit/net"
+	"github.com/maride/afl-transmit/stats"
 	"github.com/maride/afl-transmit/watchdog"
 	"log"
 )
@@ -32,6 +33,9 @@ func main() {
 	// Start watchdog for local afl instances
 	go watchdog.WatchFuzzers(outputDirectory)
 
+	// Start stat printer
+	go stats.PrintStats()
+
 	// Listen for incoming connections
 	listenErr := net.Listen(outputDirectory)
 	if listenErr != nil {
diff --git a/net/listener.go b/net/listener.go
index da7c770..2f8130a 100644
--- a/net/listener.go
+++ b/net/listener.go
@@ -4,6 +4,7 @@ import (
 	"flag"
 	"fmt"
 	"github.com/maride/afl-transmit/logistic"
+	"github.com/maride/afl-transmit/stats"
 	"io"
 	"io/ioutil"
 	"log"
@@ -82,6 +83,9 @@ func handle(conn net.Conn, outputDirectory string) {
 			log.Printf("Encountered error processing packet from %s: %s", conn.RemoteAddr().String(), unpackErr)
 		}
 
+		// Push read bytes to stats
+		stats.PushStat(stats.Stat{ReceivedBytes: uint64(len(cont))})
+
 		return
 	} else {
 		// We encountered an error on that connection
diff --git a/net/peer.go b/net/peer.go
index d4ef47b..07e54ac 100644
--- a/net/peer.go
+++ b/net/peer.go
@@ -2,6 +2,7 @@ package net
 
 import (
 	"fmt"
+	"github.com/maride/afl-transmit/stats"
 	"log"
 	"net"
 	"regexp"
@@ -44,12 +45,15 @@ func (p *Peer) SendToPeer(content []byte) {
 	}
 
 	// Send
-	_, writeErr := tcpConn.Write(content)
+	written, writeErr := tcpConn.Write(content)
 	if writeErr != nil {
 		log.Printf("Unable to write to peer %s: %s", tcpConn.RemoteAddr().String(), writeErr)
 		return
 	}
 
+	// Push written bytes to stats
+	stats.PushStat(stats.Stat{SentBytes: uint64(written)})
+
 	// Close connection
 	tcpConn.Close()
 }
diff --git a/stats/stats.go b/stats/stats.go
new file mode 100644
index 0000000..bc50114
--- /dev/null
+++ b/stats/stats.go
@@ -0,0 +1,38 @@
+package stats
+
+import (
+	"fmt"
+	"github.com/dustin/go-humanize"
+	"time"
+)
+
+// Stat bundles the metrics we collect into a single struct
+type Stat struct {
+	SentBytes uint64
+	ReceivedBytes uint64
+}
+
+// statPipe is a channel used to
+var stats Stat
+
+// PushStat pushes the given stat
+func PushStat(s Stat) {
+	stats.SentBytes += s.SentBytes
+	stats.ReceivedBytes += s.ReceivedBytes
+}
+
+// PrintStats periodically prints the collected statistics
+func PrintStats() {
+	t := time.NewTicker(2 * time.Second)
+
+	for {
+		// Wait until we get a tick
+		<-t.C
+
+		// Format numbers and write them out
+		bIn := humanize.Bytes(stats.ReceivedBytes)
+		bOut := humanize.Bytes(stats.SentBytes)
+
+		fmt.Printf("%s in / %s out\r", bIn, bOut)
+	}
+}
\ No newline at end of file