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

GODT-35: New pmapi client and manager using resty

Browse Source
This commit is contained in:
James Houlahan
2021-02-22 18:23:51 +01:00
committed by Jakub
parent 1d538e8540
commit 2284e9ede1
163 changed files with 3333 additions and 8124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
internal/transfer/mocks/mocks.go
View File

@ -1,5 +1,5 @@
// Code generated by MockGen. DO NOT EDIT.
// Source: github.com/ProtonMail/proton-bridge/internal/transfer (interfaces: PanicHandler,ClientManager,IMAPClientProvider)
// Source: github.com/ProtonMail/proton-bridge/internal/transfer (interfaces: PanicHandler,IMAPClientProvider)
// Package mocks is a generated GoMock package.
package mocks
@ -7,7 +7,6 @@ package mocks
import (
reflect "reflect"
pmapi "github.com/ProtonMail/proton-bridge/pkg/pmapi"
imap "github.com/emersion/go-imap"
sasl "github.com/emersion/go-sasl"
gomock "github.com/golang/mock/gomock"
@ -48,57 +47,6 @@ func (mr *MockPanicHandlerMockRecorder) HandlePanic() *gomock.Call {
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HandlePanic", reflect.TypeOf((*MockPanicHandler)(nil).HandlePanic))
}
// MockClientManager is a mock of ClientManager interface
type MockClientManager struct {
ctrl *gomock.Controller
recorder *MockClientManagerMockRecorder
}
// MockClientManagerMockRecorder is the mock recorder for MockClientManager
type MockClientManagerMockRecorder struct {
mock *MockClientManager
}
// NewMockClientManager creates a new mock instance
func NewMockClientManager(ctrl *gomock.Controller) *MockClientManager {
mock := &MockClientManager{ctrl: ctrl}
mock.recorder = &MockClientManagerMockRecorder{mock}
return mock
}
// EXPECT returns an object that allows the caller to indicate expected use
func (m *MockClientManager) EXPECT() *MockClientManagerMockRecorder {
return m.recorder
}
// CheckConnection mocks base method
func (m *MockClientManager) CheckConnection() error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "CheckConnection")
ret0, _ := ret[0].(error)
return ret0
}
// CheckConnection indicates an expected call of CheckConnection
func (mr *MockClientManagerMockRecorder) CheckConnection() *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckConnection", reflect.TypeOf((*MockClientManager)(nil).CheckConnection))
}
// GetClient mocks base method
func (m *MockClientManager) GetClient(arg0 string) pmapi.Client {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetClient", arg0)
ret0, _ := ret[0].(pmapi.Client)
return ret0
}
// GetClient indicates an expected call of GetClient
func (mr *MockClientManagerMockRecorder) GetClient(arg0 interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetClient", reflect.TypeOf((*MockClientManager)(nil).GetClient), arg0)
}
// MockIMAPClientProvider is a mock of IMAPClientProvider interface
type MockIMAPClientProvider struct {
ctrl *gomock.Controller

21
internal/transfer/provider_imap_utils.go
View File

@ -25,7 +25,6 @@ import (
imapID "github.com/ProtonMail/go-imap-id"
"github.com/ProtonMail/proton-bridge/internal/constants"
"github.com/ProtonMail/proton-bridge/pkg/pmapi"
"github.com/emersion/go-imap"
imapClient "github.com/emersion/go-imap/client"
"github.com/emersion/go-sasl"
@ -118,15 +117,19 @@ func (p *IMAPProvider) tryReconnect(ensureSelectedIn string) error {
return previousErr
}
err := pmapi.CheckConnection()
log.WithError(err).Debug("Connection check")
if err != nil {
time.Sleep(imapReconnectSleep)
previousErr = err
continue
}
// FIXME(conman): This should register as connection observer.
err = p.reauth()
/*
err := pmapi.CheckConnection()
log.WithError(err).Debug("Connection check")
if err != nil {
time.Sleep(imapReconnectSleep)
previousErr = err
continue
}
*/
err := p.reauth()
log.WithError(err).Debug("Reauth")
if err != nil {
time.Sleep(imapReconnectSleep)

39
internal/transfer/provider_pmapi.go
View File

@ -18,6 +18,7 @@
package transfer
import (
"context"
"sort"
"github.com/ProtonMail/gopenpgp/v2/crypto"
@ -34,25 +35,27 @@ const (
// PMAPIProvider implements import and export to/from ProtonMail server.
type PMAPIProvider struct {
clientManager ClientManager
userID string
addressID string
keyRing *crypto.KeyRing
builder *message.Builder
client pmapi.Client
userID string
addressID string
keyRing *crypto.KeyRing
builder *message.Builder
nextImportRequests map[string]*pmapi.ImportMsgReq // Key is msg transfer ID.
nextImportRequestsSize int
timeIt *timeIt
connection bool
}
// NewPMAPIProvider returns new PMAPIProvider.
func NewPMAPIProvider(clientManager ClientManager, userID, addressID string) (*PMAPIProvider, error) {
func NewPMAPIProvider(client pmapi.Client, userID, addressID string) (*PMAPIProvider, error) {
provider := &PMAPIProvider{
clientManager: clientManager,
userID: userID,
addressID: addressID,
builder: message.NewBuilder(fetchWorkers, attachWorkers, buildWorkers),
client: client,
userID: userID,
addressID: addressID,
builder: message.NewBuilder(fetchWorkers, attachWorkers, buildWorkers),
nextImportRequests: map[string]*pmapi.ImportMsgReq{},
nextImportRequestsSize: 0,
@ -61,7 +64,7 @@ func NewPMAPIProvider(clientManager ClientManager, userID, addressID string) (*P
}
if addressID != "" {
keyRing, err := clientManager.GetClient(userID).KeyRingForAddressID(addressID)
keyRing, err := client.KeyRingForAddressID(addressID)
if err != nil {
return nil, errors.Wrap(err, "failed to get key ring")
}
@ -71,10 +74,6 @@ func NewPMAPIProvider(clientManager ClientManager, userID, addressID string) (*P
return provider, nil
}
func (p *PMAPIProvider) client() pmapi.Client {
return p.clientManager.GetClient(p.userID)
}
// ID returns identifier of current setup of PMAPI provider.
// Identification is unique per user.
func (p *PMAPIProvider) ID() string {
@ -83,7 +82,7 @@ func (p *PMAPIProvider) ID() string {
// Mailboxes returns all available labels in ProtonMail account.
func (p *PMAPIProvider) Mailboxes(includeEmpty, includeAllMail bool) ([]Mailbox, error) {
labels, err := p.client().ListLabels()
labels, err := p.client.ListLabels(context.Background())
if err != nil {
return nil, err
}
@ -92,7 +91,7 @@ func (p *PMAPIProvider) Mailboxes(includeEmpty, includeAllMail bool) ([]Mailbox,
emptyLabelsMap := map[string]bool{}
if !includeEmpty {
messagesCounts, err := p.client().CountMessages(p.addressID)
messagesCounts, err := p.client.CountMessages(context.Background(), p.addressID)
if err != nil {
return nil, err
}
@ -120,7 +119,7 @@ func (p *PMAPIProvider) Mailboxes(includeEmpty, includeAllMail bool) ([]Mailbox,
ID: label.ID,
Name: label.Name,
Color: label.Color,
IsExclusive: label.Exclusive == 1,
IsExclusive: bool(label.Exclusive),
})
}
return mailboxes, nil
@ -160,10 +159,10 @@ func (l byFoldersLabels) Swap(i, j int) {
// Less sorts first folders, then labels, by user order.
func (l byFoldersLabels) Less(i, j int) bool {
if l[i].Exclusive == 1 && l[j].Exclusive == 0 {
if l[i].Exclusive && !l[j].Exclusive {
return true
}
if l[i].Exclusive == 0 && l[j].Exclusive == 1 {
if !l[i].Exclusive && l[j].Exclusive {
return false
}
return l[i].Order < l[j].Order

9
internal/transfer/provider_pmapi_source.go
View File

@ -157,7 +157,7 @@ func (p *PMAPIProvider) exportMessage(rule *Rule, progress *Progress, pmapiMsgID
body, err := p.builder.NewJobWithOptions(
context.Background(),
p.client(),
p.client,
msg.ID,
message.JobOptions{IgnoreDecryptionErrors: !skipEncryptedMessages},
).GetResult()
@ -169,14 +169,9 @@ func (p *PMAPIProvider) exportMessage(rule *Rule, progress *Progress, pmapiMsgID
return Message{Body: []byte(msg.Body)}, err
}
unread := false
if msg.Unread == 1 {
unread = true
}
return Message{
ID: msgID,
Unread: unread,
Unread: bool(msg.Unread),
Body: body,
Sources: []Mailbox{rule.SourceMailbox},
Targets: rule.TargetMailboxes,

30
internal/transfer/provider_pmapi_target.go
View File

@ -19,6 +19,7 @@ package transfer
import (
"bytes"
"context"
"fmt"
"io"
"io/ioutil"
@ -56,7 +57,7 @@ func (p *PMAPIProvider) CreateMailbox(mailbox Mailbox) (Mailbox, error) {
exclusive = 1
}
label, err := p.client().CreateLabel(&pmapi.Label{
label, err := p.client.CreateLabel(context.TODO(), &pmapi.Label{
Name: mailbox.Name,
Color: mailbox.Color,
Exclusive: exclusive,
@ -194,7 +195,7 @@ func (p *PMAPIProvider) transferMessage(rules transferRules, progress *Progress,
return
}
importMsgReqSize := len(importMsgReq.Body)
importMsgReqSize := len(importMsgReq.Message)
if p.nextImportRequestsSize+importMsgReqSize > pmapiImportBatchMaxSize || len(p.nextImportRequests) == pmapiImportBatchMaxItems {
preparedImportRequestsCh <- p.nextImportRequests
p.nextImportRequests = map[string]*pmapi.ImportMsgReq{}
@ -226,9 +227,12 @@ func (p *PMAPIProvider) generateImportMsgReq(rules transferRules, progress *Prog
}
}
unread := 0
var unread pmapi.Boolean
if msg.Unread {
unread = 1
unread = pmapi.True
} else {
unread = pmapi.False
}
labelIDs := []string{}
@ -243,12 +247,14 @@ func (p *PMAPIProvider) generateImportMsgReq(rules transferRules, progress *Prog
}
return &pmapi.ImportMsgReq{
AddressID: p.addressID,
Body: body,
Unread: unread,
Time: message.Time,
Flags: computeMessageFlags(message.Header),
LabelIDs: labelIDs,
Metadata: &pmapi.ImportMetadata{
AddressID: p.addressID,
Unread: unread,
Time: message.Time,
Flags: computeMessageFlags(message.Header),
LabelIDs: labelIDs,
},
Message: body,
}, nil
}
@ -293,7 +299,7 @@ func (p *PMAPIProvider) importMessages(progress *Progress, importRequests map[st
}
importMsgIDs := []string{}
importMsgRequests := []*pmapi.ImportMsgReq{}
importMsgRequests := pmapi.ImportMsgReqs{}
for msgID, req := range importRequests {
importMsgIDs = append(importMsgIDs, msgID)
importMsgRequests = append(importMsgRequests, req)
@ -327,7 +333,7 @@ func (p *PMAPIProvider) importMessages(progress *Progress, importRequests map[st
func (p *PMAPIProvider) importMessage(msgSourceID string, progress *Progress, req *pmapi.ImportMsgReq) (importedID string, importedErr error) {
progress.callWrap(func() error {
results, err := p.importRequest(msgSourceID, []*pmapi.ImportMsgReq{req})
results, err := p.importRequest(msgSourceID, pmapi.ImportMsgReqs{req})
if err != nil {
return errors.Wrap(err, "failed to import messages")
}

27
internal/transfer/provider_pmapi_test.go
View File

@ -19,6 +19,7 @@ package transfer
import (
"bytes"
"context"
"fmt"
"testing"
"time"
@ -33,7 +34,7 @@ func TestPMAPIProviderMailboxes(t *testing.T) {
defer m.ctrl.Finish()
setupPMAPIClientExpectationForExport(&m)
provider, err := NewPMAPIProvider(m.clientManager, "user", "addressID")
provider, err := NewPMAPIProvider(m.pmapiClient, "user", "addressID")
r.NoError(t, err)
tests := []struct {
@ -78,7 +79,7 @@ func TestPMAPIProviderTransferTo(t *testing.T) {
defer m.ctrl.Finish()
setupPMAPIClientExpectationForExport(&m)
provider, err := NewPMAPIProvider(m.clientManager, "user", "addressID")
provider, err := NewPMAPIProvider(m.pmapiClient, "user", "addressID")
r.NoError(t, err)
rules, rulesClose := newTestRules(t)
@ -96,7 +97,7 @@ func TestPMAPIProviderTransferFrom(t *testing.T) {
defer m.ctrl.Finish()
setupPMAPIClientExpectationForImport(&m)
provider, err := NewPMAPIProvider(m.clientManager, "user", "addressID")
provider, err := NewPMAPIProvider(m.pmapiClient, "user", "addressID")
r.NoError(t, err)
rules, rulesClose := newTestRules(t)
@ -114,7 +115,7 @@ func TestPMAPIProviderTransferFromDraft(t *testing.T) {
defer m.ctrl.Finish()
setupPMAPIClientExpectationForImportDraft(&m)
provider, err := NewPMAPIProvider(m.clientManager, "user", "addressID")
provider, err := NewPMAPIProvider(m.pmapiClient, "user", "addressID")
r.NoError(t, err)
rules, rulesClose := newTestRules(t)
@ -133,9 +134,9 @@ func TestPMAPIProviderTransferFromTo(t *testing.T) {
setupPMAPIClientExpectationForExport(&m)
setupPMAPIClientExpectationForImport(&m)
source, err := NewPMAPIProvider(m.clientManager, "user", "addressID")
source, err := NewPMAPIProvider(m.pmapiClient, "user", "addressID")
r.NoError(t, err)
target, err := NewPMAPIProvider(m.clientManager, "user", "addressID")
target, err := NewPMAPIProvider(m.pmapiClient, "user", "addressID")
r.NoError(t, err)
rules, rulesClose := newTestRules(t)
@ -151,22 +152,22 @@ func setupPMAPIRules(rules transferRules) {
func setupPMAPIClientExpectationForExport(m *mocks) {
m.pmapiClient.EXPECT().KeyRingForAddressID(gomock.Any()).Return(m.keyring, nil).AnyTimes()
m.pmapiClient.EXPECT().ListLabels().Return([]*pmapi.Label{
m.pmapiClient.EXPECT().ListLabels(gomock.Any()).Return([]*pmapi.Label{
{ID: "label1", Name: "Foo", Color: "blue", Exclusive: 0, Order: 2},
{ID: "label2", Name: "Bar", Color: "green", Exclusive: 0, Order: 1},
{ID: "folder1", Name: "One", Color: "red", Exclusive: 1, Order: 1},
{ID: "folder2", Name: "Two", Color: "orange", Exclusive: 1, Order: 2},
}, nil).AnyTimes()
m.pmapiClient.EXPECT().CountMessages(gomock.Any()).Return([]*pmapi.MessagesCount{
m.pmapiClient.EXPECT().CountMessages(gomock.Any(), gomock.Any()).Return([]*pmapi.MessagesCount{
{LabelID: "label1", Total: 10},
{LabelID: "label2", Total: 0},
{LabelID: "folder1", Total: 20},
}, nil).AnyTimes()
m.pmapiClient.EXPECT().ListMessages(gomock.Any()).Return([]*pmapi.Message{
m.pmapiClient.EXPECT().ListMessages(gomock.Any(), gomock.Any()).Return([]*pmapi.Message{
{ID: "msg1"},
{ID: "msg2"},
}, 2, nil).AnyTimes()
m.pmapiClient.EXPECT().GetMessage(gomock.Any()).DoAndReturn(func(msgID string) (*pmapi.Message, error) {
m.pmapiClient.EXPECT().GetMessage(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, msgID string) (*pmapi.Message, error) {
return &pmapi.Message{
ID: msgID,
Body: string(getTestMsgBody(msgID)),
@ -177,11 +178,11 @@ func setupPMAPIClientExpectationForExport(m *mocks) {
func setupPMAPIClientExpectationForImport(m *mocks) {
m.pmapiClient.EXPECT().KeyRingForAddressID(gomock.Any()).Return(m.keyring, nil).AnyTimes()
m.pmapiClient.EXPECT().Import(gomock.Any()).DoAndReturn(func(requests []*pmapi.ImportMsgReq) ([]*pmapi.ImportMsgRes, error) {
m.pmapiClient.EXPECT().Import(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, requests pmapi.ImportMsgReqs) ([]*pmapi.ImportMsgRes, error) {
results := []*pmapi.ImportMsgRes{}
for _, request := range requests {
for _, msgID := range []string{"msg1", "msg2"} {
if bytes.Contains(request.Body, []byte(msgID)) {
if bytes.Contains(request.Message, []byte(msgID)) {
results = append(results, &pmapi.ImportMsgRes{MessageID: msgID, Error: nil})
}
}
@ -192,7 +193,7 @@ func setupPMAPIClientExpectationForImport(m *mocks) {
func setupPMAPIClientExpectationForImportDraft(m *mocks) {
m.pmapiClient.EXPECT().KeyRingForAddressID(gomock.Any()).Return(m.keyring, nil).AnyTimes()
m.pmapiClient.EXPECT().CreateDraft(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(msg *pmapi.Message, parentID string, action int) (*pmapi.Message, error) {
m.pmapiClient.EXPECT().CreateDraft(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, msg *pmapi.Message, parentID string, action int) (*pmapi.Message, error) {
r.Equal(m.t, msg.Subject, "draft1")
msg.ID = "draft1"
return msg, nil

32
internal/transfer/provider_pmapi_utils.go
View File

@ -18,6 +18,7 @@
package transfer
import (
"context"
"fmt"
"io"
"time"
@ -57,13 +58,18 @@ func (p *PMAPIProvider) tryReconnect() error {
return previousErr
}
err := p.clientManager.CheckConnection()
log.WithError(err).Debug("Connection check")
if err != nil {
time.Sleep(pmapiReconnectSleep)
previousErr = err
continue
}
// FIXME(conman): This should register as a connection observer somehow...
// Maybe the entire "provider" could register as an observer and pause if it is notified of dropped connection?
/*
err := p.clientManager.CheckConnection()
log.WithError(err).Debug("Connection check")
if err != nil {
time.Sleep(pmapiReconnectSleep)
previousErr = err
continue
}
*/
break
}
@ -77,7 +83,7 @@ func (p *PMAPIProvider) listMessages(filter *pmapi.MessagesFilter) (messages []*
p.timeIt.start("listing", key)
defer p.timeIt.stop("listing", key)
messages, count, err = p.client().ListMessages(filter)
messages, count, err = p.client.ListMessages(context.TODO(), filter)
return err
})
return
@ -88,18 +94,18 @@ func (p *PMAPIProvider) getMessage(msgID string) (message *pmapi.Message, err er
p.timeIt.start("download", msgID)
defer p.timeIt.stop("download", msgID)
message, err = p.client().GetMessage(msgID)
message, err = p.client.GetMessage(context.TODO(), msgID)
return err
})
return
}
func (p *PMAPIProvider) importRequest(msgSourceID string, req []*pmapi.ImportMsgReq) (res []*pmapi.ImportMsgRes, err error) {
func (p *PMAPIProvider) importRequest(msgSourceID string, req pmapi.ImportMsgReqs) (res []*pmapi.ImportMsgRes, err error) {
err = p.ensureConnection(func() error {
p.timeIt.start("upload", msgSourceID)
defer p.timeIt.stop("upload", msgSourceID)
res, err = p.client().Import(req)
res, err = p.client.Import(context.TODO(), req)
return err
})
return
@ -110,7 +116,7 @@ func (p *PMAPIProvider) createDraft(msgSourceID string, message *pmapi.Message,
p.timeIt.start("upload", msgSourceID)
defer p.timeIt.stop("upload", msgSourceID)
draft, err = p.client().CreateDraft(message, parent, action)
draft, err = p.client.CreateDraft(context.TODO(), message, parent, action)
return err
})
return
@ -123,7 +129,7 @@ func (p *PMAPIProvider) createAttachment(msgSourceID string, att *pmapi.Attachme
p.timeIt.start("upload", key)
defer p.timeIt.stop("upload", key)
created, err = p.client().CreateAttachment(att, r, sig)
created, err = p.client.CreateAttachment(context.TODO(), att, r, sig)
return err
})
return

7
internal/transfer/transfer_test.go
View File

@ -23,7 +23,6 @@ import (
"github.com/ProtonMail/gopenpgp/v2/crypto"
transfermocks "github.com/ProtonMail/proton-bridge/internal/transfer/mocks"
"github.com/ProtonMail/proton-bridge/pkg/pmapi"
pmapimocks "github.com/ProtonMail/proton-bridge/pkg/pmapi/mocks"
gomock "github.com/golang/mock/gomock"
)
@ -33,10 +32,8 @@ type mocks struct {
ctrl *gomock.Controller
panicHandler *transfermocks.MockPanicHandler
clientManager *transfermocks.MockClientManager
imapClientProvider *transfermocks.MockIMAPClientProvider
pmapiClient *pmapimocks.MockClient
pmapiConfig *pmapi.ClientConfig
keyring *crypto.KeyRing
}
@ -49,15 +46,11 @@ func initMocks(t *testing.T) mocks {
ctrl: mockCtrl,
panicHandler: transfermocks.NewMockPanicHandler(mockCtrl),
clientManager: transfermocks.NewMockClientManager(mockCtrl),
imapClientProvider: transfermocks.NewMockIMAPClientProvider(mockCtrl),
pmapiClient: pmapimocks.NewMockClient(mockCtrl),
pmapiConfig: &pmapi.ClientConfig{},
keyring: newTestKeyring(),
}
m.clientManager.EXPECT().GetClient("user").Return(m.pmapiClient).AnyTimes()
return m
}

9
internal/transfer/types.go
View File

@ -17,10 +17,6 @@
package transfer
import (
"github.com/ProtonMail/proton-bridge/pkg/pmapi"
)
type PanicHandler interface {
HandlePanic()
}
@ -32,8 +28,3 @@ type MetricsManager interface {
Cancel()
Fail()
}
type ClientManager interface {
GetClient(userID string) pmapi.Client
CheckConnection() error
}