Silverfish/proton-bridge
1
0
Fork 1
mirror of https://github.com/ProtonMail/proton-bridge.git synced 2025-12-10 12:46:46 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat(BRIDGE-266): heartbeat telemetry update; extra integration tests;

Browse Source
This commit is contained in:
Atanas Janeshliev
2024-11-22 14:09:48 +00:00
parent cdcdd45bcf
commit 2e98d64f94
17 changed files with 411 additions and 256 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
internal/bridge/heartbeat.go
View File

@ -73,15 +73,15 @@ func (h *heartBeatState) init(bridge *Bridge, manager telemetry.HeartbeatManager
for _, user := range bridge.users {
if user.GetAddressMode() == vault.SplitMode {
splitMode = true
break
}
h.SetUserPlan(user.GetUserPlanName())
}
var nbAccount = len(bridge.users)
h.SetNbAccount(nbAccount)
var numberConnectedAccounts = len(bridge.users)
h.SetNumberConnectedAccounts(numberConnectedAccounts)
h.SetSplitMode(splitMode)
// Do not try to send if there is no user yet.
if nbAccount > 0 {
if numberConnectedAccounts > 0 {
defer h.start()
}
}, bridge.usersLock)

7
internal/bridge/identifier.go
View File

@ -17,7 +17,9 @@
package bridge
import "github.com/sirupsen/logrus"
import (
"github.com/sirupsen/logrus"
)
func (bridge *Bridge) GetCurrentUserAgent() string {
return bridge.identifier.GetUserAgent()
@ -30,6 +32,8 @@ func (bridge *Bridge) SetCurrentPlatform(platform string) {
func (bridge *Bridge) setUserAgent(name, version string) {
currentUserAgent := bridge.identifier.GetClientString()
bridge.heartbeat.SetContactedByAppleNotes(name)
bridge.identifier.SetClient(name, version)
newUserAgent := bridge.identifier.GetClientString()
@ -54,6 +58,7 @@ func (b *bridgeUserAgentUpdater) HasClient() bool {
}
func (b *bridgeUserAgentUpdater) SetClient(name, version string) {
b.heartbeat.SetContactedByAppleNotes(name)
b.identifier.SetClient(name, version)
}

7
internal/bridge/user.go
View File

@ -583,9 +583,12 @@ func (bridge *Bridge) addUserWithVault(
// Finally, save the user in the bridge.
safe.Lock(func() {
bridge.users[apiUser.ID] = user
bridge.heartbeat.SetNbAccount(len(bridge.users))
bridge.heartbeat.SetNumberConnectedAccounts(len(bridge.users))
}, bridge.usersLock)
// Set user plan if its of a higher rank.
bridge.heartbeat.SetUserPlan(user.GetUserPlanName())
// As we need at least one user to send heartbeat, try to send it.
bridge.heartbeat.start()
@ -618,7 +621,7 @@ func (bridge *Bridge) logoutUser(ctx context.Context, user *user.User, withAPI,
logUser.WithError(err).Error("Failed to logout user")
}
bridge.heartbeat.SetNbAccount(len(bridge.users))
bridge.heartbeat.SetNumberConnectedAccounts(len(bridge.users) - 1)
user.Close()
}

87
internal/plan/plan.go Normal file
View File

@ -0,0 +1,87 @@
// Copyright (c) 2024 Proton AG
//
// This file is part of Proton Mail Bridge.
//
// Proton Mail Bridge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// Proton Mail Bridge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with Proton Mail Bridge. If not, see <https://www.gnu.org/licenses/>.
package plan
import "strings"
const (
Unknown = "unknown"
Other = "other"
Business = "business"
Individual = "individual"
Group = "group"
)
var planHierarchy = map[string]int{ //nolint:gochecknoglobals
Business: 4,
Group: 3,
Individual: 2,
Other: 1,
Unknown: 0,
}
func IsHigherPriority(currentPlan, newPlan string) bool {
newRank, ok := planHierarchy[newPlan]
if !ok {
return false
}
currentRank, ok2 := planHierarchy[currentPlan]
if !ok2 {
return true // we don't have a valid plan, might as well replace it
}
return newRank > currentRank
}
func MapUserPlan(planName string) string {
if planName == "" {
return Unknown
}
switch strings.TrimSpace(strings.ToLower(planName)) {
case Individual:
return Individual
case Unknown:
return Unknown
case Business:
return Business
case Group:
return Group
case "mail2022":
return Individual
case "bundle2022":
return Individual
case "family2022":
return Group
case "visionary2022":
return Group
case "mailpro2022":
return Business
case "planbiz2024":
return Business
case "bundlepro2022":
return Business
case "bundlepro2024":
return Business
case "duo2024":
return Group
default:
return Other
}
}

3
internal/services/observability/distinction_utility.go
View File

@ -24,6 +24,7 @@ import (
"github.com/ProtonMail/gluon/async"
"github.com/ProtonMail/go-proton-api"
"github.com/ProtonMail/proton-bridge/v3/internal/plan"
"github.com/ProtonMail/proton-bridge/v3/internal/updater"
)
@ -62,7 +63,7 @@ func newDistinctionUtility(ctx context.Context, panicHandler async.PanicHandler,
observabilitySender: observabilitySender,
userPlanUnsafe: planUnknown,
userPlanUnsafe: plan.Unknown,
heartbeatData: heartbeatData{},
heartbeatTicker: time.NewTicker(updateInterval),

14
internal/services/observability/heartbeat.go
View File

@ -18,7 +18,7 @@
package observability
import (
"fmt"
"strconv"
"time"
"github.com/ProtonMail/gluon/async"
@ -87,10 +87,6 @@ func (d *distinctionUtility) sendHeartbeat() {
})
}
func formatBool(value bool) string {
return fmt.Sprintf("%t", value)
}
// generateHeartbeatUserMetric creates the heartbeat user metric and includes the relevant data.
func (d *distinctionUtility) generateHeartbeatUserMetric() proton.ObservabilityMetric {
return generateHeartbeatMetric(
@ -98,10 +94,10 @@ func (d *distinctionUtility) generateHeartbeatUserMetric() proton.ObservabilityM
d.getEmailClientUserAgent(),
getEnabled(d.settingsGetter.GetProxyAllowed()),
getEnabled(d.getBetaAccessEnabled()),
formatBool(d.heartbeatData.receivedOtherError),
formatBool(d.heartbeatData.receivedSyncError),
formatBool(d.heartbeatData.receivedEventLoopError),
formatBool(d.heartbeatData.receivedGluonError),
strconv.FormatBool(d.heartbeatData.receivedOtherError),
strconv.FormatBool(d.heartbeatData.receivedSyncError),
strconv.FormatBool(d.heartbeatData.receivedEventLoopError),
strconv.FormatBool(d.heartbeatData.receivedGluonError),
)
}

85
internal/services/observability/plan_utils.go
View File

@ -18,76 +18,9 @@
package observability
import (
"context"
"strings"
"github.com/ProtonMail/gluon/async"
"github.com/ProtonMail/go-proton-api"
"github.com/ProtonMail/proton-bridge/v3/internal/plan"
)
const (
planUnknown = "unknown"
planOther = "other"
planBusiness = "business"
planIndividual = "individual"
planGroup = "group"
)
var planHierarchy = map[string]int{ //nolint:gochecknoglobals
planBusiness: 4,
planGroup: 3,
planIndividual: 2,
planOther: 1,
planUnknown: 0,
}
type planGetter interface {
GetOrganizationData(ctx context.Context) (proton.OrganizationResponse, error)
}
func isHigherPriority(currentPlan, newPlan string) bool {
newRank, ok := planHierarchy[newPlan]
if !ok {
return false
}
currentRank, ok2 := planHierarchy[currentPlan]
if !ok2 {
return true // we don't have a valid plan, might as well replace it
}
return newRank > currentRank
}
func mapUserPlan(planName string) string {
if planName == "" {
return planUnknown
}
switch strings.TrimSpace(strings.ToLower(planName)) {
case "mail2022":
return planIndividual
case "bundle2022":
return planIndividual
case "family2022":
return planGroup
case "visionary2022":
return planGroup
case "mailpro2022":
return planBusiness
case "planbiz2024":
return planBusiness
case "bundlepro2022":
return planBusiness
case "bundlepro2024":
return planBusiness
case "duo2024":
return planGroup
default:
return planOther
}
}
func (d *distinctionUtility) setUserPlan(planName string) {
if planName == "" {
return
@ -96,24 +29,12 @@ func (d *distinctionUtility) setUserPlan(planName string) {
d.userPlanLock.Lock()
defer d.userPlanLock.Unlock()
userPlanMapped := mapUserPlan(planName)
if isHigherPriority(d.userPlanUnsafe, userPlanMapped) {
userPlanMapped := plan.MapUserPlan(planName)
if plan.IsHigherPriority(d.userPlanUnsafe, userPlanMapped) {
d.userPlanUnsafe = userPlanMapped
}
}
func (d *distinctionUtility) registerUserPlan(ctx context.Context, getter planGetter, panicHandler async.PanicHandler) {
go func() {
defer async.HandlePanic(panicHandler)
orgRes, err := getter.GetOrganizationData(ctx)
if err != nil {
return
}
d.setUserPlan(orgRes.Organization.PlanName)
}()
}
func (d *distinctionUtility) getUserPlanSafe() string {
d.userPlanLock.Lock()
defer d.userPlanLock.Unlock()

4
internal/services/observability/service.go
View File

@ -250,7 +250,7 @@ func (s *Service) addMetricsIfClients(metric ...proton.ObservabilityMetric) {
s.addMetrics(metric...)
}
func (s *Service) RegisterUserClient(userID string, protonClient *proton.Client, telemetryService *telemetry.Service) {
func (s *Service) RegisterUserClient(userID string, protonClient *proton.Client, telemetryService *telemetry.Service, userPlan string) {
s.log.Info("Registering user client, ID:", userID)
s.withUserClientStoreLock(func() {
@ -260,7 +260,7 @@ func (s *Service) RegisterUserClient(userID string, protonClient *proton.Client,
}
})
s.distinctionUtility.registerUserPlan(s.ctx, protonClient, s.panicHandler)
s.distinctionUtility.setUserPlan(userPlan)
// There may be a case where we already have metric updates stored, so try to flush;
s.sendSignal(s.signalDataArrived)

21
internal/services/observability/test_utils.go
View File

@ -20,15 +20,16 @@ package observability
import (
gluonMetrics "github.com/ProtonMail/gluon/observability/metrics"
"github.com/ProtonMail/go-proton-api"
"github.com/ProtonMail/proton-bridge/v3/internal/plan"
)
func GenerateAllUsedDistinctionMetricPermutations() []proton.ObservabilityMetric {
planValues := []string{
planUnknown,
planOther,
planBusiness,
planIndividual,
planGroup}
plan.Unknown,
plan.Other,
plan.Business,
plan.Individual,
plan.Group}
mailClientValues := []string{
emailAgentAppleMail,
emailAgentOutlook,
@ -58,11 +59,11 @@ func GenerateAllUsedDistinctionMetricPermutations() []proton.ObservabilityMetric
func GenerateAllHeartbeatMetricPermutations() []proton.ObservabilityMetric {
planValues := []string{
planUnknown,
planOther,
planBusiness,
planIndividual,
planGroup}
plan.Unknown,
plan.Other,
plan.Business,
plan.Individual,
plan.Group}
mailClientValues := []string{
emailAgentAppleMail,
emailAgentOutlook,

5
internal/services/observability/utils_test.go
View File

@ -104,8 +104,3 @@ func TestMatchUserAgent(t *testing.T) {
require.Equal(t, testCase.result, matchUserAgent(testCase.agent))
}
}
func TestFormatBool(t *testing.T) {
require.Equal(t, "false", formatBool(false))
require.Equal(t, "true", formatBool(true))
}

89
internal/telemetry/heartbeat.go
View File

@ -19,9 +19,12 @@ package telemetry
import (
"context"
"math"
"strconv"
"strings"
"time"
"github.com/ProtonMail/proton-bridge/v3/internal/plan"
"github.com/ProtonMail/proton-bridge/v3/internal/updater"
"github.com/sirupsen/logrus"
)
@ -32,70 +35,66 @@ func NewHeartbeat(manager HeartbeatManager, imapPort, smtpPort int, cacheDir, ke
manager: manager,
metrics: HeartbeatData{
MeasurementGroup: "bridge.any.usage",
Event: "bridge_heartbeat",
Event: "bridge_heartbeat_new",
Dimensions: NewHeartbeatDimensions(),
},
defaultIMAPPort: imapPort,
defaultSMTPPort: smtpPort,
defaultCache: cacheDir,
defaultKeychain: keychain,
defaultUserPlan: plan.Unknown,
}
return heartbeat
}
func (heartbeat *Heartbeat) SetRollout(val float64) {
heartbeat.metrics.Dimensions.Rollout = strconv.Itoa(int(val * 100))
heartbeat.metrics.Values.Rollout = int(math.Floor(val * 10))
}
func (heartbeat *Heartbeat) SetNbAccount(val int) {
heartbeat.metrics.Values.NbAccount = val
func (heartbeat *Heartbeat) GetRollout() int {
return heartbeat.metrics.Values.Rollout
}
func (heartbeat *Heartbeat) SetNumberConnectedAccounts(val int) {
heartbeat.metrics.Values.NumberConnectedAccounts = val
}
func (heartbeat *Heartbeat) SetAutoUpdate(val bool) {
if val {
heartbeat.metrics.Dimensions.AutoUpdate = dimensionON
} else {
heartbeat.metrics.Dimensions.AutoUpdate = dimensionOFF
}
heartbeat.metrics.Dimensions.AutoUpdateEnabled = strconv.FormatBool(val)
}
func (heartbeat *Heartbeat) SetAutoStart(val bool) {
if val {
heartbeat.metrics.Dimensions.AutoStart = dimensionON
} else {
heartbeat.metrics.Dimensions.AutoStart = dimensionOFF
}
heartbeat.metrics.Dimensions.AutoStartEnabled = strconv.FormatBool(val)
}
func (heartbeat *Heartbeat) SetBeta(val updater.Channel) {
if val == updater.EarlyChannel {
heartbeat.metrics.Dimensions.Beta = dimensionON
} else {
heartbeat.metrics.Dimensions.Beta = dimensionOFF
}
heartbeat.metrics.Dimensions.BetaEnabled = strconv.FormatBool(val == updater.EarlyChannel)
}
func (heartbeat *Heartbeat) SetDoh(val bool) {
if val {
heartbeat.metrics.Dimensions.Doh = dimensionON
} else {
heartbeat.metrics.Dimensions.Doh = dimensionOFF
}
heartbeat.metrics.Dimensions.DohEnabled = strconv.FormatBool(val)
}
func (heartbeat *Heartbeat) SetSplitMode(val bool) {
if val {
heartbeat.metrics.Dimensions.SplitMode = dimensionON
} else {
heartbeat.metrics.Dimensions.SplitMode = dimensionOFF
heartbeat.metrics.Dimensions.UseSplitMode = strconv.FormatBool(val)
}
func (heartbeat *Heartbeat) SetUserPlan(val string) {
mappedUserPlan := plan.MapUserPlan(val)
if plan.IsHigherPriority(heartbeat.metrics.Dimensions.UserPlanGroup, mappedUserPlan) {
heartbeat.metrics.Dimensions.UserPlanGroup = val
}
}
func (heartbeat *Heartbeat) SetContactedByAppleNotes(uaName string) {
uaNameLowered := strings.ToLower(uaName)
if strings.Contains(uaNameLowered, "mac") && strings.Contains(uaNameLowered, "notes") {
heartbeat.metrics.Dimensions.ContactedByAppleNotes = strconv.FormatBool(true)
}
}
func (heartbeat *Heartbeat) SetShowAllMail(val bool) {
if val {
heartbeat.metrics.Dimensions.ShowAllMail = dimensionON
} else {
heartbeat.metrics.Dimensions.ShowAllMail = dimensionOFF
}
heartbeat.metrics.Dimensions.ShowAllMail = strconv.FormatBool(val)
}
func (heartbeat *Heartbeat) SetIMAPConnectionMode(val bool) {
@ -115,35 +114,19 @@ func (heartbeat *Heartbeat) SetSMTPConnectionMode(val bool) {
}
func (heartbeat *Heartbeat) SetIMAPPort(val int) {
if val == heartbeat.defaultIMAPPort {
heartbeat.metrics.Dimensions.IMAPPort = dimensionDefault
} else {
heartbeat.metrics.Dimensions.IMAPPort = dimensionCustom
}
heartbeat.metrics.Dimensions.UseDefaultIMAPPort = strconv.FormatBool(val == heartbeat.defaultIMAPPort)
}
func (heartbeat *Heartbeat) SetSMTPPort(val int) {
if val == heartbeat.defaultSMTPPort {
heartbeat.metrics.Dimensions.SMTPPort = dimensionDefault
} else {
heartbeat.metrics.Dimensions.SMTPPort = dimensionCustom
}
heartbeat.metrics.Dimensions.UseDefaultSMTPPort = strconv.FormatBool(val == heartbeat.defaultSMTPPort)
}
func (heartbeat *Heartbeat) SetCacheLocation(val string) {
if val == heartbeat.defaultCache {
heartbeat.metrics.Dimensions.CacheLocation = dimensionDefault
} else {
heartbeat.metrics.Dimensions.CacheLocation = dimensionCustom
}
heartbeat.metrics.Dimensions.UseDefaultCacheLocation = strconv.FormatBool(val == heartbeat.defaultCache)
}
func (heartbeat *Heartbeat) SetKeyChainPref(val string) {
if val == heartbeat.defaultKeychain {
heartbeat.metrics.Dimensions.KeychainPref = dimensionDefault
} else {
heartbeat.metrics.Dimensions.KeychainPref = dimensionCustom
}
heartbeat.metrics.Dimensions.UseDefaultKeychain = strconv.FormatBool(val == heartbeat.defaultKeychain)
}
func (heartbeat *Heartbeat) SetPrevVersion(val string) {

64
internal/telemetry/heartbeat_test.go
View File

@ -22,34 +22,38 @@ import (
"testing"
"time"
"github.com/ProtonMail/proton-bridge/v3/internal/plan"
"github.com/ProtonMail/proton-bridge/v3/internal/telemetry"
"github.com/ProtonMail/proton-bridge/v3/internal/telemetry/mocks"
"github.com/golang/mock/gomock"
"github.com/stretchr/testify/require"
)
func TestHeartbeat_default_heartbeat(t *testing.T) {
withHeartbeat(t, 1143, 1025, "/tmp", "defaultKeychain", func(hb *telemetry.Heartbeat, mock *mocks.MockHeartbeatManager) {
data := telemetry.HeartbeatData{
MeasurementGroup: "bridge.any.usage",
Event: "bridge_heartbeat",
Event: "bridge_heartbeat_new",
Values: telemetry.HeartbeatValues{
NbAccount: 1,
NumberConnectedAccounts: 1,
Rollout: 1,
},
Dimensions: telemetry.HeartbeatDimensions{
AutoUpdate: "on",
AutoStart: "on",
Beta: "off",
Doh: "off",
SplitMode: "off",
ShowAllMail: "off",
IMAPConnectionMode: "ssl",
SMTPConnectionMode: "ssl",
IMAPPort: "default",
SMTPPort: "default",
CacheLocation: "default",
KeychainPref: "default",
PrevVersion: "1.2.3",
Rollout: "10",
AutoUpdateEnabled: "true",
AutoStartEnabled: "true",
BetaEnabled: "false",
DohEnabled: "false",
UseSplitMode: "false",
ShowAllMail: "false",
UseDefaultIMAPPort: "true",
UseDefaultSMTPPort: "true",
UseDefaultCacheLocation: "true",
UseDefaultKeychain: "true",
ContactedByAppleNotes: "false",
PrevVersion: "1.2.3",
IMAPConnectionMode: "ssl",
SMTPConnectionMode: "ssl",
UserPlanGroup: plan.Unknown,
},
}
@ -81,7 +85,7 @@ func withHeartbeat(t *testing.T, imap, smtp int, cache, keychain string, tests f
heartbeat := telemetry.NewHeartbeat(manager, imap, smtp, cache, keychain)
heartbeat.SetRollout(0.1)
heartbeat.SetNbAccount(1)
heartbeat.SetNumberConnectedAccounts(1)
heartbeat.SetSplitMode(false)
heartbeat.SetAutoStart(true)
heartbeat.SetAutoUpdate(true)
@ -98,3 +102,29 @@ func withHeartbeat(t *testing.T, imap, smtp int, cache, keychain string, tests f
tests(&heartbeat, manager)
}
func Test_setRollout(t *testing.T) {
hb := telemetry.Heartbeat{}
type testStruct struct {
val float64
res int
}
tests := []testStruct{
{0.02, 0},
{0.04, 0},
{0.09999, 0},
{0.1, 1},
{0.132323, 1},
{0.2, 2},
{0.25, 2},
{0.7111, 7},
{0.93, 9},
{0.999, 9},
}
for _, test := range tests {
hb.SetRollout(test.val)
require.Equal(t, test.res, hb.GetRollout())
}
}

62
internal/telemetry/types_heartbeat.go
View File

@ -21,14 +21,11 @@ import (
"context"
"time"
"github.com/ProtonMail/proton-bridge/v3/internal/plan"
"github.com/sirupsen/logrus"
)
const (
dimensionON = "on"
dimensionOFF = "off"
dimensionDefault = "default"
dimensionCustom = "custom"
dimensionSSL = "ssl"
dimensionStartTLS = "starttls"
)
@ -46,24 +43,29 @@ type HeartbeatManager interface {
}
type HeartbeatValues struct {
NbAccount int `json:"nb_account"`
NumberConnectedAccounts int `json:"numberConnectedAccounts"`
Rollout int `json:"rolloutPercentage"`
}
type HeartbeatDimensions struct {
AutoUpdate string `json:"auto_update"`
AutoStart string `json:"auto_start"`
Beta string `json:"beta"`
Doh string `json:"doh"`
SplitMode string `json:"split_mode"`
ShowAllMail string `json:"show_all_mail"`
IMAPConnectionMode string `json:"imap_connection_mode"`
SMTPConnectionMode string `json:"smtp_connection_mode"`
IMAPPort string `json:"imap_port"`
SMTPPort string `json:"smtp_port"`
CacheLocation string `json:"cache_location"`
KeychainPref string `json:"keychain_pref"`
PrevVersion string `json:"prev_version"`
Rollout string `json:"rollout"`
// Fields below correspond to bool
AutoUpdateEnabled string `json:"isAutoUpdateEnabled"`
AutoStartEnabled string `json:"isAutoStartEnabled"`
BetaEnabled string `json:"isBetaEnabled"`
DohEnabled string `json:"isDohEnabled"`
UseSplitMode string `json:"usesSplitMode"`
ShowAllMail string `json:"useAllMail"`
UseDefaultIMAPPort string `json:"useDefaultImapPort"`
UseDefaultSMTPPort string `json:"useDefaultSmtpPort"`
UseDefaultCacheLocation string `json:"useDefaultCacheLocation"`
UseDefaultKeychain string `json:"useDefaultKeychain"`
ContactedByAppleNotes string `json:"isContactedByAppleNotes"`
// Fields below are enums.
PrevVersion string `json:"prevVersion"` // Free text (exception)
IMAPConnectionMode string `json:"imapConnectionMode"`
SMTPConnectionMode string `json:"smtpConnectionMode"`
UserPlanGroup string `json:"bridgePlanGroup"`
}
type HeartbeatData struct {
@ -82,4 +84,26 @@ type Heartbeat struct {
defaultSMTPPort int
defaultCache string
defaultKeychain string
defaultUserPlan string
}
func NewHeartbeatDimensions() HeartbeatDimensions {
return HeartbeatDimensions{
AutoUpdateEnabled: "false",
AutoStartEnabled: "false",
BetaEnabled: "false",
DohEnabled: "false",
UseSplitMode: "false",
ShowAllMail: "false",
UseDefaultIMAPPort: "false",
UseDefaultSMTPPort: "false",
UseDefaultCacheLocation: "false",
UseDefaultKeychain: "false",
ContactedByAppleNotes: "false",
PrevVersion: "unknown",
IMAPConnectionMode: dimensionSSL,
SMTPConnectionMode: dimensionSSL,
UserPlanGroup: plan.Unknown,
}
}

19
internal/user/user.go
View File

@ -63,6 +63,8 @@ type User struct {
id string
log *logrus.Entry
userPlan string
vault *vault.User
client *proton.Client
reporter reporter.Reporter
@ -176,6 +178,14 @@ func newImpl(
return nil, fmt.Errorf("failed to get addresses: %w", err)
}
// Get the user's plan name.
var userPlan string
if organizationData, err := client.GetOrganizationData(ctx); err != nil {
logrus.WithError(err).Info("Failed to obtain user organization data")
} else {
userPlan = organizationData.Organization.Name
}
// Get the user's API labels.
apiLabels, err := client.GetLabels(ctx, proton.LabelTypeSystem, proton.LabelTypeFolder, proton.LabelTypeLabel)
if err != nil {
@ -197,6 +207,8 @@ func newImpl(
log: logrus.WithField("userID", apiUser.ID),
id: apiUser.ID,
userPlan: userPlan,
vault: encVault,
client: client,
reporter: reporter,
@ -317,7 +329,7 @@ func newImpl(
user.identityService.Start(ctx, user.serviceGroup)
// Add user client to observability service
observabilityService.RegisterUserClient(user.id, client, user.telemetryService)
observabilityService.RegisterUserClient(user.id, client, user.telemetryService, userPlan)
// Start Notification service
user.notificationService.Start(ctx, user.serviceGroup)
@ -416,6 +428,11 @@ func (user *User) GetAddressMode() vault.AddressMode {
return user.vault.AddressMode()
}
// GetUserPlanName returns the user's subscription plan name.
func (user *User) GetUserPlanName() string {
return user.userPlan
}
// SetAddressMode sets the user's address mode.
func (user *User) SetAddressMode(ctx context.Context, mode vault.AddressMode) error {
user.log.WithField("mode", mode).Info("Setting address mode")

184
tests/features/bridge/heartbeat.feature
View File

@ -1,6 +1,8 @@
Feature: Send Telemetry Heartbeat
Background:
Given there exists an account with username "[user:user1]" and password "password"
And there exists an account with username "[user:user2]" and password "password"
And there exists an account with username "[user:user3]" and password "password"
Then it succeeds
When bridge starts
Then it succeeds
@ -12,28 +14,30 @@ Feature: Send Telemetry 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",
"Event": "bridge_heartbeat_new",
"Values": {
"nb_account": 1
"NumberConnectedAccounts": 1,
"rolloutPercentage": 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"
"isAutoUpdateEnabled": "true",
"isAutoStartEnabled": "true",
"isBetaEnabled": "false",
"isDohEnabled": "false",
"usesSplitMode": "false",
"useAllMail": "true",
"useDefaultImapPort": "true",
"useDefaultSmtpPort": "true",
"useDefaultCacheLocation": "true",
"useDefaultKeychain": "true",
"isContactedByAppleNotes": "false",
"imapConnectionMode": "starttls",
"smtpConnectionMode": "starttls",
"prevVersion": "0.0.0",
"bridgePlanGroup": "unknown"
}
}
"""
@ -59,25 +63,27 @@ Feature: Send Telemetry Heartbeat
"""
{
"MeasurementGroup": "bridge.any.usage",
"Event": "bridge_heartbeat",
"Event": "bridge_heartbeat_new",
"Values": {
"nb_account": 1
"NumberConnectedAccounts": 1,
"rolloutPercentage": 1
},
"Dimensions": {
"auto_update": "off",
"auto_start": "off",
"beta": "off",
"doh": "on",
"split_mode": "off",
"show_all_mail": "off",
"imap_connection_mode": "ssl",
"smtp_connection_mode": "ssl",
"imap_port": "custom",
"smtp_port": "custom",
"cache_location": "custom",
"keychain_pref": "custom",
"prev_version": "0.0.0",
"rollout": "42"
"isAutoUpdateEnabled": "false",
"isAutoStartEnabled": "false",
"isBetaEnabled": "false",
"isDohEnabled": "true",
"usesSplitMode": "false",
"useAllMail": "false",
"useDefaultImapPort": "false",
"useDefaultSmtpPort": "false",
"useDefaultCacheLocation": "false",
"useDefaultKeychain": "false",
"isContactedByAppleNotes": "false",
"imapConnectionMode": "ssl",
"smtpConnectionMode": "ssl",
"prevVersion": "0.0.0",
"bridgePlanGroup": "unknown"
}
}
"""
@ -97,25 +103,105 @@ Feature: Send Telemetry Heartbeat
"""
{
"MeasurementGroup": "bridge.any.usage",
"Event": "bridge_heartbeat",
"Event": "bridge_heartbeat_new",
"Values": {
"nb_account": 1
"NumberConnectedAccounts": 1,
"rolloutPercentage": 1
},
"Dimensions": {
"auto_update": "on",
"auto_start": "on",
"beta": "off",
"doh": "off",
"split_mode": "on",
"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"
"isAutoUpdateEnabled": "true",
"isAutoStartEnabled": "true",
"isBetaEnabled": "false",
"isDohEnabled": "false",
"usesSplitMode": "true",
"useAllMail": "true",
"useDefaultImapPort": "true",
"useDefaultSmtpPort": "true",
"useDefaultCacheLocation": "true",
"useDefaultKeychain": "true",
"isContactedByAppleNotes": "false",
"imapConnectionMode": "starttls",
"smtpConnectionMode": "starttls",
"prevVersion": "0.0.0",
"bridgePlanGroup": "unknown"
}
}
"""
And bridge do not need to send heartbeat
Scenario: Multiple-users on Bridge reported correctly
Then bridge telemetry feature is enabled
When the user logs in with username "[user:user1]" and password "password"
Then it succeeds
When the user logs in with username "[user:user2]" and password "password"
Then it succeeds
When the user logs in with username "[user:user3]" and password "password"
Then it succeeds
When bridge needs to explicitly send heartbeat
Then bridge eventually sends the following heartbeat:
"""
{
"MeasurementGroup": "bridge.any.usage",
"Event": "bridge_heartbeat_new",
"Values": {
"NumberConnectedAccounts": 3,
"rolloutPercentage": 1
},
"Dimensions": {
"isAutoUpdateEnabled": "true",
"isAutoStartEnabled": "true",
"isBetaEnabled": "false",
"isDohEnabled": "false",
"usesSplitMode": "false",
"useAllMail": "true",
"useDefaultImapPort": "true",
"useDefaultSmtpPort": "true",
"useDefaultCacheLocation": "true",
"useDefaultKeychain": "true",
"isContactedByAppleNotes": "false",
"imapConnectionMode": "starttls",
"smtpConnectionMode": "starttls",
"prevVersion": "0.0.0",
"bridgePlanGroup": "unknown"
}
}
"""
And bridge do not need to send heartbeat
Scenario: Send heartbeat explicitly - apple notes tried to connect
Then bridge telemetry feature is enabled
When the user logs in with username "[user:user1]" and password "password"
Then it succeeds
When user "[user:user1]" connects IMAP client "1"
And IMAP client "1" announces its ID with name "Mac OS X Notes" and version "14.5"
When bridge needs to explicitly send heartbeat
Then bridge eventually sends the following heartbeat:
"""
{
"MeasurementGroup": "bridge.any.usage",
"Event": "bridge_heartbeat_new",
"Values": {
"NumberConnectedAccounts": 1,
"rolloutPercentage": 1
},
"Dimensions": {
"isAutoUpdateEnabled": "true",
"isAutoStartEnabled": "true",
"isBetaEnabled": "false",
"isDohEnabled": "false",
"usesSplitMode": "false",
"useAllMail": "true",
"useDefaultImapPort": "true",
"useDefaultSmtpPort": "true",
"useDefaultCacheLocation": "true",
"useDefaultKeychain": "true",
"isContactedByAppleNotes": "true",
"imapConnectionMode": "starttls",
"smtpConnectionMode": "starttls",
"prevVersion": "0.0.0",
"bridgePlanGroup": "unknown"
}
}
"""

6
tests/heartbeat_test.go
View File

@ -54,6 +54,10 @@ func (s *scenario) bridgeNeedsToSendHeartbeat() error {
return nil
}
func (s *scenario) bridgeNeedsToSendExplicitHeartbeat() error {
return s.t.heartbeat.SetLastHeartbeatSent(time.Now().Add(-24 * time.Hour))
}
func (s *scenario) bridgeDoNotNeedToSendHeartbeat() error {
last := s.t.heartbeat.GetLastHeartbeatSent()
if isAnotherDay(last, time.Now()) {
@ -73,7 +77,7 @@ func matchHeartbeat(have, want telemetry.HeartbeatData) error {
}
// Ignore rollout number
want.Dimensions.Rollout = have.Dimensions.Rollout
want.Values.Rollout = have.Values.Rollout
if have != want {
return fmt.Errorf("missing heartbeat: have %#v, want %#v", have, want)

2
tests/steps_test.go
View File

@ -207,6 +207,8 @@ func (s *scenario) steps(ctx *godog.ScenarioContext) {
// ==== TELEMETRY ====
ctx.Step(`^bridge eventually sends the following heartbeat:$`, s.bridgeEventuallySendsTheFollowingHeartbeat)
ctx.Step(`^bridge needs to send heartbeat`, s.bridgeNeedsToSendHeartbeat)
ctx.Step(`^bridge needs to explicitly send heartbeat`, s.bridgeNeedsToSendExplicitHeartbeat)
ctx.Step(`^bridge do not need to send heartbeat`, s.bridgeDoNotNeedToSendHeartbeat)
ctx.Step(`^heartbeat is not whitelisted`, s.heartbeatIsNotwhitelisted)