From baf78514dc3f3c5a2ac634053c9b6a2c6439b8a2 Mon Sep 17 00:00:00 2001
From: Khai Tran <khaitq@cystack.net>
Date: Mon, 3 Jun 2024 21:51:09 +0700
Subject: [PATCH] chore: Handle timeout

---
 src/maily/aws/ses.py |  9 ++++++---
 src/maily/aws/sqs.py |  3 +++
 src/maily/message.py | 14 +++++++++-----
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/maily/aws/ses.py b/src/maily/aws/ses.py
index 4c34068..e911995 100644
--- a/src/maily/aws/ses.py
+++ b/src/maily/aws/ses.py
@@ -95,13 +95,16 @@ def ses_send_raw_email(self, from_address, to_address, subject, message_body, at
 
             store_reply_record(mail, ses_response)
         except ClientError as e:
-            # logger.error(f'[!] ses_client_error_raw_email:::{e.response["Error"]}')
+            # Handel SES HTTP error: https://docs.aws.amazon.com/ses/latest/APIReference-V2/CommonErrors.html
+            error_code = e.response.get("Error", {}).get("Code")
+            if error_code in ["500", "408", "503"]:
+                return None
+            logger.error(f'[!] ses_client_error_raw_email:::{e.response["Error"]}')
             # logger.error(
             #     f'from_address: {from_address}\nto_address: {to_address}\ndata: {msg_with_attachments.as_string()}')
             return False
         except botocore.exceptions.ConnectionClosedError:
-            # TODO: Handle timeout
-            return False
+            return None
         return True
 
     def list_identities(self):
diff --git a/src/maily/aws/sqs.py b/src/maily/aws/sqs.py
index 5418a91..476106b 100644
--- a/src/maily/aws/sqs.py
+++ b/src/maily/aws/sqs.py
@@ -100,6 +100,9 @@ def process_message(message):
                 # Handle domain identity failed
                 if details['status_code'] in [511]:
                     results.update({"success": False})
+                # Handle AWS service timeout
+                elif details['status_code'] in [504]:
+                    results.update({"success": False})
             else:
                 results['details']['sent_success'] = True
         return results
diff --git a/src/maily/message.py b/src/maily/message.py
index 3053c0f..fb4ca9e 100644
--- a/src/maily/message.py
+++ b/src/maily/message.py
@@ -209,10 +209,13 @@ def ses_relay_email(self, from_address, to_address, subject, message_body, attac
                         reply_address=REPLY_EMAIL):
         response = ses_client.ses_send_raw_email(from_address, to_address, subject,
                                                  message_body, attachments, mail, reply_address)
-        if response:
-            return self.response(200, "Sent email to final recipient")
-        else:
+        # Need to retry
+        if response is None:
+            return self.response(504, "SES client sent Raw Email failed. It needs to be re-sent")
+        if response is False:
             return self.response(503, "SES client error on Raw Email")
+        return self.response(200, "Sent email to final recipient")
+
 
     def get_text_html_attachments(self):
         if self.sns_message_content is None:
@@ -229,9 +232,10 @@ def get_text_html_attachments(self):
                     return self.response(404, "Email not in S3")
                 logger.error(f's3_client_error_get_email: {e.response["Error"]}')
                 # we are returning a 500 so that SNS can retry the email processing
-                return self.response(503, "Cannot fetch the message content from S3")
+                # return self.response(503, "Cannot fetch the message content from S3")
+                return self.response(504, "Fetch message from S3 error - boto3.ClientError")
             except botocore.exceptions.ConnectionClosedError:
-                return self.response(503, "Cannot fetch the message content from S3")
+                return self.response(504, "Fetch message from S3 error - boto3.ConnectionClosedError")
         else:
             message_content = self.sns_message_content
         bytes_email_message = message_from_bytes(message_content, policy=policy.default)