diff --git a/internal/app/app.go b/internal/app/app.go index a2c1cb46..dcd963d1 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -268,6 +268,8 @@ func run(c *cli.Context) error { logrus.Warn("The vault is corrupt and has been wiped") b.PushError(bridge.ErrVaultCorrupt) } + // Start telemetry heartbeat process + b.StartHeartbeat(b) // Run the frontend. return runFrontend(c, crashHandler, restarter, locations, b, eventCh, quitCh, c.Int(flagParentPID)) diff --git a/internal/bridge/bridge.go b/internal/bridge/bridge.go index 2d3ebb16..4af1571f 100644 --- a/internal/bridge/bridge.go +++ b/internal/bridge/bridge.go @@ -44,7 +44,6 @@ import ( "github.com/ProtonMail/proton-bridge/v3/internal/telemetry" "github.com/ProtonMail/proton-bridge/v3/internal/user" "github.com/ProtonMail/proton-bridge/v3/internal/vault" - "github.com/ProtonMail/proton-bridge/v3/pkg/keychain" "github.com/bradenaw/juniper/xslices" "github.com/emersion/go-smtp" "github.com/go-resty/resty/v2" @@ -307,7 +306,6 @@ func newBridge( } bridge.smtpServer = newSMTPServer(bridge, tlsConfig, logSMTP) - bridge.heartbeat = telemetry.NewHeartbeat(bridge, 1143, 1025, gluonCacheDir, keychain.DefaultHelper) return bridge, nil } @@ -429,8 +427,6 @@ func (bridge *Bridge) init(tlsReporter TLSReporter) error { }) }) - // init telemetry - bridge.initHeartbeat() return nil } diff --git a/internal/bridge/heartbeat.go b/internal/bridge/heartbeat.go index da75f3a6..4fc3afcd 100644 --- a/internal/bridge/heartbeat.go +++ b/internal/bridge/heartbeat.go @@ -26,6 +26,7 @@ import ( "github.com/ProtonMail/proton-bridge/v3/internal/safe" "github.com/ProtonMail/proton-bridge/v3/internal/telemetry" "github.com/ProtonMail/proton-bridge/v3/internal/vault" + "github.com/ProtonMail/proton-bridge/v3/pkg/keychain" "github.com/sirupsen/logrus" ) @@ -76,7 +77,9 @@ func (bridge *Bridge) SetLastHeartbeatSent(timestamp time.Time) error { return bridge.vault.SetLastHeartbeatSent(timestamp) } -func (bridge *Bridge) initHeartbeat() { +func (bridge *Bridge) StartHeartbeat(manager telemetry.HeartbeatManager) { + bridge.heartbeat = telemetry.NewHeartbeat(manager, 1143, 1025, bridge.GetGluonCacheDir(), keychain.DefaultHelper) + safe.RLock(func() { var splitMode = false for _, user := range bridge.users { @@ -102,8 +105,10 @@ func (bridge *Bridge) initHeartbeat() { bridge.heartbeat.SetCacheLocation(bridge.GetGluonCacheDir()) if val, err := bridge.GetKeychainApp(); err != nil { bridge.heartbeat.SetKeyChainPref(val) + } else { + bridge.heartbeat.SetKeyChainPref(keychain.DefaultHelper) } bridge.heartbeat.SetPrevVersion(bridge.GetLastVersion().String()) - bridge.heartbeat.StartSending() + bridge.heartbeat.TrySending() } diff --git a/internal/bridge/user.go b/internal/bridge/user.go index 4bb2a455..2c34a878 100644 --- a/internal/bridge/user.go +++ b/internal/bridge/user.go @@ -571,6 +571,9 @@ func (bridge *Bridge) addUserWithVault( bridge.heartbeat.SetNbAccount(len(bridge.users)) }, bridge.usersLock) + // As we need at least one user to send heartbeat, try to send it. + bridge.heartbeat.TrySending() + return nil } diff --git a/internal/telemetry/heartbeat.go b/internal/telemetry/heartbeat.go index 6b6c7458..8d8c8f2d 100644 --- a/internal/telemetry/heartbeat.go +++ b/internal/telemetry/heartbeat.go @@ -148,7 +148,7 @@ func (heartbeat *Heartbeat) SetPrevVersion(val string) { heartbeat.metrics.Dimensions.PrevVersion = val } -func (heartbeat *Heartbeat) StartSending() { +func (heartbeat *Heartbeat) TrySending() { if heartbeat.manager.IsTelemetryAvailable() { lastSent := heartbeat.manager.GetLastHeartbeatSent() now := time.Now() diff --git a/internal/telemetry/heartbeat_test.go b/internal/telemetry/heartbeat_test.go index b08931c0..7c8b4fb1 100644 --- a/internal/telemetry/heartbeat_test.go +++ b/internal/telemetry/heartbeat_test.go @@ -57,7 +57,7 @@ func TestHeartbeat_default_heartbeat(t *testing.T) { mock.EXPECT().SendHeartbeat(&data).Return(true) mock.EXPECT().SetLastHeartbeatSent(gomock.Any()).Return(nil) - hb.StartSending() + hb.TrySending() }) } @@ -66,7 +66,7 @@ func TestHeartbeat_already_sent_heartbeat(t *testing.T) { mock.EXPECT().IsTelemetryAvailable().Return(true) mock.EXPECT().GetLastHeartbeatSent().Return(time.Now().Truncate(24 * time.Hour)) - hb.StartSending() + hb.TrySending() }) } diff --git a/tests/bdd_test.go b/tests/bdd_test.go index da93420f..e560b7a8 100644 --- a/tests/bdd_test.go +++ b/tests/bdd_test.go @@ -231,6 +231,12 @@ func TestFeatures(testingT *testing.T) { ctx.Step(`^SMTP client "([^"]*)" sends RSET$`, s.smtpClientSendsReset) ctx.Step(`^SMTP client "([^"]*)" sends the following message from "([^"]*)" to "([^"]*)":$`, s.smtpClientSendsTheFollowingMessageFromTo) ctx.Step(`^SMTP client "([^"]*)" logs out$`, s.smtpClientLogsOut) + + // ==== TELEMETRY ==== + ctx.Step(`^bridge eventually sends the following heartbeat:$`, s.bridgeEventuallySendsTheFollowingHeartbeat) + ctx.Step(`^bridge needs to send heartbeat$`, s.bridgeNeedsToSendHeartbeat) + ctx.Step(`^bridge do not need to send heartbeat$`, s.bridgeDoNotNeedToSendHeartbeat) + ctx.Step(`^heartbeat is not whitelisted$`, s.heartbeatIsNotwhitelisted) }, Options: &godog.Options{ Format: "pretty", diff --git a/tests/ctx_bridge_test.go b/tests/ctx_bridge_test.go index 92cdf0e1..de4ed616 100644 --- a/tests/ctx_bridge_test.go +++ b/tests/ctx_bridge_test.go @@ -174,6 +174,8 @@ func (t *testCtx) initBridge() (<-chan events.Event, error) { } t.bridge = bridge + t.heartbeat.setBridge(bridge) + bridge.StartHeartbeat(t.heartbeat) return t.events.collectFrom(eventCh), nil } diff --git a/tests/ctx_heartbeat_test.go b/tests/ctx_heartbeat_test.go new file mode 100644 index 00000000..fc543cff --- /dev/null +++ b/tests/ctx_heartbeat_test.go @@ -0,0 +1,72 @@ +package tests + +import ( + "errors" + "testing" + "time" + + "github.com/ProtonMail/proton-bridge/v3/internal/bridge" + "github.com/ProtonMail/proton-bridge/v3/internal/telemetry" + "github.com/stretchr/testify/assert" +) + +type heartbeatRecorder struct { + heartbeat telemetry.HeartbeatData + bridge *bridge.Bridge + reject bool + assert *assert.Assertions +} + +func newHeartbeatRecorder(tb testing.TB) *heartbeatRecorder { + return &heartbeatRecorder{ + heartbeat: telemetry.HeartbeatData{}, + bridge: nil, + reject: false, + assert: assert.New(tb), + } +} + +func (hb *heartbeatRecorder) setBridge(bridge *bridge.Bridge) { + hb.bridge = bridge +} + +func (hb *heartbeatRecorder) GetLastHeartbeatSent() time.Time { + if hb.bridge == nil { + return time.Now() + } + return hb.bridge.GetLastHeartbeatSent() +} + +func (hb *heartbeatRecorder) IsTelemetryAvailable() bool { + if hb.bridge == nil { + return false + } + return hb.bridge.IsTelemetryAvailable() +} + +func (hb *heartbeatRecorder) SendHeartbeat(metrics *telemetry.HeartbeatData) bool { + if hb.bridge == nil { + return false + } + + if len(hb.bridge.GetUserIDs()) == 0 { + return false + } + + if hb.reject { + return false + } + hb.heartbeat = *metrics + return true +} + +func (hb *heartbeatRecorder) SetLastHeartbeatSent(timestamp time.Time) error { + if hb.bridge == nil { + return errors.New("no bridge initialized") + } + return hb.bridge.SetLastHeartbeatSent(timestamp) +} + +func (hb *heartbeatRecorder) rejectSend() { + hb.reject = true +} diff --git a/tests/ctx_test.go b/tests/ctx_test.go index 004ab501..369e865f 100644 --- a/tests/ctx_test.go +++ b/tests/ctx_test.go @@ -122,15 +122,16 @@ func newTestAddr(addrID, email string) *testAddr { type testCtx struct { // These are the objects supporting the test. - dir string - api API - netCtl *proton.NetCtl - locator *locations.Locations - storeKey []byte - version *semver.Version - mocks *bridge.Mocks - events *eventCollector - reporter *reportRecorder + dir string + api API + netCtl *proton.NetCtl + locator *locations.Locations + storeKey []byte + version *semver.Version + mocks *bridge.Mocks + events *eventCollector + reporter *reportRecorder + heartbeat *heartbeatRecorder // bridge holds the bridge app under test. bridge *bridge.Bridge @@ -180,15 +181,16 @@ func newTestCtx(tb testing.TB) *testCtx { dir := tb.TempDir() t := &testCtx{ - dir: dir, - api: newTestAPI(), - netCtl: proton.NewNetCtl(), - locator: locations.New(bridge.NewTestLocationsProvider(dir), "config-name"), - storeKey: []byte("super-secret-store-key"), - version: defaultVersion, - mocks: bridge.NewMocks(tb, defaultVersion, defaultVersion), - events: newEventCollector(), - reporter: newReportRecorder(tb), + dir: dir, + api: newTestAPI(), + netCtl: proton.NewNetCtl(), + locator: locations.New(bridge.NewTestLocationsProvider(dir), "config-name"), + storeKey: []byte("super-secret-store-key"), + version: defaultVersion, + mocks: bridge.NewMocks(tb, defaultVersion, defaultVersion), + events: newEventCollector(), + reporter: newReportRecorder(tb), + heartbeat: newHeartbeatRecorder(tb), userByID: make(map[string]*testUser), userUUIDByName: make(map[string]string), diff --git a/tests/features/bridge/heartbeat.feature b/tests/features/bridge/heartbeat.feature new file mode 100644 index 00000000..a1ce7372 --- /dev/null +++ b/tests/features/bridge/heartbeat.feature @@ -0,0 +1,47 @@ +Feature: Send Telemetry Heartbeat + Background: + Given there exists an account with username "[user:user1]" and password "password" + And bridge starts + + + Scenario: Send at first start - one user + Then bridge telemetry feature is enabled + And bridge needs to send heartbeat + When the user logs in with username "[user:user1]" and password "password" + And user "[user:user1]" finishes syncing + Then bridge eventually sends the following heartbeat: + """ + { + "MeasurementGroup": "bridge.any.usage", + "Event": "bridge_heartbeat", + "Values": { + "nb_account": 1 + }, + "Dimensions": { + "auto_update": "on", + "auto_start": "on", + "beta": "off", + "doh": "off", + "split_mode": "off", + "show_all_mail": "on", + "imap_connection_mode": "starttls", + "smtp_connection_mode": "starttls", + "imap_port": "default", + "smtp_port": "default", + "cache_location": "default", + "keychain_pref": "default", + "prev_version": "0.0.0", + "rollout": 42 + } + } + """ + And bridge do not need to send heartbeat + + Scenario: GroupMeasurement rejected by API + Given heartbeat is not whitelisted + Then bridge telemetry feature is enabled + And bridge needs to send heartbeat + When the user logs in with username "[user:user1]" and password "password" + And user "[user:user1]" finishes syncing + Then bridge needs to send heartbeat + diff --git a/tests/heartbeat_test.go b/tests/heartbeat_test.go new file mode 100644 index 00000000..075f7112 --- /dev/null +++ b/tests/heartbeat_test.go @@ -0,0 +1,70 @@ +package tests + +import ( + "encoding/json" + "errors" + "fmt" + "time" + + "github.com/ProtonMail/proton-bridge/v3/internal/telemetry" + "github.com/cucumber/godog" + "github.com/sirupsen/logrus" +) + +func (s *scenario) bridgeEventuallySendsTheFollowingHeartbeat(text *godog.DocString) error { + return eventually(func() error { + err := s.bridgeSendsTheFollowingHeartbeat(text) + logrus.WithError(err).Trace("Matching eventually") + return err + }) +} + +func (s *scenario) bridgeSendsTheFollowingHeartbeat(text *godog.DocString) error { + var wantHeartbeat telemetry.HeartbeatData + err := json.Unmarshal([]byte(text.Content), &wantHeartbeat) + if err != nil { + return err + } + + return matchHeartbeat(s.t.heartbeat.heartbeat, wantHeartbeat) +} + +func (s *scenario) bridgeNeedsToSendHeartbeat() error { + last := s.t.heartbeat.GetLastHeartbeatSent() + if !isAnotherDay(last, time.Now()) { + return fmt.Errorf("heartbeat already sent at %s", last) + } + return nil +} + +func (s *scenario) bridgeDoNotNeedToSendHeartbeat() error { + last := s.t.heartbeat.GetLastHeartbeatSent() + if isAnotherDay(last, time.Now()) { + return fmt.Errorf("heartbeat needs to be sent - last %s", last) + } + return nil +} + +func (s *scenario) heartbeatIsNotwhitelisted() error { + s.t.heartbeat.rejectSend() + return nil +} + +func matchHeartbeat(have, want telemetry.HeartbeatData) error { + if have == (telemetry.HeartbeatData{}) { + return errors.New("no heartbeat send (yet)") + } + + // Ignore rollout number + want.Dimensions.Rollout = have.Dimensions.Rollout + + if have != want { + return fmt.Errorf("missing heartbeat: have %#v, want %#v", have, want) + } + + return nil +} + +func isAnotherDay(last, now time.Time) bool { + return now.Year() > last.Year() || (now.Year() == last.Year() && now.YearDay() > last.YearDay()) +}