diff --git a/datacap/client.go b/datacap/client.go new file mode 100644 index 00000000..0a80657a --- /dev/null +++ b/datacap/client.go @@ -0,0 +1,108 @@ +// Package datacap reports per-device byte usage to the local datacap sidecar +// and enforces the throttle verdict the sidecar returns. +// +// It replaces the reporting-Redis path (see the redis package): instead of +// writing `_client:` hashes to a shared Redis and reading cohort +// settings back out of a `_throttle` key, the proxy POSTs deltas to a sidecar +// on localhost that owns the cap accounting and answers with the current +// throttle state. The wire contract is the same one lantern-box speaks +// (getlantern/lantern-box tracker/datacap), so a device's traffic accumulates +// into one counter no matter which proxy flavor carried it. +package datacap + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" +) + +// DefaultHTTPTimeout bounds a single report round-trip. The sidecar is on +// localhost and answers from memory, so anything slower than this means it is +// wedged and the delta is better retried on the next cycle than left in flight. +const DefaultHTTPTimeout = 10 * time.Second + +// Report is the body of POST /data-cap/. BytesUsed is a delta since the last +// report, not a running total — the sidecar accumulates. +type Report struct { + DeviceID string `json:"deviceId"` + CountryCode string `json:"countryCode"` + Platform string `json:"platform"` + BytesUsed int64 `json:"bytesUsed"` +} + +// Status is the sidecar's answer: the throttle verdict plus enough of the +// device's cap state to render the XBQ headers clients show in their usage UI. +type Status struct { + Throttle bool `json:"throttle"` + CapLimit int64 `json:"capLimit"` + ExpiryTime int64 `json:"expiryTime"` // Unix seconds + BytesUsed int64 `json:"bytesUsed"` +} + +// Client talks to the datacap sidecar over HTTP. +type Client struct { + httpClient *http.Client + baseURL string +} + +// NewClient returns a Client posting to baseURL, e.g. "http://127.0.0.1:8078". +// A bare host:port is assumed to be plain HTTP: the sidecar listens on loopback +// (or the phost bridge address) without TLS. +func NewClient(baseURL string, timeout time.Duration) *Client { + if !strings.HasPrefix(baseURL, "http://") && !strings.HasPrefix(baseURL, "https://") { + baseURL = "http://" + baseURL + } + if timeout <= 0 { + timeout = DefaultHTTPTimeout + } + return &Client{ + httpClient: &http.Client{ + Timeout: timeout, + // A bare Transport also deliberately ignores HTTP_PROXY et al. — + // this client only ever talks to the local sidecar. The idle pool + // matches the flush fan-out so a full cycle's connections are all + // reusable instead of the excess being closed each cycle. + Transport: &http.Transport{ + MaxIdleConnsPerHost: flushConcurrency, + }, + }, + baseURL: strings.TrimSuffix(baseURL, "/"), + } +} + +// ReportUsage posts a usage delta and returns the device's updated cap state. +func (c *Client) ReportUsage(ctx context.Context, report *Report) (*Status, error) { + body, err := json.Marshal(report) + if err != nil { + return nil, fmt.Errorf("marshal report: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/data-cap/", bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("build request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("post usage: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted { + detail, _ := io.ReadAll(io.LimitReader(resp.Body, 1024)) + return nil, fmt.Errorf("sidecar returned HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(detail))) + } + + var status Status + if err := json.NewDecoder(resp.Body).Decode(&status); err != nil { + return nil, fmt.Errorf("decode status: %w", err) + } + return &status, nil +} diff --git a/datacap/tracker.go b/datacap/tracker.go new file mode 100644 index 00000000..d95ad262 --- /dev/null +++ b/datacap/tracker.go @@ -0,0 +1,368 @@ +package datacap + +import ( + "context" + "net" + "sync" + "sync/atomic" + "time" + + "github.com/getlantern/geo" + "github.com/getlantern/golog" + "github.com/getlantern/measured" + + "github.com/getlantern/http-proxy-lantern/v2/common" + "github.com/getlantern/http-proxy-lantern/v2/listeners" +) + +var log = golog.LoggerFor("datacap") + +const ( + // DefaultReportInterval is how often accumulated deltas are flushed to the + // sidecar. The measured pipeline hands us per-connection deltas on its own + // (slower) cadence; flushing more often than that just shortens the gap + // between a delta arriving and the throttle verdict coming back. + DefaultReportInterval = 10 * time.Second + + // ThrottledWriteRate is the download rate a capped device is held to, in + // bytes per second. Matches lantern-box (tracker/datacap/conn.go) so a + // device sees the same speed whichever proxy flavor it lands on. + ThrottledWriteRate int64 = 16 * 1024 // 128 Kb/s + + // flushConcurrency bounds the reports in flight during one cycle, so a + // proxy with many active devices does not open an unbounded number of + // connections to the sidecar. + flushConcurrency = 16 + + // idleDeviceTTL is how long a device with nothing pending is kept around. + // Evicting a device severs its open connections from re-rating (they keep + // the limiter pointer they were attached with, while a re-appearing device + // gets a fresh entry), so this must comfortably exceed the proxy's + // idle-close timeout (`idleclose`, ~70-90s): a connection idle longer than + // that is closed by the proxy itself, and one that is still alive is moving + // bytes, whose deltas refresh lastSeen and block eviction. At 30 minutes + // the margin over idleclose plus the measured reporting interval is >10x. + idleDeviceTTL = 30 * time.Minute +) + +// Usage is a device's cap state as of the last sidecar response. +type Usage struct { + // BytesUsed is the total consumed in the current allotment period. + BytesUsed int64 + // CapLimit is the allotment in bytes. Zero means this device is uncapped + // (no cap entry for its country/platform). + CapLimit int64 + // Expiry is when the current allotment resets. + Expiry time.Time + // AsOf is when the sidecar reported these numbers. A zero AsOf means the + // sidecar has not answered for this device yet. + AsOf time.Time + // Throttled is the sidecar's verdict at AsOf. + Throttled bool +} + +// device holds the per-device state shared between the reporting loop and the +// request filter. +type device struct { + // limiter is attached to every connection of this device and re-rated in + // place when the throttle verdict changes, so a transfer that is already + // running slows down at the moment the cap is crossed instead of at the + // next CONNECT. + limiter *listeners.RateLimiter + // unthrottledLimiter serves requests to domains excluded from the cap. It + // stays at the default rate for the device's lifetime; those requests are + // still counted, they are just never slowed to the capped rate. + unthrottledLimiter *listeners.RateLimiter + + usage Usage + + // pendingBytes is the delta not yet accepted by the sidecar. + pendingBytes int64 + // countryCode, platform key the sidecar's cap-limit lookup. countryCode is + // sticky — set once from the first delta that resolves one — matching the + // reporting-Redis behavior this replaces. + countryCode string + platform string + // lastSeen gates eviction of idle devices. + lastSeen time.Time +} + +// Tracker aggregates per-device byte usage, reports it to the sidecar, and owns +// the rate limiters the verdict is enforced through. +type Tracker struct { + client *Client + countryLookup geo.CountryLookup + // defaultRate is the ceiling every non-pro device is held to regardless of + // its cap state, to keep bandwidth hogs from monopolizing a proxy. + defaultRate int64 + reportInterval time.Duration + + mx sync.RWMutex + devices map[string]*device +} + +// TrackerOpts configures a Tracker. +type TrackerOpts struct { + Client *Client + CountryLookup geo.CountryLookup + DefaultRate int64 + ReportInterval time.Duration +} + +// NewTracker starts a Tracker and its reporting loop. +func NewTracker(opts TrackerOpts) *Tracker { + if opts.ReportInterval <= 0 { + opts.ReportInterval = DefaultReportInterval + } + if opts.CountryLookup == nil { + // Reports still carry the device and its platform; only the + // country-specific cap limit is lost. That beats panicking in the + // reporting loop. + opts.CountryLookup = geo.NoLookup{} + } + t := &Tracker{ + client: opts.Client, + countryLookup: opts.CountryLookup, + defaultRate: opts.DefaultRate, + reportInterval: opts.ReportInterval, + devices: make(map[string]*device), + } + go t.reportPeriodically() + return t +} + +// Reporter returns the callback the measured listener feeds connection deltas +// into. Deltas are folded into per-device state synchronously: the tracker lock +// is never held across a sidecar call, so this cannot block on the network. +func (t *Tracker) Reporter() listeners.MeasuredReportFN { + return func(ctx map[string]interface{}, stats *measured.Stats, deltaStats *measured.Stats, final bool) { + bytes := int64(deltaStats.SentTotal) + int64(deltaStats.RecvTotal) + if bytes == 0 { + return + } + deviceID, _ := ctx[common.DeviceID].(string) + if deviceID == "" { + // Nothing to account against. + return + } + clientIP, _ := ctx[common.ClientIP].(string) + platform, _ := ctx[common.Platform].(string) + t.accumulate(deviceID, clientIP, platform, bytes) + } +} + +// Limiter returns the shared limiter to attach to a connection from deviceID. +// unthrottled selects the limiter used for requests to domains excluded from +// the cap, which is never re-rated to the capped speed. +func (t *Tracker) Limiter(deviceID string, unthrottled bool) *listeners.RateLimiter { + d := t.deviceFor(deviceID) + if unthrottled { + return d.unthrottledLimiter + } + return d.limiter +} + +// LimiterAndUsage returns the device's shared limiter together with the last +// cap state the sidecar reported for it, in one lookup. ok is false until the +// first report for that device has been answered. +func (t *Tracker) LimiterAndUsage(deviceID string) (limiter *listeners.RateLimiter, u Usage, ok bool) { + d := t.deviceFor(deviceID) + t.mx.RLock() + u = d.usage + t.mx.RUnlock() + return d.limiter, u, !u.AsOf.IsZero() +} + +// Usage returns the last cap state the sidecar reported for deviceID. ok is +// false until the first report for that device has been answered. +func (t *Tracker) Usage(deviceID string) (Usage, bool) { + t.mx.RLock() + defer t.mx.RUnlock() + d, exists := t.devices[deviceID] + if !exists || d.usage.AsOf.IsZero() { + return Usage{}, false + } + return d.usage, true +} + +func (t *Tracker) deviceFor(deviceID string) *device { + t.mx.RLock() + d, exists := t.devices[deviceID] + t.mx.RUnlock() + if exists { + return d + } + + t.mx.Lock() + defer t.mx.Unlock() + if d, exists = t.devices[deviceID]; exists { + return d + } + d = &device{ + limiter: listeners.NewRateLimiter(t.defaultRate, t.defaultRate), + unthrottledLimiter: listeners.NewRateLimiter(t.defaultRate, t.defaultRate), + lastSeen: time.Now(), + } + t.devices[deviceID] = d + return d +} + +func (t *Tracker) reportPeriodically() { + ticker := time.NewTicker(t.reportInterval) + defer ticker.Stop() + for range ticker.C { + // Bound the whole cycle. Reports run concurrently, but a wedged + // sidecar would still hold this goroutine for a full HTTP timeout, + // and for that entire window no throttle verdict is applied. + // Deadlined reports are restored and retried on the next tick like + // any other failure. The deadline never drops below the per-request + // timeout, so a healthy-but-slow sidecar is not cut off by an + // aggressive reportInterval. + timeout := 2 * t.reportInterval + if timeout < DefaultHTTPTimeout { + timeout = DefaultHTTPTimeout + } + ctx, cancel := context.WithTimeout(context.Background(), timeout) + t.flush(ctx) + cancel() + } +} + +func (t *Tracker) accumulate(deviceID, clientIP, platform string, bytes int64) { + d := t.deviceFor(deviceID) + + // The country lookup walks the MaxMind database, so do it at most once per + // device — outside the lock — rather than on every delta. + t.mx.RLock() + needCountry := d.countryCode == "" + t.mx.RUnlock() + countryCode := "" + if needCountry && clientIP != "" { + countryCode = t.countryLookup.CountryCode(net.ParseIP(clientIP)) + } + + now := time.Now() + t.mx.Lock() + d.pendingBytes += bytes + if d.countryCode == "" { + d.countryCode = countryCode + } + if platform != "" { + d.platform = platform + } + d.lastSeen = now + t.mx.Unlock() +} + +// pendingReport is one device's delta, detached from the tracker so the HTTP +// round-trips happen without holding the lock. +type pendingReport struct { + device *device + report Report +} + +func (t *Tracker) flush(ctx context.Context) { + evictBefore := time.Now().Add(-idleDeviceTTL) + + t.mx.Lock() + reports := make([]pendingReport, 0, len(t.devices)) + for deviceID, d := range t.devices { + if d.pendingBytes == 0 { + // The flush already visits every device, so idle eviction rides + // along instead of taking a second full scan under the lock. + if d.lastSeen.Before(evictBefore) { + delete(t.devices, deviceID) + } + continue + } + reports = append(reports, pendingReport{ + device: d, + report: Report{ + DeviceID: deviceID, + CountryCode: d.countryCode, + Platform: d.platform, + BytesUsed: d.pendingBytes, + }, + }) + d.pendingBytes = 0 + } + t.mx.Unlock() + + if len(reports) == 0 { + return + } + + // Report concurrently: serially, one slow device delays the throttle + // verdict for every device behind it in the batch, and the cycle's + // duration grows with the number of active devices. + statuses := make([]*Status, len(reports)) + errs := make([]error, len(reports)) + workers := flushConcurrency + if len(reports) < workers { + workers = len(reports) + } + var next atomic.Int64 + var wg sync.WaitGroup + for w := 0; w < workers; w++ { + wg.Add(1) + go func() { + defer wg.Done() + for { + i := int(next.Add(1)) - 1 + if i >= len(reports) { + return + } + pr := reports[i] + statuses[i], errs[i] = t.client.ReportUsage(ctx, &pr.report) + if st := statuses[i]; st != nil { + // Re-rate as soon as the verdict is in — SetRates is + // lock-free, and waiting for the whole batch would let one + // slow device delay enforcement for every other device. + // Only writes back to the client are throttled: a capped + // device can keep uploading at the default rate. + if st.Throttle { + pr.device.limiter.SetRates(t.defaultRate, ThrottledWriteRate) + } else { + pr.device.limiter.SetRates(t.defaultRate, t.defaultRate) + } + } + } + }() + } + wg.Wait() + + // The usage bookkeeping feeds the XBQ headers, which tolerate a batch's + // worth of staleness — record the whole cycle under one lock acquisition + // instead of one per device. + now := time.Now() + failed := 0 + var lastErr error + t.mx.Lock() + for i, pr := range reports { + if errs[i] != nil { + failed++ + lastErr = errs[i] + // Put the bytes back so the next cycle retries them. + pr.device.pendingBytes += pr.report.BytesUsed + continue + } + st := statuses[i] + u := Usage{ + BytesUsed: st.BytesUsed, + CapLimit: st.CapLimit, + AsOf: now, + Throttled: st.Throttle, + } + if st.ExpiryTime > 0 { + u.Expiry = time.Unix(st.ExpiryTime, 0) + } + pr.device.usage = u + } + t.mx.Unlock() + + if failed > 0 { + // A wedged sidecar fails every device in the batch, so log once per + // cycle rather than once per device. + log.Errorf("Unable to report usage for %d of %d devices, will retry: %v", failed, len(reports), lastErr) + } +} diff --git a/datacap/tracker_test.go b/datacap/tracker_test.go new file mode 100644 index 00000000..cb245f97 --- /dev/null +++ b/datacap/tracker_test.go @@ -0,0 +1,235 @@ +package datacap + +import ( + "encoding/json" + "net" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/getlantern/http-proxy-lantern/v2/common" + "github.com/getlantern/measured" +) + +const testDefaultRate = int64(640000) + +type fixedCountry string + +func (c fixedCountry) CountryCode(net.IP) string { return string(c) } + +// fakeSidecar accumulates reported deltas the way the real sidecar does and +// throttles once the cap is exceeded. +type fakeSidecar struct { + *httptest.Server + + mu sync.Mutex + reports []Report + total int64 + capLimit int64 + fail bool + delay time.Duration +} + +func newFakeSidecar(capLimit int64) *fakeSidecar { + s := &fakeSidecar{capLimit: capLimit} + s.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var report Report + if err := json.NewDecoder(r.Body).Decode(&report); err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + + s.mu.Lock() + if s.fail { + s.mu.Unlock() + w.WriteHeader(http.StatusInternalServerError) + return + } + delay := s.delay + if delay > 0 && report.DeviceID == "slowpoke" { + s.mu.Unlock() + time.Sleep(delay) + s.mu.Lock() + } + s.reports = append(s.reports, report) + s.total += report.BytesUsed + status := Status{ + Throttle: s.capLimit > 0 && s.total >= s.capLimit, + CapLimit: s.capLimit, + ExpiryTime: time.Now().Add(6 * time.Hour).Unix(), + BytesUsed: s.total, + } + s.mu.Unlock() + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(status) + })) + return s +} + +func (s *fakeSidecar) snapshot() (reports []Report, total int64) { + s.mu.Lock() + defer s.mu.Unlock() + return append([]Report(nil), s.reports...), s.total +} + +func (s *fakeSidecar) setFail(fail bool) { + s.mu.Lock() + defer s.mu.Unlock() + s.fail = fail +} + +func newTestTracker(t *testing.T, sidecar *fakeSidecar) *Tracker { + t.Helper() + return NewTracker(TrackerOpts{ + Client: NewClient(sidecar.URL, time.Second), + CountryLookup: fixedCountry("ES"), + DefaultRate: testDefaultRate, + ReportInterval: 10 * time.Millisecond, + }) +} + +func report(t *Tracker, deviceID string, bytes int) { + t.Reporter()(map[string]interface{}{ + common.DeviceID: deviceID, + common.ClientIP: "1.2.3.4", + common.Platform: "android", + }, &measured.Stats{}, &measured.Stats{RecvTotal: bytes}, false) +} + +// The whole point of the shared limiter: a device that crosses its cap is +// slowed on the limiter its already-open connections hold, not on a fresh one +// handed to the next connection. +func TestThrottleAppliesToTheLimiterAlreadyHandedOut(t *testing.T) { + sidecar := newFakeSidecar(1000) + defer sidecar.Close() + tracker := newTestTracker(t, sidecar) + + limiter := tracker.Limiter("device1", false) + require.Equal(t, testDefaultRate, limiter.GetRateWrite(), "should start at the default rate") + _, known := tracker.Usage("device1") + require.False(t, known, "usage is unknown until the sidecar answers") + + report(tracker, "device1", 400) + assert.Eventually(t, func() bool { + u, ok := tracker.Usage("device1") + return ok && u.BytesUsed == 400 && !u.Throttled + }, time.Second, 5*time.Millisecond, "usage should be reported below the cap without throttling") + assert.Equal(t, testDefaultRate, limiter.GetRateWrite(), "should not be throttled below the cap") + + report(tracker, "device1", 700) + assert.Eventually(t, func() bool { + return limiter.GetRateWrite() == ThrottledWriteRate + }, time.Second, 5*time.Millisecond, "crossing the cap should re-rate the limiter that was already handed out") + assert.Equal(t, testDefaultRate, limiter.GetRateRead(), "uploads stay at the default rate when capped") + assert.Same(t, limiter, tracker.Limiter("device1", false), "the device's limiter is shared, not rebuilt") +} + +// Bytes to cap-excluded domains still count, they are just never slowed. +func TestUnthrottledLimiterIsNeverCapped(t *testing.T) { + sidecar := newFakeSidecar(100) + defer sidecar.Close() + tracker := newTestTracker(t, sidecar) + + unthrottled := tracker.Limiter("device1", true) + report(tracker, "device1", 500) + assert.Eventually(t, func() bool { + u, ok := tracker.Usage("device1") + return ok && u.Throttled + }, time.Second, 5*time.Millisecond) + + assert.Equal(t, testDefaultRate, unthrottled.GetRateWrite()) + assert.Equal(t, ThrottledWriteRate, tracker.Limiter("device1", false).GetRateWrite()) +} + +func TestDeltasAreAggregatedPerDevice(t *testing.T) { + sidecar := newFakeSidecar(0) + defer sidecar.Close() + tracker := newTestTracker(t, sidecar) + + for i := 0; i < 5; i++ { + report(tracker, "device1", 10) + report(tracker, "device2", 3) + } + + assert.Eventually(t, func() bool { + _, total := sidecar.snapshot() + return total == 65 + }, time.Second, 5*time.Millisecond) + + reports, _ := sidecar.snapshot() + byDevice := map[string]int64{} + for _, r := range reports { + byDevice[r.DeviceID] += r.BytesUsed + assert.Equal(t, "ES", r.CountryCode, "country comes from the client IP lookup") + assert.Equal(t, "android", r.Platform) + } + assert.Equal(t, int64(50), byDevice["device1"]) + assert.Equal(t, int64(15), byDevice["device2"]) + assert.Less(t, len(reports), 10, "deltas should be batched, not posted one per call") +} + +// A sidecar that is down must not silently eat usage. +func TestFailedReportsAreRetried(t *testing.T) { + sidecar := newFakeSidecar(0) + defer sidecar.Close() + sidecar.setFail(true) + tracker := newTestTracker(t, sidecar) + + report(tracker, "device1", 250) + time.Sleep(50 * time.Millisecond) + _, total := sidecar.snapshot() + require.Zero(t, total, "nothing should be recorded while the sidecar is failing") + + sidecar.setFail(false) + assert.Eventually(t, func() bool { + _, total := sidecar.snapshot() + return total == 250 + }, time.Second, 5*time.Millisecond, "the delta should be retried once the sidecar recovers") +} + +func TestDeltasWithoutADeviceIDAreIgnored(t *testing.T) { + sidecar := newFakeSidecar(0) + defer sidecar.Close() + tracker := newTestTracker(t, sidecar) + + reporter := tracker.Reporter() + for _, ctx := range []map[string]interface{}{ + {common.ClientIP: "1.2.3.4"}, // absent + {common.DeviceID: "", common.ClientIP: "1.2.3.4"}, // present but empty + {common.DeviceID: 42, common.ClientIP: "1.2.3.4"}, // wrong type + } { + reporter(ctx, &measured.Stats{}, &measured.Stats{RecvTotal: 100}, false) + } + + time.Sleep(100 * time.Millisecond) + reports, total := sidecar.snapshot() + assert.Zero(t, total) + assert.Empty(t, reports) +} + +// One unresponsive device must not hold up the throttle verdict for every other +// device in the batch. +func TestASlowDeviceDoesNotDelayTheBatch(t *testing.T) { + sidecar := newFakeSidecar(1000) + defer sidecar.Close() + sidecar.mu.Lock() + sidecar.delay = 750 * time.Millisecond + sidecar.mu.Unlock() + + tracker := newTestTracker(t, sidecar) + limiter := tracker.Limiter("device1", false) + + report(tracker, "slowpoke", 5) + report(tracker, "device1", 2000) + + assert.Eventually(t, func() bool { + return limiter.GetRateWrite() == ThrottledWriteRate + }, 500*time.Millisecond, 5*time.Millisecond, + "device1 should be throttled well before the slow device's report returns") +} diff --git a/devicefilter/devicefilter.go b/devicefilter/devicefilter.go index 66932ccb..265bc302 100644 --- a/devicefilter/devicefilter.go +++ b/devicefilter/devicefilter.go @@ -17,6 +17,7 @@ import ( "github.com/getlantern/http-proxy-lantern/v2/blacklist" "github.com/getlantern/http-proxy-lantern/v2/common" + "github.com/getlantern/http-proxy-lantern/v2/datacap" "github.com/getlantern/http-proxy-lantern/v2/domains" "github.com/getlantern/http-proxy-lantern/v2/instrument" "github.com/getlantern/http-proxy-lantern/v2/redis" @@ -30,11 +31,52 @@ var ( epoch = time.Date(2016, 1, 1, 0, 0, 0, 0, time.UTC) alwaysThrottle = listeners.NewRateLimiter(10, 10) // this is basically unusably slow, only used for malicious or really old/broken clients - - defaultThrottleRate = int64(5000 * 1024 / 8) // 5 Mbps ) -// deviceFilterPre does the device-based filtering +// DefaultThrottleRate is the ceiling every non-pro device is held to even +// before it reaches its data cap, so that no one device monopolizes a proxy. +const DefaultThrottleRate = int64(5000 * 1024 / 8) // 5 Mbps + +// checkfallbacksDeviceID is the sentinel device ID checkfallbacks sends. +const checkfallbacksDeviceID = "~~~~~~" + +// throttleSentinelDevice applies the policy for the two sentinel device-ID +// values, shared by both accounting paths so they cannot drift: no device ID +// (old lantern versions and possible cracks — throttled to a near-unusable +// rate) and the checkfallbacks marker (never throttled). It reports whether +// the request was one of the two. +func throttleSentinelDevice(inst instrument.Instrument, wc listeners.WrapConn, req *http.Request, deviceID string) bool { + switch deviceID { + case "": + inst.Throttle(req.Context(), true, "no-device-id") + wc.ControlMessage("throttle", alwaysThrottle) + return true + case checkfallbacksDeviceID: + inst.Throttle(req.Context(), false, "checkfallbacks") + return true + } + return false +} + +// setXBQHeaders attaches the XBQ/XBQv2 usage headers flashlight's bandwidth +// package renders in the client UI. This is the single definition of that wire +// format for both accounting paths — a device must see identical headers no +// matter which path its proxy was provisioned with. +// +// XBQ is //; +// XBQv2 appends /. XBQ is kept for backward +// compatibility with older clients. +func setXBQHeaders(resp *http.Response, usedBytes, capBytes int64, asOf time.Time, ttlSeconds int64) { + if resp.Header == nil { + resp.Header = make(http.Header, 1) + } + xbq := fmt.Sprintf("%d/%d/%d", usedBytes/(1024*1024), capBytes/(1024*1024), int64(asOf.Sub(epoch).Seconds())) + resp.Header.Set(common.XBQHeader, xbq) + resp.Header.Set(common.XBQHeaderv2, fmt.Sprintf("%s/%d", xbq, ttlSeconds)) +} + +// deviceFilterPre does the device-based filtering, with usage fetched from the +// reporting Redis. Its datacap-sidecar counterpart is datacapFilterPre. type deviceFilterPre struct { deviceFetcher *redis.DeviceFetcher throttleConfig throttle.Config @@ -92,13 +134,13 @@ func (f *deviceFilterPre) Apply(cs *filters.ConnectionState, req *http.Request, // keep bandwidth hogs from using too much bandwidth. Note - this does not apply to pro proxies // which don't use the devicefilter at all. throttleDefault := func(message string) { - if defaultThrottleRate <= 0 { + if DefaultThrottleRate <= 0 { f.instrument.Throttle(req.Context(), false, message) } - limiter := f.rateLimiterForDevice(lanternDeviceID, defaultThrottleRate, defaultThrottleRate) + limiter := f.rateLimiterForDevice(lanternDeviceID, DefaultThrottleRate, DefaultThrottleRate) if log.IsTraceEnabled() { log.Tracef("Throttling connection to %v per second by default", - humanize.Bytes(uint64(defaultThrottleRate))) + humanize.Bytes(uint64(DefaultThrottleRate))) } f.instrument.Throttle(req.Context(), true, "default") wc.ControlMessage("throttle", limiter) @@ -111,16 +153,7 @@ func (f *deviceFilterPre) Apply(cs *filters.ConnectionState, req *http.Request, return next(cs, req) } - if lanternDeviceID == "" { - // Old lantern versions and possible cracks do not include the device - // ID. Just throttle them. - f.instrument.Throttle(req.Context(), true, "no-device-id") - wc.ControlMessage("throttle", alwaysThrottle) - return next(cs, req) - } - if lanternDeviceID == "~~~~~~" { - // This is checkfallbacks, don't throttle it - f.instrument.Throttle(req.Context(), false, "checkfallbacks") + if throttleSentinelDevice(f.instrument, wc, req, lanternDeviceID) { return next(cs, req) } @@ -159,7 +192,7 @@ func (f *deviceFilterPre) Apply(cs *filters.ConnectionState, req *http.Request, // per connection limiter // Note - when people hit the data cap, we only throttle writes back to the client, not reads. // This way, they can continue to upload videos or other bandwidth intensive content for sharing. - limiter := f.rateLimiterForDevice(lanternDeviceID, defaultThrottleRate, settings.Rate) + limiter := f.rateLimiterForDevice(lanternDeviceID, DefaultThrottleRate, settings.Rate) if log.IsTraceEnabled() { log.Tracef("Throttling connection from device %s to %v per second", lanternDeviceID, humanize.Bytes(uint64(settings.Rate))) @@ -180,14 +213,7 @@ func (f *deviceFilterPre) Apply(cs *filters.ConnectionState, req *http.Request, if !capOn || !f.sendXBQHeader { return resp, nextCtx, err } - if resp.Header == nil { - resp.Header = make(http.Header, 1) - } - uMiB := u.Bytes / (1024 * 1024) - xbq := fmt.Sprintf("%d/%d/%d", uMiB, settings.Threshold/(1024*1024), int64(u.AsOf.Sub(epoch).Seconds())) - xbqv2 := fmt.Sprintf("%s/%d", xbq, u.TTLSeconds) - resp.Header.Set(common.XBQHeader, xbq) // for backward compatibility with older clients - resp.Header.Set(common.XBQHeaderv2, xbqv2) // for new clients that support different bandwidth cap expirations + setXBQHeaders(resp, u.Bytes, settings.Threshold, u.AsOf, u.TTLSeconds) f.instrument.XBQHeaderSent(req.Context()) return resp, nextCtx, err } @@ -204,6 +230,101 @@ func (f *deviceFilterPre) rateLimiterForDevice(deviceID string, rateLimitRead, r return limiter } +// datacapFilterPre is deviceFilterPre's counterpart for proxies whose byte +// accounting runs through the local datacap sidecar. The throttle decision +// itself is made asynchronously by the tracker; the filter attaches the +// device's shared limiter and surfaces the tracker's latest view of the device +// to the client via the XBQ headers. +type datacapFilterPre struct { + tracker *datacap.Tracker + sendXBQHeader bool + instrument instrument.Instrument +} + +// NewDatacapPre creates the filter for proxies whose byte accounting runs +// through the local datacap sidecar. Unlike the Redis path, the limiter it +// attaches is shared across all of a device's connections and is re-rated by +// the tracker as reports come back, so crossing the cap slows down transfers +// that are already in flight rather than only the next one. +func NewDatacapPre(tracker *datacap.Tracker, sendXBQHeader bool, instrument instrument.Instrument) filters.Filter { + return &datacapFilterPre{ + tracker: tracker, + sendXBQHeader: sendXBQHeader, + instrument: instrument, + } +} + +func (f *datacapFilterPre) Apply(cs *filters.ConnectionState, req *http.Request, next filters.Next) (*http.Response, *filters.ConnectionState, error) { + if log.IsTraceEnabled() { + reqStr, _ := httputil.DumpRequest(req, true) + log.Tracef("DeviceFilter Middleware received request:\n%s", reqStr) + } + + wc := cs.Downstream().(listeners.WrapConn) + deviceID := req.Header.Get(common.DeviceIdHeader) + + // Some domains are excluded from being throttled. Their bytes still count + // towards the cap (accounting is per connection, not per request), they are + // just never held to the capped rate — hence a separate limiter that the + // tracker never re-rates. + // + // This check deliberately precedes the sentinel-device guards, matching the + // Redis path: a request to an excluded domain gets the default rate even + // with a missing device ID. Moving the guards first would newly subject + // old clients to alwaysThrottle on domains we have decided not to throttle. + // Such requests share one limiter under the empty device ID, exactly as + // rateLimiterForDevice("") does on the Redis path. + if domains.ConfigForRequest(req).Unthrottled { + f.instrument.Throttle(req.Context(), true, "default") + wc.ControlMessage("throttle", f.tracker.Limiter(deviceID, true)) + return next(cs, req) + } + + if throttleSentinelDevice(f.instrument, wc, req, deviceID) { + return next(cs, req) + } + + // The limiter is shared by every connection of this device and already + // carries whatever rate the last sidecar response set, so attaching it is + // the whole of enforcement here. + limiter, u, haveUsage := f.tracker.LimiterAndUsage(deviceID) + wc.ControlMessage("throttle", limiter) + + if !haveUsage { + // The sidecar only learns of a device from its first usage report, so + // there is nothing to report to the client yet. The limiter is at the + // default rate until then. + f.instrument.Throttle(req.Context(), true, "default") + return next(cs, req) + } + + reason := "default" + if u.Throttled { + reason = "datacap" + } + f.instrument.Throttle(req.Context(), true, reason) + + resp, nextCtx, err := next(cs, req) + if resp == nil || err != nil { + return resp, nextCtx, err + } + // A zero cap limit means this device's country/platform has no cap entry, + // which the client renders as "no data cap" — same meaning as a + // non-positive Redis threshold did. + if u.CapLimit <= 0 || !f.sendXBQHeader { + return resp, nextCtx, err + } + ttlSeconds := int64(0) + if !u.Expiry.IsZero() { + if remaining := time.Until(u.Expiry); remaining > 0 { + ttlSeconds = int64(remaining.Seconds()) + } + } + setXBQHeaders(resp, u.BytesUsed, u.CapLimit, u.AsOf, ttlSeconds) + f.instrument.XBQHeaderSent(req.Context()) + return resp, nextCtx, err +} + func NewPost(bl *blacklist.Blacklist) filters.Filter { return &deviceFilterPost{ bl: bl, diff --git a/http-proxy/main.go b/http-proxy/main.go index 70b823fa..b2f86bab 100644 --- a/http-proxy/main.go +++ b/http-proxy/main.go @@ -27,6 +27,7 @@ import ( proxy "github.com/getlantern/http-proxy-lantern/v2" "github.com/getlantern/http-proxy-lantern/v2/blacklist" + "github.com/getlantern/http-proxy-lantern/v2/datacap" "github.com/getlantern/http-proxy-lantern/v2/googlefilter" "github.com/getlantern/http-proxy-lantern/v2/obfs4listener" lanternredis "github.com/getlantern/http-proxy-lantern/v2/redis" @@ -121,6 +122,12 @@ var ( reportingRedisAddr = flag.String("reportingredis", "", "The address of the reporting Redis instance in \"redis[s]://host:port\" format") + // Successor to reportingredis. When both are set datacapurl wins: a proxy + // must account and enforce from exactly one source, and the sidecar is the + // one every other proxy flavor already reports to. + datacapURL = flag.String("datacapurl", "", "Base URL of the local datacap sidecar, e.g. \"http://127.0.0.1:8078\". Enables byte accounting and data-cap throttling through the sidecar, superseding reportingredis.") + datacapReportInterval = flag.Duration("datacapreportinterval", datacap.DefaultReportInterval, "How frequently to flush accumulated per-device usage to the datacap sidecar.") + // default value of tunnelPorts matches ports in flashlight/client/client.go tunnelPorts = flag.String("tunnelports", "80,443,22,110,995,143,993,8080,8443,5222,5223,5224,5228,5229,7300,19302,19303,19304,19305,19306,19307,19308,19309", "Comma seperated list of ports allowed for HTTP CONNECT tunnel. Allow all ports if empty.") tos = flag.Int("tos", 0, "Specify a diffserv TOS to prioritize traffic. Defaults to 0 (off)") @@ -391,13 +398,16 @@ func main() { go periodicallyForceGC() var reportingRedisClient *redis.Client - if *reportingRedisAddr != "" { + switch { + case *datacapURL != "": + log.Debugf("reporting bandwidth to the datacap sidecar at %v", *datacapURL) + case *reportingRedisAddr != "": reportingRedisClient, err = lanternredis.NewClient(*reportingRedisAddr) if err != nil { log.Errorf("failed to initialize redis client, will not be able to perform bandwidth limiting: %v", err) } - } else { - log.Debug("no redis address configured for bandwidth reporting") + default: + log.Debug("neither a datacap sidecar nor a redis address configured for bandwidth reporting") } p := &proxy.Proxy{ @@ -426,6 +436,8 @@ func main() { ProxiedSitesSamplePercentage: *proxiedSitesSamplePercentage, ProxiedSitesTrackingID: *proxiedSitesTrackingId, ReportingRedisClient: reportingRedisClient, + DatacapURL: *datacapURL, + DatacapReportInterval: *datacapReportInterval, Token: *token, TunnelPorts: *tunnelPorts, Obfs4Addr: *obfs4Addr, diff --git a/http_proxy.go b/http_proxy.go index aed56d88..7e52d32b 100644 --- a/http_proxy.go +++ b/http_proxy.go @@ -31,6 +31,7 @@ import ( "github.com/getlantern/http-proxy-lantern/v2/banditcallback" "github.com/getlantern/http-proxy-lantern/v2/broflake" "github.com/getlantern/http-proxy-lantern/v2/common" + "github.com/getlantern/http-proxy-lantern/v2/datacap" "github.com/getlantern/http-proxy-lantern/v2/opsfilter" "github.com/getlantern/http-proxy-lantern/v2/otel" "github.com/getlantern/http-proxy-lantern/v2/shadowsocks" @@ -213,7 +214,14 @@ type Proxy struct { VMessAddr string VMessUUIDs []string + // DatacapURL is the base URL of the local datacap sidecar. When set, byte + // accounting and data-cap throttling run through the sidecar and the + // reporting-Redis path is left unused. + DatacapURL string + DatacapReportInterval time.Duration + throttleConfig throttle.Config + datacapTracker *datacap.Tracker instrument instrument.Instrument } @@ -255,6 +263,7 @@ func (p *Proxy) ListenAndServe(ctx context.Context) error { } p.setBenchmarkMode() p.loadThrottleConfig() + p.loadDatacapTracker() if p.ENHTTPAddr != "" { return p.ListenAndServeENHTTP() @@ -545,7 +554,7 @@ func (p *Proxy) createFilterChain(bl *blacklist.Blacklist) (filters.Chain, proxy // unauthenticated traffic there — fine because bench mode is a // local-only test setup), and before devicefilter so the // emitter still runs for pro tracks (devicefilter is gated on - // ReportingRedisClient, which pro proxies don't set, but the + // the accounting source, which pro proxies aren't given, but the // bandit still wants signal for pro arms). OnFirstOnly because // the device-id header only needs to be read once per // connection — matches the other auth-adjacent filters. @@ -555,13 +564,18 @@ func (p *Proxy) createFilterChain(bl *blacklist.Blacklist) (filters.Chain, proxy ) } - if p.ReportingRedisClient == nil { - log.Debug("Not enabling bandwidth limiting") - } else { + switch { + case p.datacapTracker != nil: + filterChain = filterChain.Append( + proxy.OnFirstOnly(devicefilter.NewDatacapPre(p.datacapTracker, !p.Pro, p.instrument)), + ) + case p.ReportingRedisClient != nil: filterChain = filterChain.Append( proxy.OnFirstOnly(devicefilter.NewPre( redis.NewDeviceFetcher(p.ReportingRedisClient), p.throttleConfig, !p.Pro, p.instrument)), ) + default: + log.Debug("Not enabling bandwidth limiting") } filterChain = filterChain.Append( @@ -722,7 +736,7 @@ func (p *Proxy) buildOTELOpts(includeProxyName bool) *otel.Opts { } func (p *Proxy) configureBandwidthReporting() *reportingConfig { - return newReportingConfig(p.CountryLookup, p.ReportingRedisClient, p.instrument, p.throttleConfig) + return newReportingConfig(p.CountryLookup, p.ReportingRedisClient, p.instrument, p.throttleConfig, p.datacapTracker) } func (p *Proxy) loadThrottleConfig() { @@ -734,6 +748,22 @@ func (p *Proxy) loadThrottleConfig() { } } +// loadDatacapTracker starts the sidecar-backed accounting pipeline. Pro tracks +// are gated server-side — the provisioner simply omits DatacapURL from their +// config — so an unset URL is the normal case there, not a misconfiguration. +func (p *Proxy) loadDatacapTracker() { + if p.Pro || p.DatacapURL == "" { + return + } + p.datacapTracker = datacap.NewTracker(datacap.TrackerOpts{ + Client: datacap.NewClient(p.DatacapURL, datacap.DefaultHTTPTimeout), + CountryLookup: p.CountryLookup, + DefaultRate: devicefilter.DefaultThrottleRate, + ReportInterval: p.DatacapReportInterval, + }) + log.Debugf("Reporting bandwidth usage to the datacap sidecar at %v", p.DatacapURL) +} + func (p *Proxy) legacyAPIHostExceptions() []string { if p.LegacyAPIHosts == "" { return nil diff --git a/listeners/bitrate.go b/listeners/bitrate.go index c78ddec0..5ea12642 100644 --- a/listeners/bitrate.go +++ b/listeners/bitrate.go @@ -3,6 +3,7 @@ package listeners import ( "net" "net/http" + "sync/atomic" "time" "github.com/getlantern/ratelimit" @@ -12,45 +13,73 @@ const ( minSleep = 5 * time.Millisecond // don't bother sleeping for less than this amount of time ) -type RateLimiter struct { +// rateBuckets pairs the token buckets with the rates they were built from so a +// re-rate swaps both together and a reader never sees a bucket that disagrees +// with its rate. +type rateBuckets struct { r *ratelimit.Bucket w *ratelimit.Bucket rateRead int64 rateWrite int64 } -func NewRateLimiter(rateRead, rateWrite int64) *RateLimiter { - l := &RateLimiter{ - rateRead: rateRead, - rateWrite: rateWrite, - } +func newRateBuckets(rateRead, rateWrite int64) *rateBuckets { + b := &rateBuckets{rateRead: rateRead, rateWrite: rateWrite} if rateRead > 0 { - l.r = ratelimit.NewBucketWithRate(float64(rateRead), rateRead) + b.r = ratelimit.NewBucketWithRate(float64(rateRead), rateRead) } if rateWrite > 0 { - l.w = ratelimit.NewBucketWithRate(float64(rateWrite), rateWrite) + b.w = ratelimit.NewBucketWithRate(float64(rateWrite), rateWrite) } + return b +} + +// RateLimiter caps read and write throughput on the connections it is attached +// to. Its rates are mutable (see SetRates): one limiter shared by every +// connection of a device can be re-rated in place, and the new rate applies to +// connections that are already open rather than only to the next one. +type RateLimiter struct { + buckets atomic.Pointer[rateBuckets] +} + +func NewRateLimiter(rateRead, rateWrite int64) *RateLimiter { + l := &RateLimiter{} + l.buckets.Store(newRateBuckets(rateRead, rateWrite)) return l } +// SetRates re-rates the limiter. It is a no-op when the rates are unchanged, so +// a caller refreshing on every reporting cycle does not continually reset the +// token buckets. +func (l *RateLimiter) SetRates(rateRead, rateWrite int64) { + if cur := l.buckets.Load(); cur.rateRead == rateRead && cur.rateWrite == rateWrite { + return + } + l.buckets.Store(newRateBuckets(rateRead, rateWrite)) +} + func (l *RateLimiter) GetRateRead() int64 { - return l.rateRead + return l.buckets.Load().rateRead } func (l *RateLimiter) GetRateWrite() int64 { - return l.rateWrite + return l.buckets.Load().rateWrite } -func (l *RateLimiter) waitRead(n int) { - d := l.wait(l.r, n) - if d > 0 { +func (b *rateBuckets) waitRead(n int) { + if b.r == nil { + return + } + if d := b.r.Take(int64(n)); d > 0 { sleep(d) } } -func (l *RateLimiter) waitWrite(n int) { - d := l.wait(l.w, n) - if d > 0 { +func (b *rateBuckets) waitWrite(n int) { + if b.w == nil { + return + } + if d := b.w.Take(int64(n)); d > 0 { sleep(d) } } @@ -64,10 +93,6 @@ func sleep(d time.Duration) { time.Sleep(d) } -func (l *RateLimiter) wait(b *ratelimit.Bucket, n int) time.Duration { - return b.Take(int64(n)) -} - type bitrateListener struct { net.Listener } @@ -76,6 +101,12 @@ func NewBitrateListener(l net.Listener) net.Listener { return &bitrateListener{l} } +// unlimited is the limiter every conn starts with until a filter attaches a +// real one. A (0,0) limiter never waits and nothing ever re-rates it — +// ControlMessage replaces the pointer wholesale — so one shared instance +// serves every conn instead of allocating two objects per accept. +var unlimited = NewRateLimiter(0, 0) + func (bl *bitrateListener) Accept() (net.Conn, error) { c, err := bl.Listener.Accept() if err != nil { @@ -83,40 +114,43 @@ func (bl *bitrateListener) Accept() (net.Conn, error) { } wc, _ := c.(WrapConnEmbeddable) - return &bitrateConn{ + brc := &bitrateConn{ WrapConnEmbeddable: wc, Conn: c, - limiter: NewRateLimiter(0, 0), - }, err + } + brc.limiter.Store(unlimited) + return brc, err } // Bitrate Conn wrapper type bitrateConn struct { WrapConnEmbeddable net.Conn - limiter *RateLimiter + limiter atomic.Pointer[RateLimiter] } func (c *bitrateConn) Read(p []byte) (n int, err error) { - if c.limiter.rateRead == 0 { + b := c.limiter.Load().buckets.Load() + if b.rateRead == 0 { return c.Conn.Read(p) } n, err = c.Conn.Read(p) if err == nil { - c.limiter.waitRead(n) + b.waitRead(n) } return } func (c *bitrateConn) Write(p []byte) (n int, err error) { - if c.limiter.rateWrite == 0 { + b := c.limiter.Load().buckets.Load() + if b.rateWrite == 0 { return c.Conn.Write(p) } n, err = c.Conn.Write(p) if err == nil { - c.limiter.waitWrite(n) + b.waitWrite(n) } return } @@ -131,7 +165,7 @@ func (c *bitrateConn) OnState(s http.ConnState) { func (c *bitrateConn) ControlMessage(msgType string, data interface{}) { // per user message always overrides the active flag if msgType == "throttle" { - c.limiter = data.(*RateLimiter) + c.limiter.Store(data.(*RateLimiter)) } if c.WrapConnEmbeddable != nil { diff --git a/listeners/bitrate_test.go b/listeners/bitrate_test.go index 5698b0ae..80932153 100644 --- a/listeners/bitrate_test.go +++ b/listeners/bitrate_test.go @@ -186,3 +186,71 @@ func BenchmarkThrottledReader(b *testing.B) { conn.Write(benchBuf) } } + +// A limiter that is re-rated must take effect on connections it is already +// attached to: that is what lets a device be slowed the moment it crosses its +// data cap instead of at its next connection. +func TestReRatingAppliesToAnOpenConn(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("Error creating listener: %v", err) + } + defer ln.Close() + + type accept struct { + conn net.Conn + err error + } + accepted := make(chan accept, 1) + bl := NewBitrateListener(ln) + go func() { + c, err := bl.Accept() + accepted <- accept{c, err} + }() + + client, err := net.Dial("tcp", ln.Addr().String()) + if err != nil { + t.Fatalf("Error connecting to local server: %v", err) + } + defer client.Close() + go io.Copy(io.Discard, client) + + res := <-accepted + if res.err != nil { + t.Fatalf("Error accepting: %v", res.err) + } + defer res.conn.Close() + + limiter := NewRateLimiter(0, 0) + res.conn.(*bitrateConn).ControlMessage("throttle", limiter) + + payload := make([]byte, 4096) + start := time.Now() + if _, err := res.conn.Write(payload); err != nil { + t.Fatalf("Error writing: %v", err) + } + assert.Less(t, time.Since(start), 100*time.Millisecond, "an unrated limiter should not delay writes") + + // Re-rate in place — the conn is never handed a new limiter. + limiter.SetRates(0, 8192) + start = time.Now() + for i := 0; i < 4; i++ { + if _, err := res.conn.Write(payload); err != nil { + t.Fatalf("Error writing: %v", err) + } + } + assert.Greater(t, time.Since(start), 500*time.Millisecond, "the open conn should be held to the new rate") +} + +func TestSetRatesKeepsBucketsWhenUnchanged(t *testing.T) { + l := NewRateLimiter(1000, 1000) + before := l.buckets.Load() + + l.SetRates(1000, 1000) + assert.Same(t, before, l.buckets.Load(), "an unchanged rate should not reset the token buckets") + + l.SetRates(1000, 500) + assert.NotSame(t, before, l.buckets.Load()) + assert.Equal(t, int64(500), l.GetRateWrite()) + assert.Equal(t, int64(1000), l.GetRateRead()) +} diff --git a/reporting.go b/reporting.go index bbffafc6..53f98c53 100644 --- a/reporting.go +++ b/reporting.go @@ -10,6 +10,7 @@ import ( "github.com/getlantern/geo" "github.com/getlantern/http-proxy-lantern/v2/common" + "github.com/getlantern/http-proxy-lantern/v2/datacap" "github.com/getlantern/http-proxy-lantern/v2/listeners" "github.com/getlantern/measured" @@ -27,7 +28,7 @@ type reportingConfig struct { wrapper func(ls net.Listener) net.Listener } -func newReportingConfig(countryLookup geo.CountryLookup, rc *rclient.Client, instrument instrument.Instrument, throttleConfig throttle.Config) *reportingConfig { +func newReportingConfig(countryLookup geo.CountryLookup, rc *rclient.Client, instrument instrument.Instrument, throttleConfig throttle.Config, datacapTracker *datacap.Tracker) *reportingConfig { proxiedBytesReporter := func(ctx map[string]interface{}, stats *measured.Stats, deltaStats *measured.Stats, final bool) { noDelta := deltaStats.SentTotal == 0 && deltaStats.RecvTotal == 0 if noDelta && !final { @@ -75,13 +76,16 @@ func newReportingConfig(countryLookup geo.CountryLookup, rc *rclient.Client, ins } var reporter listeners.MeasuredReportFN - if throttleConfig == nil { + switch { + case datacapTracker != nil: + reporter = datacapTracker.Reporter() + case throttleConfig == nil: log.Debug("No throttling configured, don't bother reporting bandwidth usage to Redis") reporter = func(ctx map[string]interface{}, stats *measured.Stats, deltaStats *measured.Stats, final bool) { // noop } - } else if rc != nil { + case rc != nil: reporter = redis.NewMeasuredReporter(countryLookup, rc, measuredReportingInterval, throttleConfig) } reporter = combineReporter(reporter, proxiedBytesReporter)