From 0e3a549b3f36a6715760bb2e25a303ba9d3dbc1a Mon Sep 17 00:00:00 2001
From: Jakub <jakub@cuth.sk>
Date: Thu, 31 Aug 2023 07:20:19 +0200
Subject: [PATCH] fix(GODT-2867): do not crash on timeout or context cancel.

---
 attachment_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 response.go        |  2 +-
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/attachment_test.go b/attachment_test.go
index c1081a1..42083ca 100644
--- a/attachment_test.go
+++ b/attachment_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"net/http"
+	"sync"
 	"testing"
 
 	"github.com/ProtonMail/go-proton-api"
@@ -42,3 +43,40 @@ func TestAttachment_429Response(t *testing.T) {
 	require.Equal(t, proton.InvalidValue, apiErr.Code)
 	require.Equal(t, "Request failed with status 429", apiErr.Message)
 }
+
+func TestAttachment_ContextCancelled(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	s := server.New()
+	defer s.Close()
+
+	m := proton.New(
+		proton.WithHostURL(s.GetHostURL()),
+		proton.WithTransport(proton.InsecureTransport()),
+	)
+
+	_, _, err := s.CreateUser("user", []byte("pass"))
+	require.NoError(t, err)
+
+	c, _, err := m.NewClientWithLogin(ctx, "user", []byte("pass"))
+	require.NoError(t, err)
+
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+
+	s.AddStatusHook(func(r *http.Request) (int, bool) {
+		wg.Wait()
+		return http.StatusTooManyRequests, true
+	})
+
+	go func() {
+		_, err = c.GetAttachment(ctx, "someID")
+		wg.Done()
+	}()
+
+	cancel()
+
+	wg.Wait()
+	require.Error(t, err)
+	require.True(t, errors.Is(err, context.Canceled))
+}
diff --git a/response.go b/response.go
index d8f160c..2e3f397 100644
--- a/response.go
+++ b/response.go
@@ -192,7 +192,7 @@ func catchDropError(_ *resty.Response, err error) bool {
 //
 // This function also closes the response body.
 func parseResponse(res *resty.Response, err error) (*resty.Response, error) {
-	if res.StatusCode() == 200 {
+	if err != nil || res.StatusCode() == 200 {
 		return res, err
 	}