refactor: move bytecounting conn in bytecounter pkg (#392)

* refactor: move bytecounting conn in bytecounter pkg

This enables other pieces of code to request bytecounting without
depending on netx or on the perverse using-the-context-to-configure-
byte-counting mechanism.

Also occurred when working on https://github.com/ooni/probe/issues/1687

* fix: add missing docs
This commit is contained in:
Simone Basso
2021-06-22 13:44:36 +02:00
committed by GitHub
parent 23bc261464
commit 760ac905d6
4 changed files with 104 additions and 33 deletions
+6 -33
View File
@@ -20,12 +20,13 @@ func (d *byteCounterDialer) DialContext(
if err != nil {
return nil, err
}
exp := contextExperimentByteCounter(ctx)
sess := contextSessionByteCounter(ctx)
if exp == nil && sess == nil {
return conn, nil // no point in wrapping
if exp := contextExperimentByteCounter(ctx); exp != nil {
conn = &bytecounter.Conn{Conn: conn, Counter: exp}
}
return &byteCounterConnWrapper{Conn: conn, exp: exp, sess: sess}, nil
if sess := contextSessionByteCounter(ctx); sess != nil {
conn = &bytecounter.Conn{Conn: conn, Counter: sess}
}
return conn, nil
}
type byteCounterSessionKey struct{}
@@ -53,31 +54,3 @@ func contextExperimentByteCounter(ctx context.Context) *bytecounter.Counter {
func WithExperimentByteCounter(ctx context.Context, counter *bytecounter.Counter) context.Context {
return context.WithValue(ctx, byteCounterExperimentKey{}, counter)
}
type byteCounterConnWrapper struct {
net.Conn
exp *bytecounter.Counter
sess *bytecounter.Counter
}
func (c *byteCounterConnWrapper) Read(p []byte) (int, error) {
count, err := c.Conn.Read(p)
if c.exp != nil {
c.exp.CountBytesReceived(count)
}
if c.sess != nil {
c.sess.CountBytesReceived(count)
}
return count, err
}
func (c *byteCounterConnWrapper) Write(p []byte) (int, error) {
count, err := c.Conn.Write(p)
if c.exp != nil {
c.exp.CountBytesSent(count)
}
if c.sess != nil {
c.sess.CountBytesSent(count)
}
return count, err
}
@@ -48,9 +48,15 @@ func TestByteCounterNormalUsage(t *testing.T) {
if err := dorequest(ctx, "http://facebook.com"); err != nil {
t.Fatal(err)
}
if exp.Received.Load() <= 0 {
t.Fatal("experiment should have received some bytes")
}
if sess.Received.Load() <= exp.Received.Load() {
t.Fatal("session should have received more than experiment")
}
if exp.Sent.Load() <= 0 {
t.Fatal("experiment should have sent some bytes")
}
if sess.Sent.Load() <= exp.Sent.Load() {
t.Fatal("session should have sent more than experiment")
}