c/ooni-probe-cli
Watch
1
0
Fork
You've already forked ooni-probe-cli
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

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

Browse source 
* 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
commit 760ac905d6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 104 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

39
internal/engine/netx/dialer/bytecounter.go
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
}