diff --git a/transport.go b/transport.go index f295220..be242d5 100644 --- a/transport.go +++ b/transport.go @@ -4,6 +4,7 @@ import ( "context" "errors" "net" + "sync/atomic" "time" "github.com/pion/stun" @@ -24,7 +25,19 @@ func (a *Agent) Accept(ctx context.Context, remoteUfrag, remotePwd string) (*Con // Conn represents the ICE connection. // At the moment the lifetime of the Conn is equal to the Agent. type Conn struct { - agent *Agent + agent *Agent + bytesReceived uint64 + bytesSent uint64 +} + +// BytesSent returns the number of bytes sent +func (c *Conn) BytesSent() uint64 { + return atomic.LoadUint64(&c.bytesSent) +} + +// BytesReceived returns the number of bytes received +func (c *Conn) BytesReceived() uint64 { + return atomic.LoadUint64(&c.bytesReceived) } func (a *Agent) connect(ctx context.Context, isControlling bool, remoteUfrag, remotePwd string) (*Conn, error) { @@ -63,7 +76,9 @@ func (c *Conn) Read(p []byte) (int, error) { return 0, err } - return c.agent.buffer.Read(p) + n, err := c.agent.buffer.Read(p) + atomic.AddUint64(&c.bytesReceived, uint64(n)) + return n, err } // Write implements the Conn Write method. @@ -82,6 +97,7 @@ func (c *Conn) Write(p []byte) (int, error) { return 0, err } + atomic.AddUint64(&c.bytesSent, uint64(len(p))) return pair.Write(p) } diff --git a/transport_test.go b/transport_test.go index 401e7d6..08a5ca6 100644 --- a/transport_test.go +++ b/transport_test.go @@ -2,6 +2,7 @@ package ice import ( "context" + "errors" "net" "sync" "testing" @@ -401,3 +402,43 @@ func randomPort(t testing.TB) int { return 0 } } + +func TestConnStats(t *testing.T) { + ca, cb := pipe() + if _, err := ca.Write(make([]byte, 10)); err != nil { + t.Fatal("unexpected error trying to write") + } + + var wg sync.WaitGroup + wg.Add(1) + go func() { + buf := make([]byte, 10) + if _, err := cb.Read(buf); err != nil { + panic(errors.New("unexpected error trying to read")) + } + wg.Done() + }() + + wg.Wait() + + if ca.BytesSent() != 10 { + t.Fatal("bytes sent don't match") + } + + if cb.BytesReceived() != 10 { + t.Fatal("bytes received don't match") + } + + err := ca.Close() + if err != nil { + // we should never get here. + panic(err) + } + + err = cb.Close() + if err != nil { + // we should never get here. + panic(err) + } + +}