Add statistics

go.mod

@@ -1,3 +1,5 @@
 module github.com/maride/afl-transmit
 
 go 1.13
+
+require github.com/dustin/go-humanize v1.0.0

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=

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 {

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

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()
 }

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)
+	}
+}