Skip to content

Handle protobuf encoding/decoding errors during uploads and album creation #43

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xob0t wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Handle protobuf encoding/decoding errors during uploads and album creation #43

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 40 additions & 7 deletions gpmc/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,17 @@
import time
from urllib.parse import parse_qs
from pathlib import Path
import logging

import requests
from requests.adapters import HTTPAdapter, Retry
from blackboxprotobuf import decode_message, encode_message
from blackboxprotobuf.lib.exceptions import BlackboxProtobufException

from . import message_types
from .exceptions import UploadRejected
from .exceptions import UploadRejected, ProtobufDecodeError, ProtobufEncodeError

logger = logging.getLogger("rich")

DEFAULT_TIMEOUT = 60
RETRIES = 10
Expand Down Expand Up @@ -233,7 +237,12 @@ def upload_file(self, file: str | Path | bytes | IO[bytes] | Generator[bytes, No

response.raise_for_status()

upload_response_decoded, _ = decode_message(response.content)
try:
upload_response_decoded, _ = decode_message(response.content)
except (BlackboxProtobufException, Exception) as e:
logger.warning(f"Failed to decode upload response: {e}")
logger.debug(f"Response content (first 500 bytes): {response.content[:500]}")
raise ProtobufDecodeError(f"Failed to decode upload response: {e}") from e
return upload_response_decoded

def commit_upload(
Expand Down Expand Up @@ -331,7 +340,12 @@ def commit_upload(
"3": bytes([1, 3]),
}

serialized_data = encode_message(proto_body, message_types.COMMIT_UPLOAD) # type: ignore
try:
serialized_data = encode_message(proto_body, message_types.COMMIT_UPLOAD) # type: ignore
except (BlackboxProtobufException, Exception) as e:
logger.error(f"Failed to encode commit_upload request for file '{file_name}'")
logger.debug(f"Proto body keys: {proto_body.keys()}")
raise ProtobufEncodeError(f"Failed to encode commit_upload request: {e}") from e

headers = {
"Accept-Encoding": "gzip",
Expand All @@ -350,10 +364,16 @@ def commit_upload(
timeout=self.timeout,
)
response.raise_for_status()
decoded_message, _ = decode_message(response.content)
try:
decoded_message, _ = decode_message(response.content)
except (BlackboxProtobufException, Exception) as e:
logger.error(f"Failed to decode commit_upload response for file '{file_name}'")
logger.debug(f"Response content (first 500 bytes): {response.content[:500]}")
raise ProtobufDecodeError(f"Failed to decode commit_upload response: {e}") from e
try:
media_key = decoded_message["1"]["3"]["1"]
except KeyError as e:
logger.warning(f"File upload rejected by api for file '{file_name}'")
raise UploadRejected("File upload rejected by api") from e
return media_key

Expand Down Expand Up @@ -422,7 +442,13 @@ def create_album(self, album_name: str, media_keys: Sequence[str]) -> str:
"8": {"3": self.model, "4": self.make, "5": self.android_api_version},
}

serialized_data = encode_message(proto_body, message_types.CREATE_ALBUM) # type: ignore
try:
serialized_data = encode_message(proto_body, message_types.CREATE_ALBUM) # type: ignore
except (BlackboxProtobufException, Exception) as e:
logger.error(f"Failed to encode create_album request for album '{album_name}' with {len(media_keys)} media items")
logger.debug(f"Proto body: {proto_body}")
logger.debug(f"Media keys: {media_keys[:5]}..." if len(media_keys) > 5 else f"Media keys: {media_keys}")
raise ProtobufEncodeError(f"Failed to encode create_album request: {e}") from e

headers = {
"Accept-Encoding": "gzip",
Expand All @@ -442,8 +468,15 @@ def create_album(self, album_name: str, media_keys: Sequence[str]) -> str:
)
response.raise_for_status()

decoded_message, _ = decode_message(response.content)
return decoded_message["1"]["1"]
try:
decoded_message, _ = decode_message(response.content)
return decoded_message["1"]["1"]
except (BlackboxProtobufException, KeyError, Exception) as e:
logger.error(f"Failed to decode create_album response for album '{album_name}'")
logger.debug(f"Response content (first 500 bytes): {response.content[:500]}")
if isinstance(e, KeyError):
raise ProtobufDecodeError(f"Unexpected response structure when creating album: {e}") from e
raise ProtobufDecodeError(f"Failed to decode create_album response: {e}") from e

def add_media_to_album(self, album_media_key: str, media_keys: Sequence[str]) -> dict:
"""Add media to an album.
Expand Down
27 changes: 19 additions & 8 deletions gpmc/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
from . import utils
from .hash_handler import calculate_sha1_hash, convert_sha1_hash
from .db_update_parser import parse_db_update
from .exceptions import ProtobufDecodeError, ProtobufEncodeError

# Make Ctrl+C work for cancelling threads
signal.signal(signal.SIGINT, signal.SIG_DFL)
Expand Down Expand Up @@ -476,8 +477,13 @@ def _upload_concurrently(self, path_hash_pairs: Mapping[Path, bytes | str], thre
try:
media_key_dict = future.result()
uploaded_files = uploaded_files | media_key_dict
except (ProtobufDecodeError, ProtobufEncodeError) as e:
self.logger.error(f"Protobuf error uploading file {file[0]}: {e}")
self.logger.info(f"Skipping {file[0]} due to protobuf encoding/decoding issue")
upload_error_count += 1
overall_progress.update(task_id=overall_task_id, description=f"[bold red] Errors: {upload_error_count}")
except Exception as e:
self.logger.error(f"Error uploading file {file}: {e}")
self.logger.error(f"Error uploading file {file[0]}: {e}")
upload_error_count += 1
overall_progress.update(task_id=overall_task_id, description=f"[bold red] Errors: {upload_error_count}")
finally:
Expand Down Expand Up @@ -562,13 +568,18 @@ def add_to_album(self, media_keys: Sequence[str], album_name: str, show_progress
current_album_key = None
for j in range(0, len(album_batch), batch_size):
batch = album_batch[j : j + batch_size]
if current_album_key is None:
# Create the album with the first batch
current_album_key = self.api.create_album(album_name=current_album_name, media_keys=batch)
album_keys.append(current_album_key)
else:
# Add to the existing album
self.api.add_media_to_album(album_media_key=current_album_key, media_keys=batch)
try:
if current_album_key is None:
# Create the album with the first batch
current_album_key = self.api.create_album(album_name=current_album_name, media_keys=batch)
album_keys.append(current_album_key)
else:
# Add to the existing album
self.api.add_media_to_album(album_media_key=current_album_key, media_keys=batch)
except (ProtobufDecodeError, ProtobufEncodeError) as e:
self.logger.error(f"Protobuf error while adding items to album '{current_album_name}': {e}")
self.logger.info(f"Batch of {len(batch)} items skipped. Media keys (first 3): {batch[:3]}")
raise
progress.update(task, advance=len(batch))
album_counter += 1
return album_keys
Expand Down
10 changes: 10 additions & 0 deletions gpmc/exceptions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,3 +4,13 @@ class CustomException(Exception):

class UploadRejected(CustomException):
pass


class ProtobufDecodeError(CustomException):
"""Raised when protobuf message decoding fails."""
pass


class ProtobufEncodeError(CustomException):
"""Raised when protobuf message encoding fails."""
pass