client/tests/auth_server_test.py

#!/usr/bin/env vpython3
# Copyright 2016 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.

import contextlib
import json
import socket
import time

# Mutates sys.path.
import test_env

# third_party/
from depot_tools import auto_stub
import requests

from libs import luci_context
from utils import authenticators
from utils import auth_server
from utils import net
from utils import oauth

import net_utils


def call_rpc(account_id, scopes=None, audience=None):
  ctx = luci_context.read('local_auth')

  if audience is None:
    assert scopes
    method = 'GetOAuthToken'
    body = {
      'account_id': account_id,
      'scopes': scopes,
      'secret': ctx['secret'],
    }
  else:
    assert scopes is None
    method = 'GetIDToken'
    body = {
      'account_id': account_id,
      'audience': audience,
      'secret': ctx['secret'],
    }

  r = requests.post(
      url='http://127.0.0.1:%d/rpc/LuciLocalAuthService.%s' % (
          ctx['rpc_port'], method),
      data=json.dumps(body),
      headers={'Content-Type': 'application/json'})

  return r.json()


@contextlib.contextmanager
def local_auth_server(token_cb, default_account_id, **overrides):
  class MockedProvider:
    def generate_access_token(self, account_id, scopes):
      return token_cb(account_id, scopes=scopes)
    def generate_id_token(self, account_id, audience):
      return token_cb(account_id, audience=audience)

  acc = lambda aid: auth_server.Account(id=aid, email=aid+'@example.com')

  s = auth_server.LocalAuthServer()
  try:
    local_auth = s.start(
        token_provider=MockedProvider(),
        accounts=(acc('acc_1'), acc('acc_2'), acc('acc_3')),
        default_account_id=default_account_id)
    local_auth.update(overrides)
    with luci_context.write(local_auth=local_auth):
      yield
  finally:
    s.stop()


