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:
@@ -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")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user