305 lines
12 KiB
Python
305 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright 2015 Google Inc. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Helper functions for composite upload tracker file functionality."""
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import print_function
|
|
from __future__ import division
|
|
from __future__ import unicode_literals
|
|
|
|
from collections import namedtuple
|
|
import errno
|
|
import json
|
|
import random
|
|
|
|
import six
|
|
|
|
import gslib
|
|
from gslib.exception import CommandException
|
|
from gslib.tracker_file import (WriteJsonDataToTrackerFile,
|
|
RaiseUnwritableTrackerFileException)
|
|
from gslib.utils.constants import UTF8
|
|
|
|
ObjectFromTracker = namedtuple('ObjectFromTracker', 'object_name generation')
|
|
|
|
|
|
class _CompositeUploadTrackerEntry(object):
|
|
"""Enum class for composite upload tracker file JSON keys."""
|
|
COMPONENTS_LIST = 'components'
|
|
COMPONENT_NAME = 'component_name'
|
|
COMPONENT_GENERATION = 'component_generation'
|
|
ENC_SHA256 = 'encryption_key_sha256'
|
|
PREFIX = 'prefix'
|
|
|
|
|
|
def ReadParallelUploadTrackerFile(tracker_file_name, logger):
|
|
"""Read the tracker file from the last parallel composite upload attempt.
|
|
|
|
If it exists, the tracker file is of the format described in
|
|
WriteParallelUploadTrackerFile or a legacy format. If the file doesn't exist
|
|
or is formatted incorrectly, then the upload will start from the beginning.
|
|
|
|
This function is not thread-safe and must be protected by a lock if
|
|
called within Command.Apply.
|
|
|
|
Args:
|
|
tracker_file_name: The name of the tracker file to read parse.
|
|
logger: logging.Logger for outputting log messages.
|
|
|
|
Returns:
|
|
enc_key_sha256: Encryption key SHA256 used to encrypt the existing
|
|
components, or None if an encryption key was not used.
|
|
component_prefix: String prefix used in naming the existing components, or
|
|
None if no prefix was found.
|
|
existing_components: A list of ObjectFromTracker objects representing
|
|
the set of files that have already been uploaded.
|
|
"""
|
|
enc_key_sha256 = None
|
|
prefix = None
|
|
existing_components = []
|
|
tracker_file = None
|
|
|
|
# If we already have a matching tracker file, get the serialization data
|
|
# so that we can resume the upload.
|
|
try:
|
|
tracker_file = open(tracker_file_name, 'r')
|
|
tracker_data = tracker_file.read()
|
|
tracker_json = json.loads(tracker_data)
|
|
enc_key_sha256 = tracker_json[_CompositeUploadTrackerEntry.ENC_SHA256]
|
|
prefix = tracker_json[_CompositeUploadTrackerEntry.PREFIX]
|
|
for component in tracker_json[_CompositeUploadTrackerEntry.COMPONENTS_LIST]:
|
|
existing_components.append(
|
|
ObjectFromTracker(
|
|
component[_CompositeUploadTrackerEntry.COMPONENT_NAME],
|
|
component[_CompositeUploadTrackerEntry.COMPONENT_GENERATION]))
|
|
except IOError as e:
|
|
# Ignore non-existent file (happens first time a upload is attempted on an
|
|
# object, or when re-starting an upload after a
|
|
# ResumableUploadStartOverException), but warn user for other errors.
|
|
if e.errno != errno.ENOENT:
|
|
logger.warn(
|
|
'Couldn\'t read upload tracker file (%s): %s. Restarting '
|
|
'parallel composite upload from scratch.', tracker_file_name,
|
|
e.strerror)
|
|
except (KeyError, ValueError) as e:
|
|
# Legacy format did not support user-supplied encryption.
|
|
enc_key_sha256 = None
|
|
(prefix, existing_components) = _ParseLegacyTrackerData(tracker_data)
|
|
finally:
|
|
if tracker_file:
|
|
tracker_file.close()
|
|
|
|
return (enc_key_sha256, prefix, existing_components)
|
|
|
|
|
|
def _ParseLegacyTrackerData(tracker_data):
|
|
"""Parses a legacy parallel composite upload tracker file.
|
|
|
|
Args:
|
|
tracker_data: Legacy tracker file contents.
|
|
|
|
Returns:
|
|
component_prefix: The prefix used in naming the existing components, or
|
|
None if no prefix was found.
|
|
existing_components: A list of ObjectFromTracker objects representing
|
|
the set of files that have already been uploaded.
|
|
"""
|
|
# Old tracker files used a non-JSON format.
|
|
# The first line represents the prefix, followed by line pairs of object_name
|
|
# and generation. Discard the last blank line.
|
|
old_tracker_data = tracker_data.split('\n')[:-1]
|
|
prefix = None
|
|
existing_components = []
|
|
if old_tracker_data:
|
|
prefix = old_tracker_data[0]
|
|
i = 1
|
|
while i < len(old_tracker_data) - 1:
|
|
(name, generation) = (old_tracker_data[i], old_tracker_data[i + 1])
|
|
if not generation:
|
|
# Cover the '' case.
|
|
generation = None
|
|
existing_components.append(ObjectFromTracker(name, generation))
|
|
i += 2
|
|
return (prefix, existing_components)
|
|
|
|
|
|
def ValidateParallelCompositeTrackerData(tracker_file_name, existing_enc_sha256,
|
|
existing_prefix, existing_components,
|
|
current_enc_key_sha256, bucket_url,
|
|
command_obj, logger, delete_func,
|
|
delete_exc_handler):
|
|
"""Validates that tracker data matches the current encryption key.
|
|
|
|
If the data does not match, makes a best-effort attempt to delete existing
|
|
temporary component objects encrypted with the old key.
|
|
|
|
Args:
|
|
tracker_file_name: String file name of tracker file.
|
|
existing_enc_sha256: Encryption key SHA256 used to encrypt the existing
|
|
components, or None if an encryption key was not used.
|
|
existing_prefix: String prefix used in naming the existing components, or
|
|
None if no prefix was found.
|
|
existing_components: A list of ObjectFromTracker objects representing
|
|
the set of files that have already been uploaded.
|
|
current_enc_key_sha256: Current Encryption key SHA256 that should be used
|
|
to encrypt objects.
|
|
bucket_url: Bucket URL in which the components exist.
|
|
command_obj: Command class for calls to Apply.
|
|
logger: logging.Logger for outputting log messages.
|
|
delete_func: command.Apply-callable function for deleting objects.
|
|
delete_exc_handler: Exception handler for delete_func.
|
|
|
|
Returns:
|
|
prefix: existing_prefix, or None if the encryption key did not match.
|
|
existing_components: existing_components, or empty list if the encryption
|
|
key did not match.
|
|
"""
|
|
if six.PY3:
|
|
if isinstance(existing_enc_sha256, str):
|
|
existing_enc_sha256 = existing_enc_sha256.encode(UTF8)
|
|
if isinstance(current_enc_key_sha256, str):
|
|
current_enc_key_sha256 = current_enc_key_sha256.encode(UTF8)
|
|
if existing_prefix and existing_enc_sha256 != current_enc_key_sha256:
|
|
try:
|
|
logger.warn(
|
|
'Upload tracker file (%s) does not match current encryption '
|
|
'key. Deleting old components and restarting upload from '
|
|
'scratch with a new tracker file that uses the current '
|
|
'encryption key.', tracker_file_name)
|
|
components_to_delete = []
|
|
for component in existing_components:
|
|
url = bucket_url.Clone()
|
|
url.object_name = component.object_name
|
|
url.generation = component.generation
|
|
|
|
command_obj.Apply(
|
|
delete_func,
|
|
components_to_delete,
|
|
delete_exc_handler,
|
|
arg_checker=gslib.command.DummyArgChecker,
|
|
parallel_operations_override=command_obj.ParallelOverrideReason.SPEED)
|
|
except: # pylint: disable=bare-except
|
|
# Regardless of why we can't clean up old components, need to proceed
|
|
# with the user's original intent to upload the file, so merely warn.
|
|
component_names = [
|
|
component.object_name for component in existing_components
|
|
]
|
|
logger.warn(
|
|
'Failed to delete some of the following temporary objects:\n%s\n'
|
|
'(Continuing on to re-upload components from scratch.)',
|
|
'\n'.join(component_names))
|
|
|
|
# Encryption keys have changed, so the old components and prefix
|
|
# cannot be used.
|
|
return (None, [])
|
|
|
|
return (existing_prefix, existing_components)
|
|
|
|
|
|
def GenerateComponentObjectPrefix(encryption_key_sha256=None):
|
|
"""Generates a random prefix for component objects.
|
|
|
|
Args:
|
|
encryption_key_sha256: Encryption key SHA256 that will be used to encrypt
|
|
the components. This is hashed into the prefix to avoid collision
|
|
during resumption with a different encryption key.
|
|
|
|
Returns:
|
|
String prefix for use in the composite upload.
|
|
"""
|
|
return str(
|
|
(random.randint(1, (10**10) - 1) + hash(encryption_key_sha256)) % 10**10)
|
|
|
|
|
|
def WriteComponentToParallelUploadTrackerFile(tracker_file_name,
|
|
tracker_file_lock,
|
|
component,
|
|
logger,
|
|
encryption_key_sha256=None):
|
|
"""Rewrites an existing tracker file with info about the uploaded component.
|
|
|
|
Follows the format described in _CreateParallelUploadTrackerFile.
|
|
|
|
Args:
|
|
tracker_file_name: Tracker file to append to.
|
|
tracker_file_lock: Thread and process-safe Lock protecting the tracker file.
|
|
component: ObjectFromTracker describing the object that was uploaded.
|
|
logger: logging.Logger for outputting log messages.
|
|
encryption_key_sha256: Encryption key SHA256 for use in this upload, if any.
|
|
"""
|
|
with tracker_file_lock:
|
|
(existing_enc_key_sha256, prefix,
|
|
existing_components) = (ReadParallelUploadTrackerFile(
|
|
tracker_file_name, logger))
|
|
if existing_enc_key_sha256 != encryption_key_sha256:
|
|
raise CommandException(
|
|
'gsutil client error: encryption key SHA256 (%s) in tracker file '
|
|
'does not match encryption key SHA256 (%s) of component %s' %
|
|
(existing_enc_key_sha256, encryption_key_sha256,
|
|
component.object_name))
|
|
newly_completed_components = [component]
|
|
completed_components = existing_components + newly_completed_components
|
|
WriteParallelUploadTrackerFile(tracker_file_name,
|
|
prefix,
|
|
completed_components,
|
|
encryption_key_sha256=encryption_key_sha256)
|
|
|
|
|
|
def WriteParallelUploadTrackerFile(tracker_file_name,
|
|
prefix,
|
|
components,
|
|
encryption_key_sha256=None):
|
|
"""Writes information about components that were successfully uploaded.
|
|
|
|
The tracker file is serialized JSON of the form:
|
|
{
|
|
"encryption_key_sha256": sha256 hash of encryption key (or null),
|
|
"prefix": Prefix used for the component objects,
|
|
"components": [
|
|
{
|
|
"component_name": Component object name,
|
|
"component_generation": Component object generation (or null),
|
|
}, ...
|
|
]
|
|
}
|
|
where N is the number of components that have been successfully uploaded.
|
|
|
|
This function is not thread-safe and must be protected by a lock if
|
|
called within Command.Apply.
|
|
|
|
Args:
|
|
tracker_file_name: The name of the parallel upload tracker file.
|
|
prefix: The generated prefix that used for uploading any existing
|
|
components.
|
|
components: A list of ObjectFromTracker objects that were uploaded.
|
|
encryption_key_sha256: Encryption key SHA256 for use in this upload, if any.
|
|
"""
|
|
if six.PY3:
|
|
if isinstance(encryption_key_sha256, bytes):
|
|
encryption_key_sha256 = encryption_key_sha256.decode('ascii')
|
|
|
|
tracker_components = []
|
|
for component in components:
|
|
tracker_components.append({
|
|
_CompositeUploadTrackerEntry.COMPONENT_NAME: component.object_name,
|
|
_CompositeUploadTrackerEntry.COMPONENT_GENERATION: component.generation
|
|
})
|
|
tracker_file_data = {
|
|
_CompositeUploadTrackerEntry.COMPONENTS_LIST: tracker_components,
|
|
_CompositeUploadTrackerEntry.ENC_SHA256: encryption_key_sha256,
|
|
_CompositeUploadTrackerEntry.PREFIX: prefix
|
|
}
|
|
WriteJsonDataToTrackerFile(tracker_file_name, tracker_file_data)
|