class LocalAuthServerTest(auto_stub.TestCase):
  epoch = 12345678

  def setUp(self):
    super(LocalAuthServerTest, self).setUp()
    self.mock_time(0)

  def mock_time(self, delta):
    self.mock(time, 'time', lambda: self.epoch + delta)

  def test_accounts_in_ctx(self):
    def token_gen(_account_id, **_kwargs):
      self.fail('must not be called')
    with local_auth_server(token_gen, 'acc_1'):
      ctx = luci_context.read('local_auth')
      ctx.pop('rpc_port')
      ctx.pop('secret')
      self.assertEqual({
          'accounts': [
              {'email': 'acc_1@example.com', 'id': 'acc_1'},
              {'email': 'acc_2@example.com', 'id': 'acc_2'},
              {'email': 'acc_3@example.com', 'id': 'acc_3'},
          ],
         'default_account_id': 'acc_1',
      }, ctx)

  def test_access_tokens(self):
    calls = []
    def token_gen(account_id, scopes=None, audience=None):
      assert audience is None
      calls.append((account_id, scopes))
      return auth_server.AccessToken('tok_%s' % account_id, time.time() + 300)

    with local_auth_server(token_gen, 'acc_1'):
      # Grab initial token.
      resp = call_rpc('acc_1', scopes=['B', 'B', 'A', 'C'])
      self.assertEqual({
          'access_token': 'tok_acc_1',
          'expiry': self.epoch + 300
      }, resp)
      self.assertEqual([('acc_1', ('A', 'B', 'C'))], calls)
      del calls[:]

      # Reuses cached token until it is close to expiration.
      self.mock_time(60)
      resp = call_rpc('acc_1', scopes=['B', 'A', 'C'])
      self.assertEqual({
          'access_token': 'tok_acc_1',
          'expiry': self.epoch + 300
      }, resp)
      self.assertFalse(calls)

      # Asking for different account gives another token.
      resp = call_rpc('acc_2', scopes=['B', 'B', 'A', 'C'])
      self.assertEqual({
          'access_token': 'tok_acc_2',
          'expiry': self.epoch + 360
      }, resp)
      self.assertEqual([('acc_2', ('A', 'B', 'C'))], calls)
      del calls[:]

      # First token has expired. Generated new one.
      self.mock_time(300)
      resp = call_rpc('acc_1', scopes=['A', 'B', 'C'])
      self.assertEqual({
          'access_token': 'tok_acc_1',
          'expiry': self.epoch + 600
      }, resp)
      self.assertEqual([('acc_1', ('A', 'B', 'C'))], calls)

  def test_id_tokens(self):
    calls = []
    def token_gen(account_id, scopes=None, audience=None):
      assert scopes is None
      calls.append((account_id, audience))
      return auth_server.AccessToken('tok_%s' % account_id, time.time() + 300)

    with local_auth_server(token_gen, 'acc_1'):
      # Grab initial token.
      resp = call_rpc('acc_1', audience='some-audience')
      self.assertEqual({
          'id_token': 'tok_acc_1',
          'expiry': self.epoch + 300
      }, resp)
      self.assertEqual([('acc_1', 'some-audience')], calls)
      del calls[:]

      # Reuses cached token until it is close to expiration.
      self.mock_time(60)
      resp = call_rpc('acc_1', audience='some-audience')
      self.assertEqual({
          'id_token': 'tok_acc_1',
          'expiry': self.epoch + 300
      }, resp)
      self.assertFalse(calls)

      # Asking for different audience gives another token.
      resp = call_rpc('acc_1', audience='another-audience')
      self.assertEqual({
          'id_token': 'tok_acc_1',
          'expiry': self.epoch + 360
      }, resp)
      self.assertEqual([('acc_1', 'another-audience')], calls)
      del calls[:]

      # Asking for different account gives another token.
      resp = call_rpc('acc_2', audience='some-audience')
      self.assertEqual({
          'id_token': 'tok_acc_2',
          'expiry': self.epoch + 360
      }, resp)
      self.assertEqual([('acc_2', 'some-audience')], calls)
      del calls[:]

      # First token has expired. Generated new one.
      self.mock_time(300)
      resp = call_rpc('acc_1', audience='some-audience')
      self.assertEqual({
          'id_token': 'tok_acc_1',
          'expiry': self.epoch + 600
      }, resp)
      self.assertEqual([('acc_1', 'some-audience')], calls)

  def test_handles_token_errors(self):
    calls = []
    def token_gen(_account_id, **_kwargs):
      calls.append(1)
      raise auth_server.TokenError(123, 'error message')

    with local_auth_server(token_gen, 'acc_1'):
      self.assertEqual({
          'error_code': 123,
          'error_message': 'error message'
      }, call_rpc('acc_1', scopes=['B', 'B', 'A', 'C']))
      self.assertEqual(1, len(calls))

      # Errors are cached. Same error is returned.
      self.assertEqual({
          'error_code': 123,
          'error_message': 'error message'
      }, call_rpc('acc_1', scopes=['B', 'B', 'A', 'C']))
      self.assertEqual(1, len(calls))

  def test_http_level_errors(self):
    def token_gen(_account_id, **_kwargs):
      self.fail('must not be called')

    with local_auth_server(token_gen, 'acc_1'):
      # Wrong URL.
      ctx = luci_context.read('local_auth')
      r = requests.post(
          url='http://127.0.0.1:%d/blah/LuciLocalAuthService.GetOAuthToken' %
              ctx['rpc_port'],
          data=json.dumps({
            'account_id': 'acc_1',
            'scopes': ['A', 'B', 'C'],
            'secret': ctx['secret'],
          }),
          headers={'Content-Type': 'application/json'})
      self.assertEqual(404, r.status_code)

      # Wrong HTTP method.
      r = requests.get(
          url='http://127.0.0.1:%d/rpc/LuciLocalAuthService.GetOAuthToken' %
              ctx['rpc_port'],
          data=json.dumps({
            'account_id': 'acc_1',
            'scopes': ['A', 'B', 'C'],
            'secret': ctx['secret'],
          }),
          headers={'Content-Type': 'application/json'})
      self.assertEqual(501, r.status_code)

      # Wrong content type.
      r = requests.post(
          url='http://127.0.0.1:%d/rpc/LuciLocalAuthService.GetOAuthToken' %
              ctx['rpc_port'],
          data=json.dumps({
            'account_id': 'acc_1',
            'scopes': ['A', 'B', 'C'],
            'secret': ctx['secret'],
          }),
          headers={'Content-Type': 'application/xml'})
      self.assertEqual(400, r.status_code)

      # Bad JSON.
      r = requests.post(
          url='http://127.0.0.1:%d/rpc/LuciLocalAuthService.GetOAuthToken' %
              ctx['rpc_port'],
          data='not a json',
          headers={'Content-Type': 'application/json'})
      self.assertEqual(400, r.status_code)

  def test_validation(self):
    def token_gen(_account_id, **_kwargs):
      self.fail('must not be called')

    with local_auth_server(token_gen, 'acc_1'):
      ctx = luci_context.read('local_auth')

      def must_fail(method, body, err, code):
        for m in ['GetOAuthToken', 'GetIDToken'] if method == '*' else [method]:
          r = requests.post(
              url='http://127.0.0.1:%d/rpc/LuciLocalAuthService.%s' % (
                  ctx['rpc_port'], m),
              data=json.dumps(body),
              headers={'Content-Type': 'application/json'})
          self.assertEqual(code, r.status_code)
          self.assertIn(err, r.text)

      cases = [
        # account_id
        (
          '*',
          {},
          '"account_id" is required',
          400,
        ),
        (
          '*',
          {'account_id': 123},
          '"account_id" must be a string',
          400,
        ),

        # secret
        (
          '*',
          {'account_id': 'acc_1', 'scopes': ['a']},
          '"secret" is required',
          400,
        ),
        (
          '*',
          {'account_id': 'acc_1', 'scopes': ['a'], 'secret': 123},
          '"secret" must be a string',
          400,
        ),
        (
          '*',
          {'account_id': 'acc_1', 'scopes': ['a'], 'secret': 'abc'},
          'Invalid "secret"',
          403,
        ),

        # The account is known.
        (
          '*',
          {'account_id': 'zzz', 'scopes': ['a'], 'secret': ctx['secret']},
          'Unrecognized account ID',
          404,
        ),

        # scopes
        (
          'GetOAuthToken',
          {'account_id': 'acc_1', 'secret': ctx['secret']},
          '"scopes" is required',
          400,
        ),
        (
          'GetOAuthToken',
          {'account_id': 'acc_1', 'secret': ctx['secret'], 'scopes': []},
          '"scopes" is required',
          400,
        ),
        (
          'GetOAuthToken',
          {'account_id': 'acc_1', 'secret': ctx['secret'], 'scopes': 'abc'},
          '"scopes" must be a list of strings',
          400,
        ),
        (
          'GetOAuthToken',
          {'account_id': 'acc_1', 'secret': ctx['secret'], 'scopes': [1]},
          '"scopes" must be a list of strings',
          400,
        ),

        # audience
        (
          'GetIDToken',
          {'account_id': 'acc_1', 'secret': ctx['secret']},
          '"audience" is required',
          400,
        ),
        (
          'GetIDToken',
          {'account_id': 'acc_1', 'secret': ctx['secret'], 'audience': ''},
          '"audience" is required',
          400,
        ),
        (
          'GetIDToken',
          {'account_id': 'acc_1', 'secret': ctx['secret'], 'audience': 123},
          '"audience" must be a string',
          400,
        ),
      ]
      for method, body, err, code in cases:
        must_fail(method, body, err, code)


