# Copyright 2016 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. """Multiprocess file credential storage. This module provides file-based storage that supports multiple credentials and cross-thread and process access. This module supersedes the functionality previously found in `multistore_file`. This module provides :class:`MultiprocessFileStorage` which: * Is tied to a single credential via a user-specified key. This key can be used to distinguish between multiple users, client ids, and/or scopes. * Can be safely accessed and refreshed across threads and processes. Process & thread safety guarantees the following behavior: * If one thread or process refreshes a credential, subsequent refreshes from other processes will re-fetch the credentials from the file instead of performing an http request. * If two processes or threads attempt to refresh concurrently, only one will be able to acquire the lock and refresh, with the deadlock caveat below. * The interprocess lock will not deadlock, instead, the if a process can not acquire the interprocess lock within ``INTERPROCESS_LOCK_DEADLINE`` it will allow refreshing the credential but will not write the updated credential to disk, This logic happens during every lock cycle - if the credentials are refreshed again it will retry locking and writing as normal. Usage ===== Before using the storage, you need to decide how you want to key the credentials. A few common strategies include: * If you're storing credentials for multiple users in a single file, use a unique identifier for each user as the key. * If you're storing credentials for multiple client IDs in a single file, use the client ID as the key. * If you're storing multiple credentials for one user, use the scopes as the key. * If you have a complicated setup, use a compound key. For example, you can use a combination of the client ID and scopes as the key. Create an instance of :class:`MultiprocessFileStorage` for each credential you want to store, for example:: filename = 'credentials' key = '{}-{}'.format(client_id, user_id) storage = MultiprocessFileStorage(filename, key) To store the credentials:: storage.put(credentials) If you're going to continue to use the credentials after storing them, be sure to call :func:`set_store`:: credentials.set_store(storage) To retrieve the credentials:: storage.get(credentials) """ import base64 import json import logging import os import threading import fasteners from six import iteritems from oauth2client import _helpers from oauth2client import client #: The maximum amount of time, in seconds, to wait when acquire the #: interprocess lock before falling back to read-only mode. INTERPROCESS_LOCK_DEADLINE = 1 logger = logging.getLogger(__name__) _backends = {} _backends_lock = threading.Lock() def _create_file_if_needed(filename): """Creates the an empty file if it does not already exist. Returns: True if the file was created, False otherwise. """ if os.path.exists(filename): return False else: # Equivalent to "touch". open(filename, 'a+b').close() logger.info('Credential file {0} created'.format(filename)) return True def _load_credentials_file(credentials_file): """Load credentials from the given file handle. The file is expected to be in this format: { "file_version": 2, "credentials": { "key": "base64 encoded json representation of credentials." } } This function will warn and return empty credentials instead of raising exceptions. Args: credentials_file: An open file handle. Returns: A dictionary mapping user-defined keys to an instance of :class:`oauth2client.client.Credentials`. """ try: credentials_file.seek(0) data = json.load(credentials_file) except Exception: logger.warning( 'Credentials file could not be loaded, will ignore and ' 'overwrite.') return {} if data.get('file_version') != 2: logger.warning( 'Credentials file is not version 2, will ignore and ' 'overwrite.') return {} credentials = {} for key, encoded_credential in iteritems(data.get('credentials', {})): try: credential_json = base64.b64decode(encoded_credential) credential = client.Credentials.new_from_json(credential_json) credentials[key] = credential except: logger.warning( 'Invalid credential {0} in file, ignoring.'.format(key)) return credentials def _write_credentials_file(credentials_file, credentials): """Writes credentials to a file. Refer to :func:`_load_credentials_file` for the format. Args: credentials_file: An open file handle, must be read/write. credentials: A dictionary mapping user-defined keys to an instance of :class:`oauth2client.client.Credentials`. """ data = {'file_version': 2, 'credentials': {}} for key, credential in iteritems(credentials): credential_json = credential.to_json() encoded_credential = _helpers._from_bytes(base64.b64encode( _helpers._to_bytes(credential_json))) data['credentials'][key] = encoded_credential credentials_file.seek(0) json.dump(data, credentials_file) credentials_file.truncate() class _MultiprocessStorageBackend(object): """Thread-local backend for multiprocess storage. Each process has only one instance of this backend per file. All threads share a single instance of this backend. This ensures that all threads use the same thread lock and process lock when accessing the file. """ def __init__(self, filename): self._file = None self._filename = filename self._process_lock = fasteners.InterProcessLock( '{0}.lock'.format(filename)) self._thread_lock = threading.Lock() self._read_only = False self._credentials = {} def _load_credentials(self): """(Re-)loads the credentials from the file.""" if not self._file: return loaded_credentials = _load_credentials_file(self._file) self._credentials.update(loaded_credentials) logger.debug('Read credential file') def _write_credentials(self): if self._read_only: logger.debug('In read-only mode, not writing credentials.') return _write_credentials_file(self._file, self._credentials) logger.debug('Wrote credential file {0}.'.format(self._filename)) def acquire_lock(self): self._thread_lock.acquire() locked = self._process_lock.acquire(timeout=INTERPROCESS_LOCK_DEADLINE) if locked: _create_file_if_needed(self._filename) self._file = open(self._filename, 'r+') self._read_only = False else: logger.warn( 'Failed to obtain interprocess lock for credentials. ' 'If a credential is being refreshed, other processes may ' 'not see the updated access token and refresh as well.') if os.path.exists(self._filename): self._file = open(self._filename, 'r') else: self._file = None self._read_only = True self._load_credentials() def release_lock(self): if self._file is not None: self._file.close() self._file = None if not self._read_only: self._process_lock.release() self._thread_lock.release() def _refresh_predicate(self, credentials): if credentials is None: return True elif credentials.invalid: return True elif credentials.access_token_expired: return True else: return False def locked_get(self, key): # Check if the credential is already in memory. credentials = self._credentials.get(key, None) # Use the refresh predicate to determine if the entire store should be # reloaded. This basically checks if the credentials are invalid # or expired. This covers the situation where another process has # refreshed the credentials and this process doesn't know about it yet. # In that case, this process won't needlessly refresh the credentials. if self._refresh_predicate(credentials): self._load_credentials() credentials = self._credentials.get(key, None) return credentials def locked_put(self, key, credentials): self._load_credentials() self._credentials[key] = credentials self._write_credentials() def locked_delete(self, key): self._load_credentials() self._credentials.pop(key, None) self._write_credentials() def _get_backend(filename): """A helper method to get or create a backend with thread locking. This ensures that only one backend is used per-file per-process, so that thread and process locks are appropriately shared. Args: filename: The full path to the credential storage file. Returns: An instance of :class:`_MultiprocessStorageBackend`. """ filename = os.path.abspath(filename) with _backends_lock: if filename not in _backends: _backends[filename] = _MultiprocessStorageBackend(filename) return _backends[filename] class MultiprocessFileStorage(client.Storage): """Multiprocess file credential storage. Args: filename: The path to the file where credentials will be stored. key: An arbitrary string used to uniquely identify this set of credentials. For example, you may use the user's ID as the key or a combination of the client ID and user ID. """ def __init__(self, filename, key): self._key = key self._backend = _get_backend(filename) def acquire_lock(self): self._backend.acquire_lock() def release_lock(self): self._backend.release_lock() def locked_get(self): """Retrieves the current credentials from the store. Returns: An instance of :class:`oauth2client.client.Credentials` or `None`. """ credential = self._backend.locked_get(self._key) if credential is not None: credential.set_store(self) return credential def locked_put(self, credentials): """Writes the given credentials to the store. Args: credentials: an instance of :class:`oauth2client.client.Credentials`. """ return self._backend.locked_put(self._key, credentials) def locked_delete(self): """Deletes the current credentials from the store.""" return self._backend.locked_delete(self._key)