diff --git a/knockr.go b/knockr.go index 48450e5..7442b02 100644 --- a/knockr.go +++ b/knockr.go @@ -188,12 +188,17 @@ func proxy(c net.Conn) { fmt.Println("[ERR] Proxy connection to server failed") fmt.Println(" Error is ", err) } else { - // A bit of cluttery code for statistics purposes, sorry. - go func() { - delta_traffic_in, _ := io.Copy(c, ln) - traffic_in += delta_traffic_in - }() - delta_traffic_out, _ := io.Copy(ln, c) - traffic_out += delta_traffic_out + go proxy_writefunc(c, ln, &traffic_in) + proxy_writefunc(ln, c, &traffic_out) + } +} + +func proxy_writefunc(a net.Conn, b net.Conn, written_bytes *int64) { + var delta int64 = 0 + var err error + + for err == nil { + delta, err = io.CopyN(a, b, 1024) + *written_bytes += delta } }