class LocalAuthHttpServiceTest(auto_stub.TestCase):
  """Tests for LocalAuthServer and LuciContextAuthenticator."""
  epoch = 12345678

  def setUp(self):
    super(LocalAuthHttpServiceTest, self).setUp()
    self.mock_time(0)

  def mock_time(self, delta):
    self.mock(time, 'time', lambda: self.epoch + delta)

  @staticmethod
  def mocked_http_service(
      url='http://example.com',
      perform_request=None):

    class MockedRequestEngine:
      def perform_request(self, request):
        return perform_request(request) if perform_request else None

    return net.HttpService(
        url,
        authenticator=authenticators.LuciContextAuthenticator(),
        engine=MockedRequestEngine())

  def test_works(self):
    service_url = 'http://example.com'
    request_url = '/some_request'
    response = b'True'
    token = 'notasecret'

    def token_gen(account_id, scopes=None, audience=None):
      self.assertEqual('acc_1', account_id)
      self.assertEqual(1, len(scopes))
      self.assertEqual(oauth.OAUTH_SCOPES, scopes[0])
      self.assertIsNone(audience)
      return auth_server.AccessToken(token, time.time() + 300)

    def handle_request(request):
      self.assertTrue(
          request.get_full_url().startswith(service_url + request_url))
      self.assertEqual(b'', request.body)
      self.assertEqual('Bearer %s' % token, request.headers['Authorization'])
      return net_utils.make_fake_response(response, request.get_full_url())

    with local_auth_server(token_gen, 'acc_1'):
      service = self.mocked_http_service(perform_request=handle_request)
      self.assertEqual(service.request(request_url, data={}).read(), response)

  def test_bad_secret(self):
    service_url = 'http://example.com'
    request_url = '/some_request'
    response = b'False'

    def token_gen(_account_id, **_kwargs):
      self.fail('must not be called')

    def handle_request(request):
      self.assertTrue(
          request.get_full_url().startswith(service_url + request_url))
      self.assertEqual(b'', request.body)
      self.assertIsNone(request.headers.get('Authorization'))
      return net_utils.make_fake_response(response, request.get_full_url())

    with local_auth_server(token_gen, 'acc_1', secret='invalid'):
      service = self.mocked_http_service(perform_request=handle_request)
      self.assertEqual(service.request(request_url, data={}).read(), response)

  def test_bad_port(self):
    request_url = '/some_request'

    def token_gen(_account_id, **_kwargs):
      self.fail('must not be called')

    def handle_request(_request):
      self.fail('must not be called')

    # This little dance should pick an unused port, bind it and then close it,
    # trusting that the OS will not reallocate it between now and when the http
    # client attempts to use it as a local_auth service. This is better than
    # picking a static port number, as there's at least some guarantee that the
    # port WASN'T in use before this test ran.
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind(('localhost', 0))
    port = sock.getsockname()[1]
    sock.close()
    with local_auth_server(token_gen, 'acc_1', rpc_port=port):
      service = self.mocked_http_service(perform_request=handle_request)
      with self.assertRaises(ConnectionRefusedError):
        self.assertRaises(service.request(request_url, data={}).read())

  def test_expired_token(self):
    service_url = 'http://example.com'
    request_url = '/some_request'
    response = b'False'
    token = 'notasecret'

    def token_gen(account_id, scopes=None, audience=None):
      self.assertEqual('acc_1', account_id)
      self.assertEqual(1, len(scopes))
      self.assertEqual(oauth.OAUTH_SCOPES, scopes[0])
      self.assertIsNone(audience)
      return auth_server.AccessToken(token, time.time())

    def handle_request(request):
      self.assertTrue(
          request.get_full_url().startswith(service_url + request_url))
      self.assertEqual(b'', request.body)
      self.assertIsNone(request.headers.get('Authorization'))
      return net_utils.make_fake_response(response, request.get_full_url())

    with local_auth_server(token_gen, 'acc_1'):
      service = self.mocked_http_service(perform_request=handle_request)
      self.assertEqual(service.request(request_url, data={}).read(), response)


if __name__ == '__main__':
  # Terminate HTTP server in tests 50x faster. Impacts performance though so
  # do it only in tests.
  auth_server._HTTPServer.poll_interval = 0.01
  test_env.main()