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.
IsFile(const string & filename)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.
GetEnvironmentVariableFileName(string * filename)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'.
GetWellKnownFileName(string * filename)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
GoogleAuthProvider()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
GoogleAuthProvider(std::unique_ptr<OAuthClient> oauth_client,std::unique_ptr<HttpRequest::Factory> http_request_factory,Env * env,int64 initial_retry_delay_usec)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
GetToken(string * t)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
GetTokenFromFiles()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, ¤t_token_, &expiration_timestamp_sec_));
198 } else if (json.isMember("private_key")) {
199 TF_RETURN_IF_ERROR(oauth_client_->GetTokenFromServiceAccountJson(
200 json, kOAuthV4Url, kOAuthScope, ¤t_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
GetTokenFromGce()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, ¤t_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
GetTokenForTesting()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