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

fix(GODT-2774): Add external context to telemetry tasks

Browse Source 
This ensures they get cancelled if the parent context becomes invalid
This commit is contained in:
Leander Beernaert
2023-07-06 12:00:19 +02:00
parent 1c2cb4f439
commit 40c48ba804
18 changed files with 67 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
go.mod
View File

@ -5,7 +5,7 @@ go 1.20
require (
github.com/0xAX/notificator v0.0.0-20220220101646-ee9b8921e557
github.com/Masterminds/semver/v3 v3.2.0
github.com/ProtonMail/gluon v0.16.1-0.20230705085038-02a8a9f2a454
github.com/ProtonMail/gluon v0.16.1-0.20230706110757-a9327fb18611
github.com/ProtonMail/go-autostart v0.0.0-20210130080809-00ed301c8e9a
github.com/ProtonMail/go-proton-api v0.4.1-0.20230704060229-a77a437ec052
github.com/ProtonMail/gopenpgp/v2 v2.7.1-proton

2
go.sum
View File

@ -27,6 +27,8 @@ github.com/ProtonMail/gluon v0.16.1-0.20230704083024-d901d16834de h1:th289W4w6aE
github.com/ProtonMail/gluon v0.16.1-0.20230704083024-d901d16834de/go.mod h1:Og5/Dz1MiGpCJn51XujZwxiLG7WzvvjE5PRpZBQmAHo=
github.com/ProtonMail/gluon v0.16.1-0.20230705085038-02a8a9f2a454 h1:dyvL9eYmjGo06CasWfGDVzlbunqKaQASMcLLT1D2irI=
github.com/ProtonMail/gluon v0.16.1-0.20230705085038-02a8a9f2a454/go.mod h1:Og5/Dz1MiGpCJn51XujZwxiLG7WzvvjE5PRpZBQmAHo=
github.com/ProtonMail/gluon v0.16.1-0.20230706110757-a9327fb18611 h1:QVydPr/+pgz5xihc2ujNNV+qnq3oTidIXvF0PgkcY6U=
github.com/ProtonMail/gluon v0.16.1-0.20230706110757-a9327fb18611/go.mod h1:Og5/Dz1MiGpCJn51XujZwxiLG7WzvvjE5PRpZBQmAHo=
github.com/ProtonMail/go-autostart v0.0.0-20210130080809-00ed301c8e9a h1:D+aZah+k14Gn6kmL7eKxoo/4Dr/lK3ChBcwce2+SQP4=
github.com/ProtonMail/go-autostart v0.0.0-20210130080809-00ed301c8e9a/go.mod h1:oTGdE7/DlWIr23G0IKW3OXK9wZ5Hw1GGiaJFccTvZi4=
github.com/ProtonMail/go-crypto v0.0.0-20230321155629-9a39f2531310/go.mod h1:8TI4H3IbrackdNgv+92dI+rhpCaLqM0IfpgCgenFvRE=

10
internal/bridge/heartbeat.go
View File

@ -32,7 +32,7 @@ import (
const HeartbeatCheckInterval = time.Hour
func (bridge *Bridge) IsTelemetryAvailable() bool {
func (bridge *Bridge) IsTelemetryAvailable(ctx context.Context) bool {
var flag = true
if bridge.GetTelemetryDisabled() {
return false
@ -40,14 +40,14 @@ func (bridge *Bridge) IsTelemetryAvailable() bool {
safe.RLock(func() {
for _, user := range bridge.users {
flag = flag && user.IsTelemetryEnabled(context.Background())
flag = flag && user.IsTelemetryEnabled(ctx)
}
}, bridge.usersLock)
return flag
}
func (bridge *Bridge) SendHeartbeat(heartbeat *telemetry.HeartbeatData) bool {
func (bridge *Bridge) SendHeartbeat(ctx context.Context, heartbeat *telemetry.HeartbeatData) bool {
data, err := json.Marshal(heartbeat)
if err != nil {
if err := bridge.reporter.ReportMessageWithContext("Cannot parse heartbeat data.", reporter.Context{
@ -62,7 +62,7 @@ func (bridge *Bridge) SendHeartbeat(heartbeat *telemetry.HeartbeatData) bool {
safe.RLock(func() {
for _, user := range bridge.users {
if err := user.SendTelemetry(context.Background(), data); err == nil {
if err := user.SendTelemetry(ctx, data); err == nil {
sent = true
break
}
@ -87,7 +87,7 @@ func (bridge *Bridge) StartHeartbeat(manager telemetry.HeartbeatManager) {
bridge.goHeartbeat = bridge.tasks.PeriodicOrTrigger(HeartbeatCheckInterval, 0, func(ctx context.Context) {
logrus.Debug("Checking for heartbeat")
bridge.heartbeat.TrySending()
bridge.heartbeat.TrySending(ctx)
})
bridge.heartbeat.SetRollout(bridge.GetUpdateRollout())

2
internal/bridge/mocks.go
View File

@ -50,7 +50,7 @@ func NewMocks(tb testing.TB, version, minAuto *semver.Version) *Mocks {
mocks.CrashHandler.EXPECT().HandlePanic(gomock.Any()).AnyTimes()
// this is called at start of heartbeat process.
mocks.Heartbeat.EXPECT().IsTelemetryAvailable().AnyTimes()
mocks.Heartbeat.EXPECT().IsTelemetryAvailable(gomock.Any()).AnyTimes()
return mocks
}

17
internal/bridge/mocks/telemetry_mocks.go
View File

@ -5,6 +5,7 @@
package mocks
import (
context "context"
reflect "reflect"
time "time"
@ -50,31 +51,31 @@ func (mr *MockHeartbeatManagerMockRecorder) GetLastHeartbeatSent() *gomock.Call
}
// IsTelemetryAvailable mocks base method.
func (m *MockHeartbeatManager) IsTelemetryAvailable() bool {
func (m *MockHeartbeatManager) IsTelemetryAvailable(arg0 context.Context) bool {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "IsTelemetryAvailable")
ret := m.ctrl.Call(m, "IsTelemetryAvailable", arg0)
ret0, _ := ret[0].(bool)
return ret0
}
// IsTelemetryAvailable indicates an expected call of IsTelemetryAvailable.
func (mr *MockHeartbeatManagerMockRecorder) IsTelemetryAvailable() *gomock.Call {
func (mr *MockHeartbeatManagerMockRecorder) IsTelemetryAvailable(arg0 interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsTelemetryAvailable", reflect.TypeOf((*MockHeartbeatManager)(nil).IsTelemetryAvailable))
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsTelemetryAvailable", reflect.TypeOf((*MockHeartbeatManager)(nil).IsTelemetryAvailable), arg0)
}
// SendHeartbeat mocks base method.
func (m *MockHeartbeatManager) SendHeartbeat(arg0 *telemetry.HeartbeatData) bool {
func (m *MockHeartbeatManager) SendHeartbeat(arg0 context.Context, arg1 *telemetry.HeartbeatData) bool {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "SendHeartbeat", arg0)
ret := m.ctrl.Call(m, "SendHeartbeat", arg0, arg1)
ret0, _ := ret[0].(bool)
return ret0
}
// SendHeartbeat indicates an expected call of SendHeartbeat.
func (mr *MockHeartbeatManagerMockRecorder) SendHeartbeat(arg0 interface{}) *gomock.Call {
func (mr *MockHeartbeatManagerMockRecorder) SendHeartbeat(arg0, arg1 interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SendHeartbeat", reflect.TypeOf((*MockHeartbeatManager)(nil).SendHeartbeat), arg0)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SendHeartbeat", reflect.TypeOf((*MockHeartbeatManager)(nil).SendHeartbeat), arg0, arg1)
}
// SetLastHeartbeatSent mocks base method.

2
internal/bridge/settings.go
View File

@ -297,7 +297,7 @@ func (bridge *Bridge) SetColorScheme(colorScheme string) error {
// Note: it does not clear the keychain. The only entry in the keychain is the vault password,
// which we need at next startup to decrypt the vault.
func (bridge *Bridge) FactoryReset(ctx context.Context) {
useTelemetry := bridge.IsTelemetryAvailable()
useTelemetry := bridge.IsTelemetryAvailable(ctx)
// Delete all the users.
safe.Lock(func() {

3
internal/bridge/smtp_backend.go
View File

@ -18,6 +18,7 @@
package bridge
import (
"context"
"fmt"
"io"
"strings"
@ -61,7 +62,7 @@ func (s *smtpSession) AuthPlain(username, password string) error {
s.Bridge.setUserAgent(useragent.UnknownClient, useragent.DefaultVersion)
}
user.SendConfigStatusSuccess()
user.SendConfigStatusSuccess(context.Background())
return nil
}

4
internal/bridge/user.go
View File

@ -243,7 +243,7 @@ func (bridge *Bridge) LogoutUser(ctx context.Context, userID string) error {
func (bridge *Bridge) DeleteUser(ctx context.Context, userID string) error {
logrus.WithField("userID", userID).Info("Deleting user")
useTelemetry := bridge.IsTelemetryAvailable()
useTelemetry := bridge.IsTelemetryAvailable(ctx)
return safe.LockRet(func() error {
if !bridge.vault.HasUser(userID) {
@ -602,7 +602,7 @@ func (bridge *Bridge) logoutUser(ctx context.Context, user *user.User, withAPI,
// if this is actually a remove account
if withData && withAPI {
user.SendConfigStatusAbort(withTelemetry)
user.SendConfigStatusAbort(ctx, withTelemetry)
}
logrus.WithFields(logrus.Fields{

7
internal/telemetry/heartbeat.go
View File

@ -18,6 +18,7 @@
package telemetry
import (
"context"
"strconv"
"time"
@ -149,12 +150,12 @@ func (heartbeat *Heartbeat) SetPrevVersion(val string) {
heartbeat.metrics.Dimensions.PrevVersion = val
}
func (heartbeat *Heartbeat) TrySending() {
if heartbeat.manager.IsTelemetryAvailable() {
func (heartbeat *Heartbeat) TrySending(ctx context.Context) {
if heartbeat.manager.IsTelemetryAvailable(ctx) {
lastSent := heartbeat.manager.GetLastHeartbeatSent()
now := time.Now()
if now.Year() > lastSent.Year() || (now.Year() == lastSent.Year() && now.YearDay() > lastSent.YearDay()) {
if !heartbeat.manager.SendHeartbeat(&heartbeat.metrics) {
if !heartbeat.manager.SendHeartbeat(ctx, &heartbeat.metrics) {
heartbeat.log.WithFields(logrus.Fields{
"metrics": heartbeat.metrics,
}).Error("Failed to send heartbeat")

11
internal/telemetry/heartbeat_test.go
View File

@ -18,6 +18,7 @@
package telemetry_test
import (
"context"
"testing"
"time"
@ -52,21 +53,21 @@ func TestHeartbeat_default_heartbeat(t *testing.T) {
},
}
mock.EXPECT().IsTelemetryAvailable().Return(true)
mock.EXPECT().IsTelemetryAvailable(context.Background()).Return(true)
mock.EXPECT().GetLastHeartbeatSent().Return(time.Date(2022, 6, 4, 0, 0, 0, 0, time.UTC))
mock.EXPECT().SendHeartbeat(&data).Return(true)
mock.EXPECT().SendHeartbeat(context.Background(), &data).Return(true)
mock.EXPECT().SetLastHeartbeatSent(gomock.Any()).Return(nil)
hb.TrySending()
hb.TrySending(context.Background())
})
}
func TestHeartbeat_already_sent_heartbeat(t *testing.T) {
withHeartbeat(t, 1143, 1025, "/tmp", "defaultKeychain", func(hb *telemetry.Heartbeat, mock *mocks.MockHeartbeatManager) {
mock.EXPECT().IsTelemetryAvailable().Return(true)
mock.EXPECT().IsTelemetryAvailable(context.Background()).Return(true)
mock.EXPECT().GetLastHeartbeatSent().Return(time.Now().Truncate(24 * time.Hour))
hb.TrySending()
hb.TrySending(context.Background())
})
}

17
internal/telemetry/mocks/mocks.go
View File

@ -5,6 +5,7 @@
package mocks
import (
context "context"
reflect "reflect"
time "time"
@ -50,31 +51,31 @@ func (mr *MockHeartbeatManagerMockRecorder) GetLastHeartbeatSent() *gomock.Call
}
// IsTelemetryAvailable mocks base method.
func (m *MockHeartbeatManager) IsTelemetryAvailable() bool {
func (m *MockHeartbeatManager) IsTelemetryAvailable(arg0 context.Context) bool {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "IsTelemetryAvailable")
ret := m.ctrl.Call(m, "IsTelemetryAvailable", arg0)
ret0, _ := ret[0].(bool)
return ret0
}
// IsTelemetryAvailable indicates an expected call of IsTelemetryAvailable.
func (mr *MockHeartbeatManagerMockRecorder) IsTelemetryAvailable() *gomock.Call {
func (mr *MockHeartbeatManagerMockRecorder) IsTelemetryAvailable(arg0 interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsTelemetryAvailable", reflect.TypeOf((*MockHeartbeatManager)(nil).IsTelemetryAvailable))
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsTelemetryAvailable", reflect.TypeOf((*MockHeartbeatManager)(nil).IsTelemetryAvailable), arg0)
}
// SendHeartbeat mocks base method.
func (m *MockHeartbeatManager) SendHeartbeat(arg0 *telemetry.HeartbeatData) bool {
func (m *MockHeartbeatManager) SendHeartbeat(arg0 context.Context, arg1 *telemetry.HeartbeatData) bool {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "SendHeartbeat", arg0)
ret := m.ctrl.Call(m, "SendHeartbeat", arg0, arg1)
ret0, _ := ret[0].(bool)
return ret0
}
// SendHeartbeat indicates an expected call of SendHeartbeat.
func (mr *MockHeartbeatManagerMockRecorder) SendHeartbeat(arg0 interface{}) *gomock.Call {
func (mr *MockHeartbeatManagerMockRecorder) SendHeartbeat(arg0, arg1 interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SendHeartbeat", reflect.TypeOf((*MockHeartbeatManager)(nil).SendHeartbeat), arg0)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SendHeartbeat", reflect.TypeOf((*MockHeartbeatManager)(nil).SendHeartbeat), arg0, arg1)
}
// SetLastHeartbeatSent mocks base method.

5
internal/telemetry/types_heartbeat.go
View File

@ -18,6 +18,7 @@
package telemetry
import (
"context"
"time"
"github.com/sirupsen/logrus"
@ -33,12 +34,12 @@ const (
)
type Availability interface {
IsTelemetryAvailable() bool
IsTelemetryAvailable(ctx context.Context) bool
}
type HeartbeatManager interface {
Availability
SendHeartbeat(heartbeat *HeartbeatData) bool
SendHeartbeat(ctx context.Context, heartbeat *HeartbeatData) bool
GetLastHeartbeatSent() time.Time
SetLastHeartbeatSent(time.Time) error
}

26
internal/user/config_status.go
View File

@ -25,12 +25,12 @@ import (
"github.com/ProtonMail/proton-bridge/v3/internal/configstatus"
)
func (user *User) SendConfigStatusSuccess() {
func (user *User) SendConfigStatusSuccess(ctx context.Context) {
if user.configStatus.IsFromFailure() {
user.SendConfigStatusRecovery()
user.SendConfigStatusRecovery(ctx)
return
}
if !user.telemetryManager.IsTelemetryAvailable() {
if !user.telemetryManager.IsTelemetryAvailable(ctx) {
return
}
if !user.configStatus.IsPending() {
@ -49,7 +49,7 @@ func (user *User) SendConfigStatusSuccess() {
return
}
if err := user.SendTelemetry(context.Background(), data); err == nil {
if err := user.SendTelemetry(ctx, data); err == nil {
user.log.Info("Configuration Status Success event sent.")
if err := user.configStatus.ApplySuccess(); err != nil {
user.log.WithError(err).Error("Failed to ApplySuccess on config_status.")
@ -57,7 +57,7 @@ func (user *User) SendConfigStatusSuccess() {
}
}
func (user *User) SendConfigStatusAbort(withTelemetry bool) {
func (user *User) SendConfigStatusAbort(ctx context.Context, withTelemetry bool) {
if err := user.configStatus.Remove(); err != nil {
user.log.WithError(err).Error("Failed to remove config_status file.")
}
@ -80,17 +80,17 @@ func (user *User) SendConfigStatusAbort(withTelemetry bool) {
return
}
if err := user.SendTelemetry(context.Background(), data); err == nil {
if err := user.SendTelemetry(ctx, data); err == nil {
user.log.Info("Configuration Status Abort event sent.")
}
}
func (user *User) SendConfigStatusRecovery() {
func (user *User) SendConfigStatusRecovery(ctx context.Context) {
if !user.configStatus.IsFromFailure() {
user.SendConfigStatusSuccess()
user.SendConfigStatusSuccess(ctx)
return
}
if !user.telemetryManager.IsTelemetryAvailable() {
if !user.telemetryManager.IsTelemetryAvailable(ctx) {
return
}
if !user.configStatus.IsPending() {
@ -109,7 +109,7 @@ func (user *User) SendConfigStatusRecovery() {
return
}
if err := user.SendTelemetry(context.Background(), data); err == nil {
if err := user.SendTelemetry(ctx, data); err == nil {
user.log.Info("Configuration Status Recovery event sent.")
if err := user.configStatus.ApplySuccess(); err != nil {
user.log.WithError(err).Error("Failed to ApplySuccess on config_status.")
@ -117,8 +117,8 @@ func (user *User) SendConfigStatusRecovery() {
}
}
func (user *User) SendConfigStatusProgress() {
if !user.telemetryManager.IsTelemetryAvailable() {
func (user *User) SendConfigStatusProgress(ctx context.Context) {
if !user.telemetryManager.IsTelemetryAvailable(ctx) {
return
}
if !user.configStatus.IsPending() {
@ -143,7 +143,7 @@ func (user *User) SendConfigStatusProgress() {
return
}
if err := user.SendTelemetry(context.Background(), data); err == nil {
if err := user.SendTelemetry(ctx, data); err == nil {
user.log.Info("Configuration Status Progress event sent.")
if err := user.configStatus.ApplyProgress(); err != nil {
user.log.WithError(err).Error("Failed to ApplyProgress on config_status.")

4
internal/user/imap.go
View File

@ -76,7 +76,7 @@ func newIMAPConnector(user *User, addrID string) *imapConnector {
}
// Authorize returns whether the given username/password combination are valid for this connector.
func (conn *imapConnector) Authorize(username string, password []byte) bool {
func (conn *imapConnector) Authorize(ctx context.Context, username string, password []byte) bool {
addrID, err := conn.CheckAuth(username, password)
if err != nil {
return false
@ -86,7 +86,7 @@ func (conn *imapConnector) Authorize(username string, password []byte) bool {
return false
}
conn.User.SendConfigStatusSuccess()
conn.User.SendConfigStatusSuccess(ctx)
return true
}

2
internal/user/user.go
View File

@ -180,7 +180,7 @@ func New(
// Check for status_progress when triggered.
user.goStatusProgress = user.tasks.PeriodicOrTrigger(configstatus.ProgressCheckInterval, 0, func(ctx context.Context) {
user.SendConfigStatusProgress()
user.SendConfigStatusProgress(ctx)
})
defer user.goStatusProgress()

2
internal/user/user_test.go
View File

@ -148,7 +148,7 @@ func withUser(tb testing.TB, ctx context.Context, _ *server.Server, m *proton.Ma
ctl := gomock.NewController(tb)
defer ctl.Finish()
manager := mocks.NewMockHeartbeatManager(ctl)
manager.EXPECT().IsTelemetryAvailable().AnyTimes()
manager.EXPECT().IsTelemetryAvailable(context.Background()).AnyTimes()
user, err := New(ctx, vaultUser, client, nil, apiUser, nil, true, vault.DefaultMaxSyncMemory, tb.TempDir(), manager)
require.NoError(tb, err)
defer user.Close()

2
tests/bridge_test.go
View File

@ -314,7 +314,7 @@ func (s *scenario) bridgeTelemetryFeatureDisabled() error {
}
func (s *scenario) checkTelemetry(expect bool) error {
res := s.t.bridge.IsTelemetryAvailable()
res := s.t.bridge.IsTelemetryAvailable(context.Background())
if res != expect {
return fmt.Errorf("expected telemetry feature %v but got %v ", expect, res)
}

7
tests/ctx_heartbeat_test.go
View File

@ -18,6 +18,7 @@
package tests
import (
"context"
"errors"
"testing"
"time"
@ -54,14 +55,14 @@ func (hb *heartbeatRecorder) GetLastHeartbeatSent() time.Time {
return hb.bridge.GetLastHeartbeatSent()
}
func (hb *heartbeatRecorder) IsTelemetryAvailable() bool {
func (hb *heartbeatRecorder) IsTelemetryAvailable(ctx context.Context) bool {
if hb.bridge == nil {
return false
}
return hb.bridge.IsTelemetryAvailable()
return hb.bridge.IsTelemetryAvailable(ctx)
}
func (hb *heartbeatRecorder) SendHeartbeat(metrics *telemetry.HeartbeatData) bool {
func (hb *heartbeatRecorder) SendHeartbeat(_ context.Context, metrics *telemetry.HeartbeatData) bool {
if hb.bridge == nil {
return false
}