Only allow explicitly whitelisted users to issue lease requests to Machine Provider

appengine/machine_provider/handlers_endpoints_test.py

 #!/usr/bin/env python
 # Copyright 2015 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 """Unit tests for handlers_endpoints.py."""
 
 import datetime
 import json
 import unittest
 
 import test_env
 test_env.setup_test_env()
 
 from google.appengine import runtime
 
 from protorpc.remote import protojson
 import webtest
 
-from components import auth_testing
 from components import utils
 from components.machine_provider import rpc_messages
 from test_support import test_case
 
 import acl
 import handlers_endpoints
 import models
 
 
 def rpc_to_json(rpc_message):
@@ skipping 523 matching lines @@
 class MachineProviderReleaseTest(test_case.EndpointsTestCase):
   """Tests for handlers_endpoints.MachineProviderEndpoints.release."""
   api_service_cls = handlers_endpoints.MachineProviderEndpoints
 
   def setUp(self):
     super(MachineProviderReleaseTest, self).setUp()
     app = handlers_endpoints.create_endpoints_app()
     self.app = webtest.TestApp(app)
 
   def test_release(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     self.mock(
         handlers_endpoints.MachineProviderEndpoints,
         '_release',
         lambda *args, **kwargs: None,
     )
 
     request = rpc_to_json(rpc_messages.LeaseReleaseRequest(
         request_id='request-id',
     ))
-    auth_testing.mock_get_current_identity(self)
 
     response = jsonish_dict_to_rpc(
         self.call_api('release', request).json,
         rpc_messages.LeaseReleaseResponse,
     )
     self.assertEqual(response.client_request_id, 'request-id')
     self.assertFalse(response.error)
 
 
 class MachineProviderBatchedReleaseTest(test_case.EndpointsTestCase):
   """Tests for handlers_endpoints.MachineProviderEndpoints.batched_release."""
   api_service_cls = handlers_endpoints.MachineProviderEndpoints
 
   def setUp(self):
     super(MachineProviderBatchedReleaseTest, self).setUp()
     app = handlers_endpoints.create_endpoints_app()
     self.app = webtest.TestApp(app)
 
   def test_batch(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     ts = utils.utcnow()
     self.mock(utils, 'utcnow', lambda *args, **kwargs: ts)
 
     release_requests = rpc_to_json(rpc_messages.BatchedLeaseReleaseRequest(
         requests=[
             rpc_messages.LeaseReleaseRequest(
                 request_id='request-id',
             ),
         ],
     ))
-    auth_testing.mock_get_current_identity(self)
 
     release_responses = jsonish_dict_to_rpc(
         self.call_api('batched_release', release_requests).json,
         rpc_messages.BatchedLeaseReleaseResponse,
     )
     self.assertEqual(len(release_responses.responses), 1)
     self.assertEqual(
         release_responses.responses[0].client_request_id, 'request-id')
     self.assertEqual(
         release_responses.responses[0].error,
         rpc_messages.LeaseReleaseRequestError.NOT_FOUND,
     )
 
   def test_deadline_exceeded(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     class utcnow(object):
       def __init__(self, init_ts):
         self.last_ts = init_ts
       def __call__(self, *args, **kwargs):
         self.last_ts = self.last_ts + datetime.timedelta(seconds=60)
         return self.last_ts
     self.mock(utils, 'utcnow', utcnow(utils.utcnow()))
 
     release_requests = rpc_to_json(rpc_messages.BatchedLeaseReleaseRequest(
         requests=[
             rpc_messages.LeaseReleaseRequest(
                 request_id='request-id',
             ),
         ],
     ))
-    auth_testing.mock_get_current_identity(self)
 
     release_responses = jsonish_dict_to_rpc(
         self.call_api('batched_release', release_requests).json,
         rpc_messages.BatchedLeaseReleaseResponse,
     )
     self.assertEqual(len(release_responses.responses), 1)
     self.assertEqual(
         release_responses.responses[0].client_request_id, 'request-id')
     self.assertEqual(
         release_responses.responses[0].error,
         rpc_messages.LeaseReleaseRequestError.DEADLINE_EXCEEDED,
     )
 
   def test_exception(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     ts = utils.utcnow()
     self.mock(utils, 'utcnow', lambda *args, **kwargs: ts)
 
     def _release(*args, **kwargs):
       raise runtime.apiproxy_errors.CancelledError
     self.mock(handlers_endpoints.MachineProviderEndpoints, '_release', _release)
 
     release_requests = rpc_to_json(rpc_messages.BatchedLeaseReleaseRequest(
         requests=[
             rpc_messages.LeaseReleaseRequest(
                 request_id='request-id',
             ),
         ],
     ))
-    auth_testing.mock_get_current_identity(self)
 
     release_responses = jsonish_dict_to_rpc(
         self.call_api('batched_release', release_requests).json,
         rpc_messages.BatchedLeaseReleaseResponse,
     )
     self.assertEqual(len(release_responses.responses), 1)
     self.assertEqual(
         release_responses.responses[0].client_request_id, 'request-id')
     self.assertEqual(
         release_responses.responses[0].error,
         rpc_messages.LeaseReleaseRequestError.TRANSIENT_ERROR,
     )
 
 
 class MachineProviderBatchedLeaseTest(test_case.EndpointsTestCase):
   """Tests for handlers_endpoints.MachineProviderEndpoints.batched_lease."""
   api_service_cls = handlers_endpoints.MachineProviderEndpoints
 
   def setUp(self):
     super(MachineProviderBatchedLeaseTest, self).setUp()
     app = handlers_endpoints.create_endpoints_app()
     self.app = webtest.TestApp(app)
 
   def test_batch(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     ts = utils.utcnow()
     self.mock(utils, 'utcnow', lambda *args, **kwargs: ts)
 
     lease_requests = rpc_to_json(rpc_messages.BatchedLeaseRequest(requests=[
         rpc_messages.LeaseRequest(
             dimensions=rpc_messages.Dimensions(
                 os_family=rpc_messages.OSFamily.LINUX,
             ),
             duration=1,
             request_id='request-id',
         ),
     ]))
-    auth_testing.mock_get_current_identity(self)
 
     lease_responses = jsonish_dict_to_rpc(
         self.call_api('batched_lease', lease_requests).json,
         rpc_messages.BatchedLeaseResponse,
     )
     self.assertEqual(len(lease_responses.responses), 1)
     self.assertEqual(
         lease_responses.responses[0].client_request_id, 'request-id')
     self.assertFalse(lease_responses.responses[0].error)
 
   def test_deadline_exceeded(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     class utcnow(object):
       def __init__(self, init_ts):
         self.last_ts = init_ts
       def __call__(self, *args, **kwargs):
         self.last_ts = self.last_ts + datetime.timedelta(seconds=60)
         return self.last_ts
     self.mock(utils, 'utcnow', utcnow(utils.utcnow()))
 
     lease_requests = rpc_to_json(rpc_messages.BatchedLeaseRequest(requests=[
         rpc_messages.LeaseRequest(
             dimensions=rpc_messages.Dimensions(
                 os_family=rpc_messages.OSFamily.LINUX,
             ),
             duration=1,
             request_id='request-id',
         ),
     ]))
-    auth_testing.mock_get_current_identity(self)
 
     lease_responses = jsonish_dict_to_rpc(
         self.call_api('batched_lease', lease_requests).json,
         rpc_messages.BatchedLeaseResponse,
     )
     self.assertEqual(len(lease_responses.responses), 1)
     self.assertEqual(
         lease_responses.responses[0].client_request_id, 'request-id')
     self.assertEqual(
         lease_responses.responses[0].error,
         rpc_messages.LeaseRequestError.DEADLINE_EXCEEDED,
     )
 
   def test_exception(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     ts = utils.utcnow()
     self.mock(utils, 'utcnow', lambda *args, **kwargs: ts)
 
     def _lease(*args, **kwargs):
       raise runtime.apiproxy_errors.CancelledError
     self.mock(handlers_endpoints.MachineProviderEndpoints, '_lease', _lease)
 
     lease_requests = rpc_to_json(rpc_messages.BatchedLeaseRequest(requests=[
         rpc_messages.LeaseRequest(
             dimensions=rpc_messages.Dimensions(
                 os_family=rpc_messages.OSFamily.LINUX,
             ),
             duration=1,
             request_id='request-id',
         ),
     ]))
-    auth_testing.mock_get_current_identity(self)
 
     lease_responses = jsonish_dict_to_rpc(
         self.call_api('batched_lease', lease_requests).json,
         rpc_messages.BatchedLeaseResponse,
     )
     self.assertEqual(len(lease_responses.responses), 1)
     self.assertEqual(
         lease_responses.responses[0].client_request_id, 'request-id')
     self.assertEqual(
         lease_responses.responses[0].error,
         rpc_messages.LeaseRequestError.TRANSIENT_ERROR,
     )
 
 
 class MachineProviderLeaseTest(test_case.EndpointsTestCase):
   """Tests for handlers_endpoints.MachineProviderEndpoints.lease."""
   api_service_cls = handlers_endpoints.MachineProviderEndpoints
 
   def setUp(self):
     super(MachineProviderLeaseTest, self).setUp()
     app = handlers_endpoints.create_endpoints_app()
     self.app = webtest.TestApp(app)
 
   def test_lease_duration(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     lease_request = rpc_to_json(rpc_messages.LeaseRequest(
         dimensions=rpc_messages.Dimensions(
             os_family=rpc_messages.OSFamily.LINUX,
         ),
         duration=1,
         request_id='abc',
         pubsub_topic='topic',
     ))
-    auth_testing.mock_get_current_identity(self)
 
     lease_response = jsonish_dict_to_rpc(
         self.call_api('lease', lease_request).json,
         rpc_messages.LeaseResponse,
     )
     self.assertFalse(lease_response.error)
 
   def test_lease_duration_zero(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     lease_request = rpc_to_json(rpc_messages.LeaseRequest(
         dimensions=rpc_messages.Dimensions(
             os_family=rpc_messages.OSFamily.LINUX,
         ),
         duration=0,
         request_id='abc',
     ))
-    auth_testing.mock_get_current_identity(self)
 
     lease_response = jsonish_dict_to_rpc(
         self.call_api('lease', lease_request).json,
         rpc_messages.LeaseResponse,
     )
     self.assertEqual(
         lease_response.error,
         rpc_messages.LeaseRequestError.LEASE_LENGTH_UNSPECIFIED,
     )
 
   def test_lease_duration_negative(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     lease_request = rpc_to_json(rpc_messages.LeaseRequest(
         dimensions=rpc_messages.Dimensions(
             os_family=rpc_messages.OSFamily.LINUX,
         ),
         duration=-1,
         request_id='abc',
     ))
-    auth_testing.mock_get_current_identity(self)
 
     lease_response = jsonish_dict_to_rpc(
         self.call_api('lease', lease_request).json,
-        rpc_messages.LeaseResponse,
-    )
-    self.assertEqual(
-        lease_response.error,
-        rpc_messages.LeaseRequestError.NONPOSITIVE_DEADLINE,
-    )
-
-  def test_lease_duration_negative(self):
-    lease_request = rpc_to_json(rpc_messages.LeaseRequest(
-        dimensions=rpc_messages.Dimensions(
-            os_family=rpc_messages.OSFamily.LINUX,
-        ),
-        duration=-1,
-        request_id='abc',
-    ))
-    auth_testing.mock_get_current_identity(self)
-
-    lease_response = jsonish_dict_to_rpc(
-        self.call_api('lease', lease_request).json,
         rpc_messages.LeaseResponse,
     )
     self.assertEqual(
         lease_response.error,
         rpc_messages.LeaseRequestError.NONPOSITIVE_DEADLINE,
     )
 
   def test_lease_duration_and_lease_expiration_ts(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     lease_request = rpc_to_json(rpc_messages.LeaseRequest(
         dimensions=rpc_messages.Dimensions(
             os_family=rpc_messages.OSFamily.LINUX,
         ),
         duration=1,
         lease_expiration_ts=9999999999,
         request_id='abc',
     ))
-    auth_testing.mock_get_current_identity(self)
 
     lease_response = jsonish_dict_to_rpc(
         self.call_api('lease', lease_request).json,
         rpc_messages.LeaseResponse,
     )
     self.assertEqual(
         lease_response.error,
         rpc_messages.LeaseRequestError.MUTUAL_EXCLUSION_ERROR,
     )
 
   def test_lease_timestamp(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     lease_request = rpc_to_json(rpc_messages.LeaseRequest(
         dimensions=rpc_messages.Dimensions(
             os_family=rpc_messages.OSFamily.LINUX,
         ),
         lease_expiration_ts=9999999999,
         request_id='abc',
     ))
-    auth_testing.mock_get_current_identity(self)
 
     lease_response = jsonish_dict_to_rpc(
         self.call_api('lease', lease_request).json,
         rpc_messages.LeaseResponse,
     )
     self.assertFalse(lease_response.error)
 
   def test_lease_timestamp_passed(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     lease_request = rpc_to_json(rpc_messages.LeaseRequest(
         dimensions=rpc_messages.Dimensions(
             os_family=rpc_messages.OSFamily.LINUX,
         ),
         lease_expiration_ts=1,
         request_id='abc',
     ))
-    auth_testing.mock_get_current_identity(self)
 
     lease_response = jsonish_dict_to_rpc(
         self.call_api('lease', lease_request).json,
         rpc_messages.LeaseResponse,
     )
     self.assertEqual(
         lease_response.error,
         rpc_messages.LeaseRequestError.LEASE_EXPIRATION_TS_ERROR,
     )
 
   def test_duplicate(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     lease_request = rpc_to_json(rpc_messages.LeaseRequest(
         dimensions=rpc_messages.Dimensions(
             os_family=rpc_messages.OSFamily.OSX,
         ),
         duration=3,
         request_id='asdf',
     ))
-    auth_testing.mock_get_current_identity(self)
 
     lease_response_1 = jsonish_dict_to_rpc(
         self.call_api('lease', lease_request).json,
         rpc_messages.LeaseResponse,
     )
     lease_response_2 = jsonish_dict_to_rpc(
         self.call_api('lease', lease_request).json,
         rpc_messages.LeaseResponse,
     )
     self.assertFalse(lease_response_1.error)
     self.assertFalse(lease_response_2.error)
     self.assertEqual(
         lease_response_1.request_hash,
         lease_response_2.request_hash,
     )
 
   def test_request_id_reuse(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     lease_request_1 = rpc_to_json(rpc_messages.LeaseRequest(
         dimensions=rpc_messages.Dimensions(
             os_family=rpc_messages.OSFamily.WINDOWS,
         ),
         duration=7,
         request_id='qwerty',
     ))
     lease_request_2 = rpc_to_json(rpc_messages.LeaseRequest(
         dimensions=rpc_messages.Dimensions(
             os_family=rpc_messages.OSFamily.WINDOWS,
         ),
         duration=189,
         request_id='qwerty',
     ))
-    auth_testing.mock_get_current_identity(self)
 
     lease_response_1 = jsonish_dict_to_rpc(
         self.call_api('lease', lease_request_1).json,
         rpc_messages.LeaseResponse,
     )
     lease_response_2 = jsonish_dict_to_rpc(
         self.call_api('lease', lease_request_2).json,
         rpc_messages.LeaseResponse,
     )
     self.assertFalse(lease_response_1.error)
     self.assertEqual(
         lease_response_2.error,
         rpc_messages.LeaseRequestError.REQUEST_ID_REUSE,
     )
     self.assertNotEqual(
         lease_response_1.request_hash,
         lease_response_2.request_hash,
     )
 
   def test_invalid_topic(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     lease_request = rpc_to_json(rpc_messages.LeaseRequest(
         dimensions=rpc_messages.Dimensions(
             os_family=rpc_messages.OSFamily.WINDOWS,
         ),
         duration=9,
         pubsub_topic='../../a-different-project/topics/my-topic',
         request_id='123',
     ))
-    auth_testing.mock_get_current_identity(self)
 
     lease_response = jsonish_dict_to_rpc(
         self.call_api('lease', lease_request).json,
         rpc_messages.LeaseResponse,
     )
     self.assertEqual(
         lease_response.error,
         rpc_messages.LeaseRequestError.INVALID_TOPIC,
     )
 
   def test_invalid_project(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     lease_request = rpc_to_json(rpc_messages.LeaseRequest(
         dimensions=rpc_messages.Dimensions(
             os_family=rpc_messages.OSFamily.WINDOWS,
         ),
         duration=9,
         pubsub_topic='my-topic',
         pubsub_project='../../a-different-project/topics/my-other-topic',
         request_id='123',
     ))
-    auth_testing.mock_get_current_identity(self)
 
     lease_response = jsonish_dict_to_rpc(
         self.call_api('lease', lease_request).json,
         rpc_messages.LeaseResponse,
     )
     self.assertEqual(
         lease_response.error,
         rpc_messages.LeaseRequestError.INVALID_PROJECT,
     )
 
   def test_project_without_topic(self):
+    def is_group_member(group):
+      return group == 'machine-provider-users'
+    self.mock(acl.auth, 'is_group_member', is_group_member)
     lease_request = rpc_to_json(rpc_messages.LeaseRequest(
         dimensions=rpc_messages.Dimensions(
             os_family=rpc_messages.OSFamily.WINDOWS,
         ),
         duration=9,
         pubsub_project='my-project',
         request_id='123',
     ))
-    auth_testing.mock_get_current_identity(self)
 
     lease_response = jsonish_dict_to_rpc(
         self.call_api('lease', lease_request).json,
         rpc_messages.LeaseResponse,
     )
     self.assertEqual(
         lease_response.error,
         rpc_messages.LeaseRequestError.UNSPECIFIED_TOPIC,
     )
 
 
 if __name__ == '__main__':
   unittest.main()