mirror of
https://github.com/ProtonMail/proton-bridge.git
synced 2025-12-20 09:06:45 +00:00
fix(GODT-2822): Handle 429 during message download
When we run into 429 during a message download, do not cancel the whole batch and switch to a sequential downloader to avoid API overload.
This commit is contained in:
Leander Beernaert
400
internal/user/sync_downloader_test.go
Normal file
400
internal/user/sync_downloader_test.go
Normal file
@ -0,0 +1,400 @@
|
||||
// Copyright (c) 2023 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 user
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/ProtonMail/gluon/async"
|
||||
"github.com/ProtonMail/go-proton-api"
|
||||
"github.com/ProtonMail/proton-bridge/v3/internal/user/mocks"
|
||||
"github.com/golang/mock/gomock"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSyncDownloader_Stage1_429(t *testing.T) {
|
||||
// Check 429 is correctly caught and download state recorded correctly
|
||||
// Message 1: All ok
|
||||
// Message 2: Message failed
|
||||
// Message 3: One attachment failed.
|
||||
mockCtrl := gomock.NewController(t)
|
||||
messageDownloader := mocks.NewMockMessageDownloader(mockCtrl)
|
||||
panicHandler := &async.NoopPanicHandler{}
|
||||
ctx := context.Background()
|
||||
|
||||
requests := downloadRequest{
|
||||
ids: []string{"Msg1", "Msg2", "Msg3"},
|
||||
expectedSize: 0,
|
||||
err: nil,
|
||||
}
|
||||
|
||||
messageDownloader.EXPECT().GetMessage(gomock.Any(), gomock.Eq("Msg1")).Times(1).Return(proton.Message{
|
||||
MessageMetadata: proton.MessageMetadata{
|
||||
ID: "MsgID1",
|
||||
NumAttachments: 1,
|
||||
},
|
||||
Attachments: []proton.Attachment{
|
||||
{
|
||||
ID: "Attachment1_1",
|
||||
},
|
||||
},
|
||||
}, nil)
|
||||
|
||||
messageDownloader.EXPECT().GetMessage(gomock.Any(), gomock.Eq("Msg2")).Times(1).Return(proton.Message{}, &proton.APIError{Status: 429})
|
||||
messageDownloader.EXPECT().GetMessage(gomock.Any(), gomock.Eq("Msg3")).Times(1).Return(proton.Message{
|
||||
MessageMetadata: proton.MessageMetadata{
|
||||
ID: "MsgID3",
|
||||
NumAttachments: 2,
|
||||
},
|
||||
Attachments: []proton.Attachment{
|
||||
{
|
||||
ID: "Attachment3_1",
|
||||
},
|
||||
{
|
||||
ID: "Attachment3_2",
|
||||
},
|
||||
},
|
||||
}, nil)
|
||||
|
||||
const attachmentData = "attachment data"
|
||||
|
||||
messageDownloader.EXPECT().GetAttachmentInto(gomock.Any(), gomock.Eq("Attachment1_1"), gomock.Any()).Times(1).DoAndReturn(func(_ context.Context, _ string, r io.ReaderFrom) error {
|
||||
_, err := r.ReadFrom(strings.NewReader(attachmentData))
|
||||
return err
|
||||
})
|
||||
|
||||
messageDownloader.EXPECT().GetAttachmentInto(gomock.Any(), gomock.Eq("Attachment3_1"), gomock.Any()).Times(1).Return(&proton.APIError{Status: 429})
|
||||
messageDownloader.EXPECT().GetAttachmentInto(gomock.Any(), gomock.Eq("Attachment3_2"), gomock.Any()).Times(1).DoAndReturn(func(_ context.Context, _ string, r io.ReaderFrom) error {
|
||||
_, err := r.ReadFrom(strings.NewReader(attachmentData))
|
||||
return err
|
||||
})
|
||||
|
||||
attachmentDownloader := newAttachmentDownloader(ctx, panicHandler, messageDownloader, 1)
|
||||
defer attachmentDownloader.close()
|
||||
|
||||
result, err := downloadMessageStage1(ctx, panicHandler, requests, messageDownloader, attachmentDownloader, 1)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 3, len(result))
|
||||
// Check message 1
|
||||
require.Equal(t, result[0].State, downloadStateFinished)
|
||||
require.Equal(t, result[0].Message.ID, "MsgID1")
|
||||
require.NotEmpty(t, result[0].Message.AttData)
|
||||
require.NotEqual(t, attachmentData, result[0].Message.AttData[0])
|
||||
require.NotNil(t, result[0].Message.AttData[0])
|
||||
require.Nil(t, result[0].err)
|
||||
|
||||
// Check message 2
|
||||
require.Equal(t, result[1].State, downloadStateZero)
|
||||
require.Empty(t, result[1].Message.ID)
|
||||
require.NotNil(t, result[1].err)
|
||||
|
||||
require.Equal(t, result[2].State, downloadStateHasMessage)
|
||||
require.Equal(t, result[2].Message.ID, "MsgID3")
|
||||
require.Equal(t, 2, len(result[2].Message.AttData))
|
||||
require.NotNil(t, result[2].err)
|
||||
require.Nil(t, result[2].Message.AttData[0])
|
||||
require.NotEqual(t, attachmentData, result[2].Message.AttData[1])
|
||||
require.NotNil(t, result[2].err)
|
||||
}
|
||||
|
||||
func TestSyncDownloader_Stage2_Everything200(t *testing.T) {
|
||||
mockCtrl := gomock.NewController(t)
|
||||
messageDownloader := mocks.NewMockMessageDownloader(mockCtrl)
|
||||
ctx := context.Background()
|
||||
|
||||
downloadResult := []downloadResult{
|
||||
{
|
||||
ID: "Msg1",
|
||||
State: downloadStateFinished,
|
||||
},
|
||||
{
|
||||
ID: "Msg2",
|
||||
State: downloadStateFinished,
|
||||
},
|
||||
}
|
||||
|
||||
result, err := downloadMessagesStage2(ctx, downloadResult, messageDownloader, time.Millisecond)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, len(result))
|
||||
}
|
||||
|
||||
func TestSyncDownloader_Stage2_Not429(t *testing.T) {
|
||||
mockCtrl := gomock.NewController(t)
|
||||
messageDownloader := mocks.NewMockMessageDownloader(mockCtrl)
|
||||
ctx := context.Background()
|
||||
|
||||
msgErr := fmt.Errorf("something not 429")
|
||||
downloadResult := []downloadResult{
|
||||
{
|
||||
ID: "Msg1",
|
||||
State: downloadStateFinished,
|
||||
},
|
||||
{
|
||||
ID: "Msg2",
|
||||
State: downloadStateHasMessage,
|
||||
err: msgErr,
|
||||
},
|
||||
{
|
||||
ID: "Msg3",
|
||||
State: downloadStateFinished,
|
||||
},
|
||||
}
|
||||
|
||||
_, err := downloadMessagesStage2(ctx, downloadResult, messageDownloader, time.Millisecond)
|
||||
require.Error(t, err)
|
||||
require.Equal(t, msgErr, err)
|
||||
}
|
||||
|
||||
func TestSyncDownloader_Stage2_API500(t *testing.T) {
|
||||
mockCtrl := gomock.NewController(t)
|
||||
messageDownloader := mocks.NewMockMessageDownloader(mockCtrl)
|
||||
ctx := context.Background()
|
||||
|
||||
msgErr := &proton.APIError{Status: 500}
|
||||
downloadResult := []downloadResult{
|
||||
{
|
||||
ID: "Msg2",
|
||||
State: downloadStateHasMessage,
|
||||
err: msgErr,
|
||||
},
|
||||
{
|
||||
ID: "Msg3",
|
||||
State: downloadStateFinished,
|
||||
},
|
||||
}
|
||||
|
||||
_, err := downloadMessagesStage2(ctx, downloadResult, messageDownloader, time.Millisecond)
|
||||
require.Error(t, err)
|
||||
require.Equal(t, msgErr, err)
|
||||
}
|
||||
|
||||
func TestSyncDownloader_Stage2_Some429(t *testing.T) {
|
||||
mockCtrl := gomock.NewController(t)
|
||||
messageDownloader := mocks.NewMockMessageDownloader(mockCtrl)
|
||||
ctx := context.Background()
|
||||
|
||||
const attachmentData1 = "attachment data 1"
|
||||
const attachmentData2 = "attachment data 2"
|
||||
const attachmentData3 = "attachment data 3"
|
||||
const attachmentData4 = "attachment data 4"
|
||||
|
||||
err429 := &proton.APIError{Status: 429}
|
||||
downloadResult := []downloadResult{
|
||||
{
|
||||
// Full message , but missing 1 of 2 attachments
|
||||
ID: "Msg1",
|
||||
Message: proton.FullMessage{
|
||||
Message: proton.Message{
|
||||
MessageMetadata: proton.MessageMetadata{
|
||||
ID: "Msg1",
|
||||
NumAttachments: 2,
|
||||
},
|
||||
Attachments: []proton.Attachment{
|
||||
{
|
||||
ID: "A3",
|
||||
},
|
||||
{
|
||||
ID: "A4",
|
||||
},
|
||||
},
|
||||
},
|
||||
AttData: [][]byte{
|
||||
nil,
|
||||
[]byte(attachmentData4),
|
||||
},
|
||||
},
|
||||
State: downloadStateHasMessage,
|
||||
err: err429,
|
||||
},
|
||||
{
|
||||
// Full message, but missing all attachments
|
||||
ID: "Msg2",
|
||||
Message: proton.FullMessage{
|
||||
Message: proton.Message{
|
||||
MessageMetadata: proton.MessageMetadata{
|
||||
ID: "Msg2",
|
||||
NumAttachments: 2,
|
||||
},
|
||||
Attachments: []proton.Attachment{
|
||||
{
|
||||
ID: "A1",
|
||||
},
|
||||
{
|
||||
ID: "A2",
|
||||
},
|
||||
},
|
||||
},
|
||||
AttData: nil,
|
||||
},
|
||||
State: downloadStateHasMessage,
|
||||
err: err429,
|
||||
},
|
||||
{
|
||||
// Missing everything
|
||||
ID: "Msg3",
|
||||
State: downloadStateZero,
|
||||
Message: proton.FullMessage{
|
||||
Message: proton.Message{MessageMetadata: proton.MessageMetadata{ID: "Msg3"}},
|
||||
},
|
||||
err: err429,
|
||||
},
|
||||
}
|
||||
|
||||
{
|
||||
// Simulate 2 failures for message 3 body.
|
||||
firstCall := messageDownloader.EXPECT().GetMessage(gomock.Any(), gomock.Eq("Msg3")).Times(2).Return(proton.Message{}, err429)
|
||||
messageDownloader.EXPECT().GetMessage(gomock.Any(), gomock.Eq("Msg3")).After(firstCall).Times(1).Return(proton.Message{
|
||||
MessageMetadata: proton.MessageMetadata{
|
||||
ID: "Msg3",
|
||||
},
|
||||
}, nil)
|
||||
}
|
||||
|
||||
{
|
||||
// Simulate failures for message 2 attachments.
|
||||
firstCall := messageDownloader.EXPECT().GetAttachmentInto(gomock.Any(), gomock.Eq("A1"), gomock.Any()).Times(2).Return(err429)
|
||||
messageDownloader.EXPECT().GetAttachmentInto(gomock.Any(), gomock.Eq("A1"), gomock.Any()).After(firstCall).Times(1).DoAndReturn(func(_ context.Context, _ string, r io.ReaderFrom) error {
|
||||
_, err := r.ReadFrom(strings.NewReader(attachmentData1))
|
||||
return err
|
||||
})
|
||||
messageDownloader.EXPECT().GetAttachmentInto(gomock.Any(), gomock.Eq("A2"), gomock.Any()).Times(1).DoAndReturn(func(_ context.Context, _ string, r io.ReaderFrom) error {
|
||||
_, err := r.ReadFrom(strings.NewReader(attachmentData2))
|
||||
return err
|
||||
})
|
||||
}
|
||||
|
||||
{
|
||||
messageDownloader.EXPECT().GetAttachmentInto(gomock.Any(), gomock.Eq("A3"), gomock.Any()).Times(1).DoAndReturn(func(_ context.Context, _ string, r io.ReaderFrom) error {
|
||||
_, err := r.ReadFrom(strings.NewReader(attachmentData3))
|
||||
return err
|
||||
})
|
||||
}
|
||||
|
||||
messages, err := downloadMessagesStage2(ctx, downloadResult, messageDownloader, time.Millisecond)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 3, len(messages))
|
||||
|
||||
require.Equal(t, messages[0].Message.ID, "Msg1")
|
||||
require.Equal(t, messages[1].Message.ID, "Msg2")
|
||||
require.Equal(t, messages[2].Message.ID, "Msg3")
|
||||
|
||||
// check attachments
|
||||
require.Equal(t, attachmentData3, string(messages[0].AttData[0]))
|
||||
require.Equal(t, attachmentData4, string(messages[0].AttData[1]))
|
||||
require.Equal(t, attachmentData1, string(messages[1].AttData[0]))
|
||||
require.Equal(t, attachmentData2, string(messages[1].AttData[1]))
|
||||
require.Empty(t, messages[2].AttData)
|
||||
}
|
||||
|
||||
func TestSyncDownloader_Stage2_ErrorOnNon429MessageDownload(t *testing.T) {
|
||||
mockCtrl := gomock.NewController(t)
|
||||
messageDownloader := mocks.NewMockMessageDownloader(mockCtrl)
|
||||
ctx := context.Background()
|
||||
|
||||
err429 := &proton.APIError{Status: 429}
|
||||
err500 := &proton.APIError{Status: 500}
|
||||
downloadResult := []downloadResult{
|
||||
{
|
||||
// Missing everything
|
||||
ID: "Msg3",
|
||||
State: downloadStateZero,
|
||||
Message: proton.FullMessage{
|
||||
Message: proton.Message{MessageMetadata: proton.MessageMetadata{ID: "Msg3"}},
|
||||
},
|
||||
err: err429,
|
||||
},
|
||||
{
|
||||
// Full message , but missing 1 of 2 attachments
|
||||
ID: "Msg1",
|
||||
Message: proton.FullMessage{
|
||||
Message: proton.Message{
|
||||
MessageMetadata: proton.MessageMetadata{
|
||||
ID: "Msg1",
|
||||
NumAttachments: 2,
|
||||
},
|
||||
Attachments: []proton.Attachment{
|
||||
{
|
||||
ID: "A3",
|
||||
},
|
||||
{
|
||||
ID: "A4",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
State: downloadStateHasMessage,
|
||||
err: err429,
|
||||
},
|
||||
}
|
||||
|
||||
{
|
||||
// Simulate 2 failures for message 3 body,
|
||||
messageDownloader.EXPECT().GetMessage(gomock.Any(), gomock.Eq("Msg3")).Times(1).Return(proton.Message{}, err500)
|
||||
}
|
||||
|
||||
messages, err := downloadMessagesStage2(ctx, downloadResult, messageDownloader, time.Millisecond)
|
||||
require.Error(t, err)
|
||||
require.Empty(t, 0, messages)
|
||||
}
|
||||
|
||||
func TestSyncDownloader_Stage2_ErrorOnNon429AttachmentDownload(t *testing.T) {
|
||||
mockCtrl := gomock.NewController(t)
|
||||
messageDownloader := mocks.NewMockMessageDownloader(mockCtrl)
|
||||
ctx := context.Background()
|
||||
|
||||
err429 := &proton.APIError{Status: 429}
|
||||
err500 := &proton.APIError{Status: 500}
|
||||
downloadResult := []downloadResult{
|
||||
{
|
||||
// Full message , but missing 1 of 2 attachments
|
||||
ID: "Msg1",
|
||||
Message: proton.FullMessage{
|
||||
Message: proton.Message{
|
||||
MessageMetadata: proton.MessageMetadata{
|
||||
ID: "Msg1",
|
||||
NumAttachments: 2,
|
||||
},
|
||||
Attachments: []proton.Attachment{
|
||||
{
|
||||
ID: "A3",
|
||||
},
|
||||
{
|
||||
ID: "A4",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
State: downloadStateHasMessage,
|
||||
err: err429,
|
||||
},
|
||||
}
|
||||
|
||||
// 429 for first attachment
|
||||
messageDownloader.EXPECT().GetAttachmentInto(gomock.Any(), gomock.Eq("A3"), gomock.Any()).Times(1).Return(err429)
|
||||
// 500 for second attachment
|
||||
messageDownloader.EXPECT().GetAttachmentInto(gomock.Any(), gomock.Eq("A4"), gomock.Any()).Times(1).Return(err500)
|
||||
|
||||
messages, err := downloadMessagesStage2(ctx, downloadResult, messageDownloader, time.Millisecond)
|
||||
require.Error(t, err)
|
||||
require.Empty(t, 0, messages)
|
||||
}
|
||||
Reference in New Issue
Block a user