314c3c934d
This diff replaces engine/netx code with netxlite code in the engine/session.go file. To this end, we needed to move some code from engine/netx to netxlite. While there, we did review and improve the unit tests. A notable change in this diff is (or seems to be) that in engine/session.go we're not filtering for bogons anymore so that, in principle, we could believe a resolver returning to us bogon IP addresses for OONI services. However, I did not bother with changing bogons filtering because the sessionresolver package is already filtering for bogons, so it is actually okay to avoid doing that again the session.go code. See: https://github.com/ooni/probe-cli/blob/v3.15.0-alpha.1/internal/engine/internal/sessionresolver/resolvermaker.go#L88 There are two reference issues for this cleanup: 1. https://github.com/ooni/probe/issues/2115 2. https://github.com/ooni/probe/issues/2121
45 lines
1.4 KiB
Go
45 lines
1.4 KiB
Go
package bytecounter
|
|
|
|
//
|
|
// Implicit byte counting based on context
|
|
//
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
)
|
|
|
|
type byteCounterSessionKey struct{}
|
|
|
|
// ContextSessionByteCounter retrieves the session byte counter from the context
|
|
func ContextSessionByteCounter(ctx context.Context) *Counter {
|
|
counter, _ := ctx.Value(byteCounterSessionKey{}).(*Counter)
|
|
return counter
|
|
}
|
|
|
|
// WithSessionByteCounter assigns the session byte counter to the context.
|
|
func WithSessionByteCounter(ctx context.Context, counter *Counter) context.Context {
|
|
return context.WithValue(ctx, byteCounterSessionKey{}, counter)
|
|
}
|
|
|
|
type byteCounterExperimentKey struct{}
|
|
|
|
// ContextExperimentByteCounter retrieves the experiment byte counter from the context
|
|
func ContextExperimentByteCounter(ctx context.Context) *Counter {
|
|
counter, _ := ctx.Value(byteCounterExperimentKey{}).(*Counter)
|
|
return counter
|
|
}
|
|
|
|
// WithExperimentByteCounter assigns the experiment byte counter to the context.
|
|
func WithExperimentByteCounter(ctx context.Context, counter *Counter) context.Context {
|
|
return context.WithValue(ctx, byteCounterExperimentKey{}, counter)
|
|
}
|
|
|
|
// MaybeWrapWithContextByteCounters wraps a conn with the byte counters
|
|
// that have previosuly been configured into a context.
|
|
func MaybeWrapWithContextByteCounters(ctx context.Context, conn net.Conn) net.Conn {
|
|
conn = MaybeWrap(conn, ContextExperimentByteCounter(ctx))
|
|
conn = MaybeWrap(conn, ContextSessionByteCounter(ctx))
|
|
return conn
|
|
}
|