1# Copyright 2016 Google Inc. All rights reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15"""Multiprocess file credential storage. 16 17This module provides file-based storage that supports multiple credentials and 18cross-thread and process access. 19 20This module supersedes the functionality previously found in `multistore_file`. 21 22This module provides :class:`MultiprocessFileStorage` which: 23 * Is tied to a single credential via a user-specified key. This key can be 24 used to distinguish between multiple users, client ids, and/or scopes. 25 * Can be safely accessed and refreshed across threads and processes. 26 27Process & thread safety guarantees the following behavior: 28 * If one thread or process refreshes a credential, subsequent refreshes 29 from other processes will re-fetch the credentials from the file instead 30 of performing an http request. 31 * If two processes or threads attempt to refresh concurrently, only one 32 will be able to acquire the lock and refresh, with the deadlock caveat 33 below. 34 * The interprocess lock will not deadlock, instead, the if a process can 35 not acquire the interprocess lock within ``INTERPROCESS_LOCK_DEADLINE`` 36 it will allow refreshing the credential but will not write the updated 37 credential to disk, This logic happens during every lock cycle - if the 38 credentials are refreshed again it will retry locking and writing as 39 normal. 40 41Usage 42===== 43 44Before using the storage, you need to decide how you want to key the 45credentials. A few common strategies include: 46 47 * If you're storing credentials for multiple users in a single file, use 48 a unique identifier for each user as the key. 49 * If you're storing credentials for multiple client IDs in a single file, 50 use the client ID as the key. 51 * If you're storing multiple credentials for one user, use the scopes as 52 the key. 53 * If you have a complicated setup, use a compound key. For example, you 54 can use a combination of the client ID and scopes as the key. 55 56Create an instance of :class:`MultiprocessFileStorage` for each credential you 57want to store, for example:: 58 59 filename = 'credentials' 60 key = '{}-{}'.format(client_id, user_id) 61 storage = MultiprocessFileStorage(filename, key) 62 63To store the credentials:: 64 65 storage.put(credentials) 66 67If you're going to continue to use the credentials after storing them, be sure 68to call :func:`set_store`:: 69 70 credentials.set_store(storage) 71 72To retrieve the credentials:: 73 74 storage.get(credentials) 75 76""" 77 78import base64 79import json 80import logging 81import os 82import threading 83 84import fasteners 85from six import iteritems 86 87from oauth2client import _helpers 88from oauth2client import client 89 90 91#: The maximum amount of time, in seconds, to wait when acquire the 92#: interprocess lock before falling back to read-only mode. 93INTERPROCESS_LOCK_DEADLINE = 1 94 95logger = logging.getLogger(__name__) 96_backends = {} 97_backends_lock = threading.Lock() 98 99 100def _create_file_if_needed(filename): 101 """Creates the an empty file if it does not already exist. 102 103 Returns: 104 True if the file was created, False otherwise. 105 """ 106 if os.path.exists(filename): 107 return False 108 else: 109 # Equivalent to "touch". 110 open(filename, 'a+b').close() 111 logger.info('Credential file {0} created'.format(filename)) 112 return True 113 114 115def _load_credentials_file(credentials_file): 116 """Load credentials from the given file handle. 117 118 The file is expected to be in this format: 119 120 { 121 "file_version": 2, 122 "credentials": { 123 "key": "base64 encoded json representation of credentials." 124 } 125 } 126 127 This function will warn and return empty credentials instead of raising 128 exceptions. 129 130 Args: 131 credentials_file: An open file handle. 132 133 Returns: 134 A dictionary mapping user-defined keys to an instance of 135 :class:`oauth2client.client.Credentials`. 136 """ 137 try: 138 credentials_file.seek(0) 139 data = json.load(credentials_file) 140 except Exception: 141 logger.warning( 142 'Credentials file could not be loaded, will ignore and ' 143 'overwrite.') 144 return {} 145 146 if data.get('file_version') != 2: 147 logger.warning( 148 'Credentials file is not version 2, will ignore and ' 149 'overwrite.') 150 return {} 151 152 credentials = {} 153 154 for key, encoded_credential in iteritems(data.get('credentials', {})): 155 try: 156 credential_json = base64.b64decode(encoded_credential) 157 credential = client.Credentials.new_from_json(credential_json) 158 credentials[key] = credential 159 except: 160 logger.warning( 161 'Invalid credential {0} in file, ignoring.'.format(key)) 162 163 return credentials 164 165 166def _write_credentials_file(credentials_file, credentials): 167 """Writes credentials to a file. 168 169 Refer to :func:`_load_credentials_file` for the format. 170 171 Args: 172 credentials_file: An open file handle, must be read/write. 173 credentials: A dictionary mapping user-defined keys to an instance of 174 :class:`oauth2client.client.Credentials`. 175 """ 176 data = {'file_version': 2, 'credentials': {}} 177 178 for key, credential in iteritems(credentials): 179 credential_json = credential.to_json() 180 encoded_credential = _helpers._from_bytes(base64.b64encode( 181 _helpers._to_bytes(credential_json))) 182 data['credentials'][key] = encoded_credential 183 184 credentials_file.seek(0) 185 json.dump(data, credentials_file) 186 credentials_file.truncate() 187 188 189class _MultiprocessStorageBackend(object): 190 """Thread-local backend for multiprocess storage. 191 192 Each process has only one instance of this backend per file. All threads 193 share a single instance of this backend. This ensures that all threads 194 use the same thread lock and process lock when accessing the file. 195 """ 196 197 def __init__(self, filename): 198 self._file = None 199 self._filename = filename 200 self._process_lock = fasteners.InterProcessLock( 201 '{0}.lock'.format(filename)) 202 self._thread_lock = threading.Lock() 203 self._read_only = False 204 self._credentials = {} 205 206 def _load_credentials(self): 207 """(Re-)loads the credentials from the file.""" 208 if not self._file: 209 return 210 211 loaded_credentials = _load_credentials_file(self._file) 212 self._credentials.update(loaded_credentials) 213 214 logger.debug('Read credential file') 215 216 def _write_credentials(self): 217 if self._read_only: 218 logger.debug('In read-only mode, not writing credentials.') 219 return 220 221 _write_credentials_file(self._file, self._credentials) 222 logger.debug('Wrote credential file {0}.'.format(self._filename)) 223 224 def acquire_lock(self): 225 self._thread_lock.acquire() 226 locked = self._process_lock.acquire(timeout=INTERPROCESS_LOCK_DEADLINE) 227 228 if locked: 229 _create_file_if_needed(self._filename) 230 self._file = open(self._filename, 'r+') 231 self._read_only = False 232 233 else: 234 logger.warn( 235 'Failed to obtain interprocess lock for credentials. ' 236 'If a credential is being refreshed, other processes may ' 237 'not see the updated access token and refresh as well.') 238 if os.path.exists(self._filename): 239 self._file = open(self._filename, 'r') 240 else: 241 self._file = None 242 self._read_only = True 243 244 self._load_credentials() 245 246 def release_lock(self): 247 if self._file is not None: 248 self._file.close() 249 self._file = None 250 251 if not self._read_only: 252 self._process_lock.release() 253 254 self._thread_lock.release() 255 256 def _refresh_predicate(self, credentials): 257 if credentials is None: 258 return True 259 elif credentials.invalid: 260 return True 261 elif credentials.access_token_expired: 262 return True 263 else: 264 return False 265 266 def locked_get(self, key): 267 # Check if the credential is already in memory. 268 credentials = self._credentials.get(key, None) 269 270 # Use the refresh predicate to determine if the entire store should be 271 # reloaded. This basically checks if the credentials are invalid 272 # or expired. This covers the situation where another process has 273 # refreshed the credentials and this process doesn't know about it yet. 274 # In that case, this process won't needlessly refresh the credentials. 275 if self._refresh_predicate(credentials): 276 self._load_credentials() 277 credentials = self._credentials.get(key, None) 278 279 return credentials 280 281 def locked_put(self, key, credentials): 282 self._load_credentials() 283 self._credentials[key] = credentials 284 self._write_credentials() 285 286 def locked_delete(self, key): 287 self._load_credentials() 288 self._credentials.pop(key, None) 289 self._write_credentials() 290 291 292def _get_backend(filename): 293 """A helper method to get or create a backend with thread locking. 294 295 This ensures that only one backend is used per-file per-process, so that 296 thread and process locks are appropriately shared. 297 298 Args: 299 filename: The full path to the credential storage file. 300 301 Returns: 302 An instance of :class:`_MultiprocessStorageBackend`. 303 """ 304 filename = os.path.abspath(filename) 305 306 with _backends_lock: 307 if filename not in _backends: 308 _backends[filename] = _MultiprocessStorageBackend(filename) 309 return _backends[filename] 310 311 312class MultiprocessFileStorage(client.Storage): 313 """Multiprocess file credential storage. 314 315 Args: 316 filename: The path to the file where credentials will be stored. 317 key: An arbitrary string used to uniquely identify this set of 318 credentials. For example, you may use the user's ID as the key or 319 a combination of the client ID and user ID. 320 """ 321 def __init__(self, filename, key): 322 self._key = key 323 self._backend = _get_backend(filename) 324 325 def acquire_lock(self): 326 self._backend.acquire_lock() 327 328 def release_lock(self): 329 self._backend.release_lock() 330 331 def locked_get(self): 332 """Retrieves the current credentials from the store. 333 334 Returns: 335 An instance of :class:`oauth2client.client.Credentials` or `None`. 336 """ 337 credential = self._backend.locked_get(self._key) 338 339 if credential is not None: 340 credential.set_store(self) 341 342 return credential 343 344 def locked_put(self, credentials): 345 """Writes the given credentials to the store. 346 347 Args: 348 credentials: an instance of 349 :class:`oauth2client.client.Credentials`. 350 """ 351 return self._backend.locked_put(self._key, credentials) 352 353 def locked_delete(self): 354 """Deletes the current credentials from the store.""" 355 return self._backend.locked_delete(self._key) 356