From 8b9265ad96acb398185de5395793671e364ea988 Mon Sep 17 00:00:00 2001
From: James Houlahan <james.houlahan@proton.ch>
Date: Thu, 9 Feb 2023 16:24:40 +0100
Subject: [PATCH] feat(GODT-2283): Limit max import size to 30MB (bump GPA to
 v0.4.0)

---
 go.mod                       |  2 +-
 go.sum                       |  4 ++--
 internal/bridge/sync_test.go |  7 +++++--
 internal/user/imap.go        |  9 +++++++--
 tests/ctx_helper_test.go     | 11 ++++++++---
 5 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/go.mod b/go.mod
index cf7bbd2b..d46e2155 100644
--- a/go.mod
+++ b/go.mod
@@ -7,7 +7,7 @@ require (
 	github.com/Masterminds/semver/v3 v3.1.1
 	github.com/ProtonMail/gluon v0.14.2-0.20230206162331-cf36d870802b
 	github.com/ProtonMail/go-autostart v0.0.0-20210130080809-00ed301c8e9a
-	github.com/ProtonMail/go-proton-api v0.3.1-0.20230209105544-b185be77ce15
+	github.com/ProtonMail/go-proton-api v0.4.0
 	github.com/ProtonMail/go-rfc5322 v0.11.0
 	github.com/ProtonMail/gopenpgp/v2 v2.4.10
 	github.com/PuerkitoBio/goquery v1.8.0
diff --git a/go.sum b/go.sum
index 3b64f79c..df711095 100644
--- a/go.sum
+++ b/go.sum
@@ -41,8 +41,8 @@ github.com/ProtonMail/go-message v0.0.0-20210611055058-fabeff2ec753/go.mod h1:NB
 github.com/ProtonMail/go-mime v0.0.0-20220302105931-303f85f7fe0f/go.mod h1:NYt+V3/4rEeDuaev/zw1zCq8uqVEuPHzDPo3OZrlGJ4=
 github.com/ProtonMail/go-mime v0.0.0-20220429130430-2192574d760f h1:4IWzKjHzZxdrW9k4zl/qCwenOVHDbVDADPPHFLjs0Oc=
 github.com/ProtonMail/go-mime v0.0.0-20220429130430-2192574d760f/go.mod h1:qRZgbeASl2a9OwmsV85aWwRqic0NHPh+9ewGAzb4cgM=
-github.com/ProtonMail/go-proton-api v0.3.1-0.20230209105544-b185be77ce15 h1:OiyETrhilrJGS4WXqKXdTO7LnzaDb00/4F3W4JqICrM=
-github.com/ProtonMail/go-proton-api v0.3.1-0.20230209105544-b185be77ce15/go.mod h1:JUo5IQG0hNuPRuDpOUsCOvtee6UjTEHHF1QN2i8RSos=
+github.com/ProtonMail/go-proton-api v0.4.0 h1:Tw8Ieuc355ljPqpIzh/uttpE+5ia0z8GA/ca5iyl/9w=
+github.com/ProtonMail/go-proton-api v0.4.0/go.mod h1:JUo5IQG0hNuPRuDpOUsCOvtee6UjTEHHF1QN2i8RSos=
 github.com/ProtonMail/go-rfc5322 v0.11.0 h1:o5Obrm4DpmQEffvgsVqG6S4BKwC1Wat+hYwjIp2YcCY=
 github.com/ProtonMail/go-rfc5322 v0.11.0/go.mod h1:6oOKr0jXvpoE6pwTx/HukigQpX2J9WUf6h0auplrFTw=
 github.com/ProtonMail/go-srp v0.0.5 h1:xhUioxZgDbCnpo9JehyFhwwsn9JLWkUGfB0oiKXgiGg=
diff --git a/internal/bridge/sync_test.go b/internal/bridge/sync_test.go
index 357cbab1..329a42c5 100644
--- a/internal/bridge/sync_test.go
+++ b/internal/bridge/sync_test.go
@@ -431,7 +431,7 @@ func createMessages(ctx context.Context, t *testing.T, c *proton.Client, addrID,
 	_, ok := addrKRs[addrID]
 	require.True(t, ok)
 
-	res, err := stream.Collect(ctx, c.ImportMessages(
+	str, err := c.ImportMessages(
 		ctx,
 		addrKRs[addrID],
 		runtime.NumCPU(),
@@ -446,7 +446,10 @@ func createMessages(ctx context.Context, t *testing.T, c *proton.Client, addrID,
 				Message: message,
 			}
 		})...,
-	))
+	)
+	require.NoError(t, err)
+
+	res, err := stream.Collect(ctx, str)
 	require.NoError(t, err)
 
 	return xslices.Map(res, func(res proton.ImportRes) string {
diff --git a/internal/user/imap.go b/internal/user/imap.go
index 8be80a4c..84fc8943 100644
--- a/internal/user/imap.go
+++ b/internal/user/imap.go
@@ -538,7 +538,7 @@ func (conn *imapConnector) importMessage(
 
 				messageID = msg.ID
 			} else {
-				res, err := stream.Collect(ctx, conn.client.ImportMessages(ctx, addrKR, 1, 1, []proton.ImportReq{{
+				str, err := conn.client.ImportMessages(ctx, addrKR, 1, 1, []proton.ImportReq{{
 					Metadata: proton.ImportMetadata{
 						AddressID: conn.addrID,
 						LabelIDs:  labelIDs,
@@ -546,7 +546,12 @@ func (conn *imapConnector) importMessage(
 						Flags:     flags,
 					},
 					Message: literal,
-				}}...))
+				}}...)
+				if err != nil {
+					return fmt.Errorf("failed to prepare message for import: %w", err)
+				}
+
+				res, err := stream.Collect(ctx, str)
 				if err != nil {
 					return fmt.Errorf("failed to import message: %w", err)
 				}
diff --git a/tests/ctx_helper_test.go b/tests/ctx_helper_test.go
index 3c803b5e..04c4c51b 100644
--- a/tests/ctx_helper_test.go
+++ b/tests/ctx_helper_test.go
@@ -124,14 +124,19 @@ func (t *testCtx) withAddrKR(
 func (t *testCtx) createMessages(ctx context.Context, username, addrID string, req []proton.ImportReq) error {
 	return t.withClient(ctx, username, func(ctx context.Context, c *proton.Client) error {
 		return t.withAddrKR(ctx, c, username, addrID, func(ctx context.Context, addrKR *crypto.KeyRing) error {
-			if _, err := stream.Collect(ctx, c.ImportMessages(
+			str, err := c.ImportMessages(
 				ctx,
 				addrKR,
 				runtime.NumCPU(),
 				runtime.NumCPU(),
 				req...,
-			)); err != nil {
-				return err
+			)
+			if err != nil {
+				return fmt.Errorf("failed to prepare messages for import: %w", err)
+			}
+
+			if _, err := stream.Collect(ctx, str); err != nil {
+				return fmt.Errorf("failed to import messages: %w", err)
 			}
 
 			return nil