Home | History | Annotate | Download | only in cloud 
      1 /* Copyright 2016 The TensorFlow Authors. 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 
     16 #include "tensorflow/core/platform/cloud/google_auth_provider.h"
     17 #ifndef _WIN32
     18 #include <pwd.h>
     19 #include <unistd.h>
     20 #else
     21 #include <sys/types.h>
     22 #endif
     23 #include <fstream>
     24 #include "include/json/json.h"
     25 #include "tensorflow/core/lib/core/errors.h"
     26 #include "tensorflow/core/lib/io/path.h"
     27 #include "tensorflow/core/lib/strings/base64.h"
     28 #include "tensorflow/core/platform/cloud/curl_http_request.h"
     29 #include "tensorflow/core/platform/cloud/retrying_utils.h"
     30 #include "tensorflow/core/platform/env.h"
     31 
     32 namespace tensorflow {
     33 
     34 namespace {
     35 
     36 // The environment variable pointing to the file with local
     37 // Application Default Credentials.
     38 constexpr char kGoogleApplicationCredentials[] =
     39     "GOOGLE_APPLICATION_CREDENTIALS";
     40 
     41 // The environment variable to override token generation for testing.
     42 constexpr char kGoogleAuthTokenForTesting[] = "GOOGLE_AUTH_TOKEN_FOR_TESTING";
     43 
     44 // The environment variable which can override '~/.config/gcloud' if set.
     45 constexpr char kCloudSdkConfig[] = "CLOUDSDK_CONFIG";
     46 
     47 // The default path to the gcloud config folder, relative to the home folder.
     48 constexpr char kGCloudConfigFolder[] = ".config/gcloud/";
     49 
     50 // The name of the well-known credentials JSON file in the gcloud config folder.
     51 constexpr char kWellKnownCredentialsFile[] =
     52     "application_default_credentials.json";
     53 
     54 // The minimum time delta between now and the token expiration time
     55 // for the token to be re-used.
     56 constexpr int kExpirationTimeMarginSec = 60;
     57 
     58 // The URL to retrieve the auth bearer token via OAuth with a refresh token.
     59 constexpr char kOAuthV3Url[] = "https://www.googleapis.com/oauth2/v3/token";
     60 
     61 // The URL to retrieve the auth bearer token via OAuth with a private key.
     62 constexpr char kOAuthV4Url[] = "https://www.googleapis.com/oauth2/v4/token";
     63 
     64 // The URL to retrieve the auth bearer token when running in Google Compute
     65 // Engine.
     66 constexpr char kGceTokenUrl[] =
     67     "http://metadata/computeMetadata/v1/instance/service-accounts/default/"
     68     "token";
     69 
     70 // The authentication token scope to request.
     71 constexpr char kOAuthScope[] = "https://www.googleapis.com/auth/cloud-platform";
     72 
     73 // The default initial delay between retries with exponential backoff.
     74 constexpr int kInitialRetryDelayUsec = 500000;  // 0.5 sec
     75 
     76 /// Returns whether the given path points to a readable file.
     77 bool IsFile(const string& filename) {
     78   std::ifstream fstream(filename.c_str());
     79   return fstream.good();
     80 }
     81 
     82 /// Returns the credentials file name from the env variable.
     83 Status GetEnvironmentVariableFileName(string* filename) {
     84   if (!filename) {
     85     return errors::FailedPrecondition("'filename' cannot be nullptr.");
     86   }
     87   const char* result = std::getenv(kGoogleApplicationCredentials);
     88   if (!result || !IsFile(result)) {
     89     return errors::NotFound(strings::StrCat("$", kGoogleApplicationCredentials,
     90                                             " is not set or corrupt."));
     91   }
     92   *filename = result;
     93   return Status::OK();
     94 }
     95 
     96 /// Returns the well known file produced by command 'gcloud auth login'.
     97 Status GetWellKnownFileName(string* filename) {
     98   if (!filename) {
     99     return errors::FailedPrecondition("'filename' cannot be nullptr.");
    100   }
    101   string config_dir;
    102   const char* config_dir_override = std::getenv(kCloudSdkConfig);
    103   if (config_dir_override) {
    104     config_dir = config_dir_override;
    105   } else {
    106     // Determine the home dir path.
    107     const char* home_dir = std::getenv("HOME");
    108     if (!home_dir) {
    109       return errors::FailedPrecondition("Could not read $HOME.");
    110     }
    111     config_dir = io::JoinPath(home_dir, kGCloudConfigFolder);
    112   }
    113   auto result = io::JoinPath(config_dir, kWellKnownCredentialsFile);
    114   if (!IsFile(result)) {
    115     return errors::NotFound(
    116         "Could not find the credentials file in the standard gcloud location.");
    117   }
    118   *filename = result;
    119   return Status::OK();
    120 }
    121 
    122 }  // namespace
    123 
    124 GoogleAuthProvider::GoogleAuthProvider()
    125     : GoogleAuthProvider(
    126           std::unique_ptr<OAuthClient>(new OAuthClient()),
    127           std::unique_ptr<HttpRequest::Factory>(new CurlHttpRequest::Factory()),
    128           Env::Default(), kInitialRetryDelayUsec) {}
    129 
    130 GoogleAuthProvider::GoogleAuthProvider(
    131     std::unique_ptr<OAuthClient> oauth_client,
    132     std::unique_ptr<HttpRequest::Factory> http_request_factory, Env* env,
    133     int64 initial_retry_delay_usec)
    134     : oauth_client_(std::move(oauth_client)),
    135       http_request_factory_(std::move(http_request_factory)),
    136       env_(env),
    137       initial_retry_delay_usec_(initial_retry_delay_usec) {}
    138 
    139 Status GoogleAuthProvider::GetToken(string* t) {
    140   mutex_lock lock(mu_);
    141   const uint64 now_sec = env_->NowSeconds();
    142 
    143   if (!current_token_.empty() &&
    144       now_sec + kExpirationTimeMarginSec < expiration_timestamp_sec_) {
    145     *t = current_token_;
    146     return Status::OK();
    147   }
    148 
    149   if (GetTokenForTesting().ok()) {
    150     *t = current_token_;
    151     return Status::OK();
    152   }
    153 
    154   auto token_from_files_status = GetTokenFromFiles();
    155   auto token_from_gce_status =
    156       token_from_files_status.ok() ? Status::OK() : GetTokenFromGce();
    157 
    158   if (token_from_files_status.ok() || token_from_gce_status.ok()) {
    159     *t = current_token_;
    160     return Status::OK();
    161   }
    162 
    163   LOG(WARNING)
    164       << "All attempts to get a Google authentication bearer token failed, "
    165       << "returning an empty token. Retrieving token from files failed with \""
    166       << token_from_files_status.ToString() << "\"."
    167       << " Retrieving token from GCE failed with \""
    168       << token_from_gce_status.ToString() << "\".";
    169 
    170   // Public objects can still be accessed with an empty bearer token,
    171   // so return an empty token instead of failing.
    172   *t = "";
    173 
    174   // From now on, always return the empty token.
    175   expiration_timestamp_sec_ = UINT64_MAX;
    176   current_token_ = "";
    177 
    178   return Status::OK();
    179 }
    180 
    181 Status GoogleAuthProvider::GetTokenFromFiles() {
    182   string credentials_filename;
    183   if (!GetEnvironmentVariableFileName(&credentials_filename).ok() &&
    184       !GetWellKnownFileName(&credentials_filename).ok()) {
    185     return errors::NotFound("Could not locate the credentials file.");
    186   }
    187 
    188   Json::Value json;
    189   Json::Reader reader;
    190   std::ifstream credentials_fstream(credentials_filename);
    191   if (!reader.parse(credentials_fstream, json)) {
    192     return errors::FailedPrecondition(
    193         "Couldn't parse the JSON credentials file.");
    194   }
    195   if (json.isMember("refresh_token")) {
    196     TF_RETURN_IF_ERROR(oauth_client_->GetTokenFromRefreshTokenJson(
    197         json, kOAuthV3Url, &current_token_, &expiration_timestamp_sec_));
    198   } else if (json.isMember("private_key")) {
    199     TF_RETURN_IF_ERROR(oauth_client_->GetTokenFromServiceAccountJson(
    200         json, kOAuthV4Url, kOAuthScope, &current_token_,
    201         &expiration_timestamp_sec_));
    202   } else {
    203     return errors::FailedPrecondition(
    204         "Unexpected content of the JSON credentials file.");
    205   }
    206   return Status::OK();
    207 }
    208 
    209 Status GoogleAuthProvider::GetTokenFromGce() {
    210   const auto get_token_from_gce = [this]() {
    211     std::unique_ptr<HttpRequest> request(http_request_factory_->Create());
    212     std::vector<char> response_buffer;
    213     const uint64 request_timestamp_sec = env_->NowSeconds();
    214     request->SetUri(kGceTokenUrl);
    215     request->AddHeader("Metadata-Flavor", "Google");
    216     request->SetResultBuffer(&response_buffer);
    217     TF_RETURN_IF_ERROR(request->Send());
    218     StringPiece response =
    219         StringPiece(&response_buffer[0], response_buffer.size());
    220 
    221     TF_RETURN_IF_ERROR(oauth_client_->ParseOAuthResponse(
    222         response, request_timestamp_sec, &current_token_,
    223         &expiration_timestamp_sec_));
    224     return Status::OK();
    225   };
    226   return RetryingUtils::CallWithRetries(get_token_from_gce,
    227                                         initial_retry_delay_usec_);
    228 }
    229 
    230 Status GoogleAuthProvider::GetTokenForTesting() {
    231   const char* token = std::getenv(kGoogleAuthTokenForTesting);
    232   if (!token) {
    233     return errors::NotFound("The env variable for testing was not set.");
    234   }
    235   expiration_timestamp_sec_ = UINT64_MAX;
    236   current_token_ = token;
    237   return Status::OK();
    238 }
    239 
    240 }  // namespace tensorflow
    241