1 /*
2 * Copyright 2022 Google LLC
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16 #include "fcp/client/federated_select.h"
17
18 #include <fstream>
19 #include <memory>
20 #include <sstream>
21 #include <string>
22 #include <vector>
23
24 #include "google/protobuf/any.pb.h"
25 #include "google/protobuf/text_format.h"
26 #include "gmock/gmock.h"
27 #include "gtest/gtest.h"
28 #include "absl/status/status.h"
29 #include "absl/synchronization/blocking_counter.h"
30 #include "absl/time/time.h"
31 #include "fcp/base/monitoring.h"
32 #include "fcp/client/client_runner.h"
33 #include "fcp/client/diag_codes.pb.h"
34 #include "fcp/client/engine/example_iterator_factory.h"
35 #include "fcp/client/http/http_client.h"
36 #include "fcp/client/http/in_memory_request_response.h"
37 #include "fcp/client/http/testing/test_helpers.h"
38 #include "fcp/client/interruptible_runner.h"
39 #include "fcp/client/stats.h"
40 #include "fcp/client/test_helpers.h"
41 #include "fcp/protos/plan.pb.h"
42 #include "fcp/testing/testing.h"
43
44 namespace fcp::client {
45 namespace {
46
47 using ::fcp::IsCode;
48 using ::fcp::client::ExampleIterator;
49 using ::fcp::client::engine::ExampleIteratorFactory;
50 using ::fcp::client::http::FakeHttpResponse;
51 using ::fcp::client::http::HeaderList;
52 using ::fcp::client::http::HttpRequest;
53 using ::fcp::client::http::HttpRequestHandle;
54 using ::fcp::client::http::MockHttpClient;
55 using ::fcp::client::http::SimpleHttpRequestMatcher;
56 using ::fcp::client::http::internal::CompressWithGzip;
57 using ::google::internal::federated::plan::ExampleSelector;
58 using ::google::internal::federated::plan::SlicesSelector;
59 using ::testing::_;
60 using ::testing::Gt;
61 using ::testing::HasSubstr;
62 using ::testing::InSequence;
63 using ::testing::MockFunction;
64 using ::testing::NiceMock;
65 using ::testing::Not;
66 using ::testing::Return;
67 using ::testing::StrictMock;
68
CreateExampleSelector(const std::string & served_at_id,std::vector<int32_t> keys)69 ExampleSelector CreateExampleSelector(const std::string& served_at_id,
70 std::vector<int32_t> keys) {
71 ExampleSelector example_selector;
72 *example_selector.mutable_collection_uri() = "internal:/federated_select";
73 SlicesSelector slices_selector;
74 *slices_selector.mutable_served_at_id() = served_at_id;
75 slices_selector.mutable_keys()->Add(keys.begin(), keys.end());
76 example_selector.mutable_criteria()->PackFrom(slices_selector);
77 return example_selector;
78 }
79
FileExists(const std::string & path)80 bool FileExists(const std::string& path) {
81 std::ifstream istream(path);
82 return istream.good();
83 }
84
ReadFile(const std::string & path)85 std::string ReadFile(const std::string& path) {
86 std::ifstream istream(path);
87 FCP_CHECK(istream);
88 std::stringstream stringstream;
89 stringstream << istream.rdbuf();
90 return stringstream.str();
91 }
92
93 class HttpFederatedSelectManagerTest : public ::testing::Test {
94 protected:
HttpFederatedSelectManagerTest()95 HttpFederatedSelectManagerTest()
96 : fedselect_manager_(
97 &mock_log_manager_, &files_impl_, &mock_http_client_,
98 mock_should_abort_.AsStdFunction(),
99 InterruptibleRunner::TimingConfig{
100 .polling_period = absl::ZeroDuration(),
101 .graceful_shutdown_period = absl::InfiniteDuration(),
102 .extended_shutdown_period = absl::InfiniteDuration()}) {}
103
SetUp()104 void SetUp() override {
105 EXPECT_CALL(mock_flags_, enable_federated_select())
106 .WillRepeatedly(Return(true));
107 }
108
TearDown()109 void TearDown() override {
110 // Regardless of the outcome of the test (or the protocol interaction being
111 // tested), network usage must always be reflected in the network stats
112 // methods.
113 HttpRequestHandle::SentReceivedBytes sent_received_bytes =
114 mock_http_client_.TotalSentReceivedBytes();
115 NetworkStats network_stats = fedselect_manager_.GetNetworkStats();
116 EXPECT_EQ(network_stats.bytes_downloaded,
117 sent_received_bytes.received_bytes);
118 EXPECT_EQ(network_stats.bytes_uploaded, sent_received_bytes.sent_bytes);
119 // If any network traffic occurred, we expect to see some time reflected in
120 // the duration.
121 if (network_stats.bytes_uploaded > 0) {
122 EXPECT_THAT(network_stats.network_duration, Gt(absl::ZeroDuration()));
123 }
124 }
125
126 NiceMock<MockLogManager> mock_log_manager_;
127 MockFlags mock_flags_;
128 fcp::client::FilesImpl files_impl_;
129 StrictMock<MockHttpClient> mock_http_client_;
130 NiceMock<MockFunction<bool()>> mock_should_abort_;
131
132 HttpFederatedSelectManager fedselect_manager_;
133 };
134
TEST_F(HttpFederatedSelectManagerTest,IteratorFactoryShouldHandleValidSelector)135 TEST_F(HttpFederatedSelectManagerTest,
136 IteratorFactoryShouldHandleValidSelector) {
137 // Should be handled by the factory.
138 ExampleSelector selector =
139 CreateExampleSelector(/*served_at_id=*/"foo", /*keys=*/{1});
140
141 std::unique_ptr<ExampleIteratorFactory> iterator_factory =
142 fedselect_manager_.CreateExampleIteratorFactoryForUriTemplate(
143 "https://foo.bar");
144
145 EXPECT_TRUE(iterator_factory->CanHandle(selector));
146 }
147
TEST_F(HttpFederatedSelectManagerTest,IteratorFactoryShouldNotHandleUnrelatedSelector)148 TEST_F(HttpFederatedSelectManagerTest,
149 IteratorFactoryShouldNotHandleUnrelatedSelector) {
150 // Should not be handled by the factory.
151 ExampleSelector selector;
152 *selector.mutable_collection_uri() = "internal:/foo";
153
154 std::unique_ptr<ExampleIteratorFactory> iterator_factory =
155 fedselect_manager_.CreateExampleIteratorFactoryForUriTemplate(
156 "https://foo.bar");
157
158 EXPECT_FALSE(iterator_factory->CanHandle(selector));
159 EXPECT_THAT(iterator_factory->CreateExampleIterator(selector),
160 IsCode(INVALID_ARGUMENT));
161 }
162
TEST_F(HttpFederatedSelectManagerTest,IteratorFactoryShouldNotCollectStats)163 TEST_F(HttpFederatedSelectManagerTest, IteratorFactoryShouldNotCollectStats) {
164 std::unique_ptr<ExampleIteratorFactory> iterator_factory =
165 fedselect_manager_.CreateExampleIteratorFactoryForUriTemplate(
166 "https://foo.bar");
167
168 EXPECT_FALSE(iterator_factory->ShouldCollectStats());
169 }
170
TEST_F(HttpFederatedSelectManagerTest,EmptyUriTemplateShouldFailAllExampleQueries)171 TEST_F(HttpFederatedSelectManagerTest,
172 EmptyUriTemplateShouldFailAllExampleQueries) {
173 ExampleSelector selector =
174 CreateExampleSelector(/*served_at_id=*/"foo", /*keys=*/{1});
175
176 EXPECT_CALL(
177 mock_log_manager_,
178 LogDiag(ProdDiagCode::FEDSELECT_SLICE_HTTP_FETCH_REQUESTED_BUT_DISABLED));
179
180 // Create an iterator factory using an empty base URI (indicating that the
181 // server didn't provide us with a federated select URI template, i.e. the
182 // feature is disabled on the server or the plan doesn't use the Federated
183 // Select feature).
184 std::unique_ptr<ExampleIteratorFactory> iterator_factory =
185 fedselect_manager_.CreateExampleIteratorFactoryForUriTemplate(
186 /*uri_template=*/"");
187 absl::StatusOr<std::unique_ptr<ExampleIterator>> iterator =
188 iterator_factory->CreateExampleIterator(selector);
189
190 // The iterator creation should have failed, since the URI template was empty.
191 EXPECT_THAT(iterator, IsCode(INVALID_ARGUMENT));
192 EXPECT_THAT(iterator.status().message(), HasSubstr("disabled"));
193 // Even though creating an iterator should fail since the feature is disabled,
194 // the iterator factory should still handle all internal:/federated_select
195 // example queries (rather than let them bubble up to the default
196 // environment-provided example iterator factory, which won't know how to
197 // handle them anyway).
198 EXPECT_TRUE(iterator_factory->CanHandle(selector));
199 }
200
201 /** Tests the "happy" path. Two separate federated select slice fetching queries
202 * are received by a single iterator factory, each serving different slice data
203 * to the client. All slice fetches are successful and should result in the
204 * correct data being returned, in the right order.*/
TEST_F(HttpFederatedSelectManagerTest,SuccessfullyFetchMultipleSlicesAcrossMultipleIterators)205 TEST_F(HttpFederatedSelectManagerTest,
206 SuccessfullyFetchMultipleSlicesAcrossMultipleIterators) {
207 const std::string uri_template =
208 "https://foo.bar/{served_at_id}/baz/{key_base10}/bazz";
209
210 // Create an iterator factory with a valid (non-empty) URI template.
211 std::unique_ptr<ExampleIteratorFactory> iterator_factory =
212 fedselect_manager_.CreateExampleIteratorFactoryForUriTemplate(
213 uri_template);
214
215 // Once the first iterator is created we expect the following slice fetch
216 // requests to be issued immediately.
217 const std::string expected_key1_data = "key1_data";
218 EXPECT_CALL(
219 mock_http_client_,
220 PerformSingleRequest(SimpleHttpRequestMatcher(
221 // Note that the '/' in "id/X" is *not* URI-escaped.
222 "https://foo.bar/id/X/baz/1/bazz", HttpRequest::Method::kGet, _, "")))
223 .WillOnce(
224 Return(FakeHttpResponse(200, HeaderList(), expected_key1_data)));
225
226 const std::string expected_key2_data = "key2_data";
227 EXPECT_CALL(mock_http_client_, PerformSingleRequest(SimpleHttpRequestMatcher(
228 "https://foo.bar/id/X/baz/2/bazz",
229 HttpRequest::Method::kGet, _, "")))
230 .WillOnce(
231 Return(FakeHttpResponse(200, HeaderList(), expected_key2_data)));
232
233 {
234 InSequence in_sequence;
235 EXPECT_CALL(mock_log_manager_,
236 LogDiag(ProdDiagCode::FEDSELECT_SLICE_HTTP_FETCH_REQUESTED));
237 EXPECT_CALL(mock_log_manager_,
238 LogDiag(ProdDiagCode::FEDSELECT_SLICE_HTTP_FETCH_SUCCEEDED));
239 }
240
241 absl::StatusOr<std::unique_ptr<ExampleIterator>> iterator1 =
242 iterator_factory->CreateExampleIterator(CreateExampleSelector(
243 // Note that we use a served_at_id value with a '/' in it. It should
244 // *not* get URI-escaped, as per the FederatedSelectUriInfo docs.
245 /*served_at_id=*/"id/X",
246 // Also note that we request slices 2 and 1, in that exact order. I.e.
247 // the first slice data we receive should be slice 2, and the second
248 // slice data we receive should be slice 1.
249 /*keys=*/{2, 1}));
250
251 // The iterator creation should have succeeded.
252 ASSERT_OK(iterator1);
253
254 // Reading the data for each of the slices should now succeed.
255 absl::StatusOr<std::string> first_slice = (*iterator1)->Next();
256 ASSERT_OK(first_slice);
257 ASSERT_TRUE(FileExists(*first_slice));
258 EXPECT_THAT(ReadFile(*first_slice), expected_key2_data);
259
260 absl::StatusOr<std::string> second_slice = (*iterator1)->Next();
261 ASSERT_OK(second_slice);
262 ASSERT_TRUE(FileExists(*second_slice));
263 EXPECT_THAT(ReadFile(*second_slice), expected_key1_data);
264
265 // We should now have reached the end of the first iterator.
266 EXPECT_THAT((*iterator1)->Next(), IsCode(OUT_OF_RANGE));
267
268 // Closing the iterator should not fail/crash.
269 (*iterator1)->Close();
270 // The slice files we saw earlier (possibly all the same file) should now be
271 // deleted.
272 ASSERT_FALSE(FileExists(*first_slice));
273 ASSERT_FALSE(FileExists(*second_slice));
274
275 const std::string expected_key99_data = "key99_data";
276 EXPECT_CALL(mock_http_client_, PerformSingleRequest(SimpleHttpRequestMatcher(
277 "https://foo.bar/id/Y/baz/99/bazz",
278 HttpRequest::Method::kGet, _, "")))
279 .WillOnce(
280 Return(FakeHttpResponse(200, HeaderList(), expected_key99_data)));
281
282 {
283 InSequence in_sequence;
284 EXPECT_CALL(mock_log_manager_,
285 LogDiag(ProdDiagCode::FEDSELECT_SLICE_HTTP_FETCH_REQUESTED));
286 EXPECT_CALL(mock_log_manager_,
287 LogDiag(ProdDiagCode::FEDSELECT_SLICE_HTTP_FETCH_SUCCEEDED));
288 }
289
290 absl::StatusOr<std::unique_ptr<ExampleIterator>> iterator2 =
291 iterator_factory->CreateExampleIterator(
292 CreateExampleSelector(/*served_at_id=*/"id/Y", /*keys=*/{99}));
293
294 // The iterator creation should have succeeded.
295 ASSERT_OK(iterator2);
296
297 // Reading the data for the slices should now succeed.
298 absl::StatusOr<std::string> third_slice = (*iterator2)->Next();
299 ASSERT_OK(third_slice);
300 ASSERT_TRUE(FileExists(*third_slice));
301 EXPECT_THAT(ReadFile(*third_slice), expected_key99_data);
302
303 // We purposely do not close the 2nd iterator, nor iterate it all the way to
304 // the end until we receive OUT_OF_RANGE, but instead simply destroy it. This
305 // should have the same effect as closing it, and cause the file to be
306 // deleted.
307 *iterator2 = nullptr;
308 ASSERT_FALSE(FileExists(*third_slice));
309 }
310
311 /** Tests the case where the fetched resources are compressed using the
312 * "Content-Type: ...+gzip" approach. The data should be decompressed before
313 * being returned.
314 */
TEST_F(HttpFederatedSelectManagerTest,SuccessfullyFetchCompressedSlice)315 TEST_F(HttpFederatedSelectManagerTest, SuccessfullyFetchCompressedSlice) {
316 const std::string uri_template =
317 "https://foo.bar/{served_at_id}/{key_base10}";
318
319 // Create an iterator factory with a valid (non-empty) URI template.
320 std::unique_ptr<ExampleIteratorFactory> iterator_factory =
321 fedselect_manager_.CreateExampleIteratorFactoryForUriTemplate(
322 uri_template);
323
324 // Once the first iterator is created we expect the following slice fetch
325 // requests to be issued immediately.
326 const std::string expected_key1_data = "key1_data";
327 EXPECT_CALL(mock_http_client_,
328 PerformSingleRequest(SimpleHttpRequestMatcher(
329 "https://foo.bar/id-X/1", HttpRequest::Method::kGet, _, "")))
330 .WillOnce(Return(FakeHttpResponse(
331 200, HeaderList{{"Content-Type", "application/octet-stream+gzip"}},
332 *CompressWithGzip(expected_key1_data))));
333
334 absl::StatusOr<std::unique_ptr<ExampleIterator>> iterator =
335 iterator_factory->CreateExampleIterator(CreateExampleSelector(
336 /*served_at_id=*/"id-X", /*keys=*/{1}));
337
338 // The iterator creation should have succeeded.
339 ASSERT_OK(iterator);
340
341 // Reading the data for the slice should now succeed and return the expected
342 // (uncompressed) data.
343 absl::StatusOr<std::string> slice = (*iterator)->Next();
344 ASSERT_OK(slice);
345 ASSERT_TRUE(FileExists(*slice));
346 EXPECT_THAT(ReadFile(*slice), expected_key1_data);
347 }
348
349 /** Tests the case where the URI template contains the substitution strings more
350 * than once. The client should replace *all* of them, not just the first one.
351 */
TEST_F(HttpFederatedSelectManagerTest,SuccessfullyFetchFromUriTemplateWithMultipleTemplateEntries)352 TEST_F(HttpFederatedSelectManagerTest,
353 SuccessfullyFetchFromUriTemplateWithMultipleTemplateEntries) {
354 const std::string uri_template =
355 "https://{served_at_id}.foo.bar/{key_base10}{served_at_id}/baz/"
356 "{key_base10}/bazz";
357
358 // Create an iterator factory with a valid (non-empty) URI template.
359 std::unique_ptr<ExampleIteratorFactory> iterator_factory =
360 fedselect_manager_.CreateExampleIteratorFactoryForUriTemplate(
361 uri_template);
362
363 // Once the first iterator is created we expect the following slice fetch
364 // requests to be issued immediately.
365 const std::string expected_key1_data = "key1_data";
366 EXPECT_CALL(mock_http_client_, PerformSingleRequest(SimpleHttpRequestMatcher(
367 "https://id-X.foo.bar/1id-X/baz/1/bazz",
368 HttpRequest::Method::kGet, _, "")))
369 .WillOnce(
370 Return(FakeHttpResponse(200, HeaderList(), expected_key1_data)));
371
372 absl::StatusOr<std::unique_ptr<ExampleIterator>> iterator =
373 iterator_factory->CreateExampleIterator(CreateExampleSelector(
374 /*served_at_id=*/"id-X", /*keys=*/{1}));
375
376 // The iterator creation should have succeeded.
377 ASSERT_OK(iterator);
378
379 // Reading the data should now succeed.
380 absl::StatusOr<std::string> slice = (*iterator)->Next();
381 ASSERT_OK(slice);
382 ASSERT_TRUE(FileExists(*slice));
383 EXPECT_THAT(ReadFile(*slice), expected_key1_data);
384
385 // We should now have reached the end of the first iterator.
386 EXPECT_THAT((*iterator)->Next(), IsCode(OUT_OF_RANGE));
387
388 // Closing the iterator should not fail/crash.
389 (*iterator)->Close();
390 // The slice files we saw earlier (possibly all the same file) should now be
391 // deleted.
392 ASSERT_FALSE(FileExists(*slice));
393 }
394
395 /** Tests the case where the URI template contains the substitution strings more
396 * than once. The client should replace *all* of them, not just the first one.
397 */
TEST_F(HttpFederatedSelectManagerTest,ErrorDuringFetch)398 TEST_F(HttpFederatedSelectManagerTest, ErrorDuringFetch) {
399 const std::string uri_template =
400 "https://foo.bar/{served_at_id}/{key_base10}";
401
402 // Create an iterator factory with a valid (non-empty) URI template.
403 std::unique_ptr<ExampleIteratorFactory> iterator_factory =
404 fedselect_manager_.CreateExampleIteratorFactoryForUriTemplate(
405 uri_template);
406
407 // Once the first iterator is created we expect the following slice fetch
408 // requests to be issued immediately. We'll make the 2nd slice's HTTP request
409 // return an error.
410 EXPECT_CALL(mock_http_client_, PerformSingleRequest(SimpleHttpRequestMatcher(
411 "https://foo.bar/id-X/998",
412 HttpRequest::Method::kGet, _, "")))
413 .WillOnce(Return(FakeHttpResponse(200, HeaderList(), "")));
414
415 EXPECT_CALL(mock_http_client_, PerformSingleRequest(SimpleHttpRequestMatcher(
416 "https://foo.bar/id-X/999",
417 HttpRequest::Method::kGet, _, "")))
418 .WillOnce(Return(FakeHttpResponse(404, HeaderList(), "")));
419
420 {
421 InSequence in_sequence;
422 EXPECT_CALL(mock_log_manager_,
423 LogDiag(ProdDiagCode::FEDSELECT_SLICE_HTTP_FETCH_REQUESTED));
424 EXPECT_CALL(mock_log_manager_,
425 LogDiag(ProdDiagCode::FEDSELECT_SLICE_HTTP_FETCH_FAILED));
426 }
427
428 absl::StatusOr<std::unique_ptr<ExampleIterator>> iterator =
429 iterator_factory->CreateExampleIterator(CreateExampleSelector(
430 /*served_at_id=*/"id-X", /*keys=*/{998, 999}));
431
432 // The iterator creation should fail, since if we can't fetch all of the
433 // slices successfully there is no way the plan can continue executing.
434 // The error code should be UNAVAILABLE and the error message should include
435 // the original NOT_FOUND error code that HTTP's 404 maps to, as well as the
436 // URI template to aid debugging.
437 EXPECT_THAT(iterator, IsCode(UNAVAILABLE));
438 EXPECT_THAT(iterator.status().message(), HasSubstr("fetch request failed"));
439 EXPECT_THAT(iterator.status().message(), HasSubstr(uri_template));
440 EXPECT_THAT(iterator.status().message(), HasSubstr("NOT_FOUND"));
441 // The error message should not contain the exact original HTTP code (since we
442 // expect the HTTP layer's error *messages* to not be included in the message
443 // returned to the plan).
444 EXPECT_THAT(iterator.status().message(), Not(HasSubstr("404")));
445 // The error message should not contain the slice IDs either.
446 EXPECT_THAT(iterator.status().message(), Not(HasSubstr("998")));
447 EXPECT_THAT(iterator.status().message(), Not(HasSubstr("999")));
448 }
449
450 } // anonymous namespace
451 } // namespace fcp::client
452