Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[auth] Allow use of cloud access tokens for hail batch authorization #13131

Merged
merged 14 commits into from
Sep 5, 2023
Merged
66 changes: 59 additions & 7 deletions auth/auth/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
from gear.cloud_config import get_global_config
from gear.profiling import install_profiler_if_requested
from hailtop import httpx
from hailtop.auth import AzureFlow, Flow, GoogleFlow
from hailtop.config import get_deploy_config
from hailtop.hail_logging import AccessLogger
from hailtop.tls import internal_server_ssl_context
Expand All @@ -50,7 +51,6 @@
PreviouslyDeletedUser,
UnknownUser,
)
from .flow import get_flow_client
danking marked this conversation as resolved.
Show resolved Hide resolved

log = logging.getLogger('auth')

Expand Down Expand Up @@ -510,6 +510,12 @@ async def rest_login(request: web.Request) -> web.Response:
)


@routes.get('/api/v1alpha/oauth2-client')
async def hailctl_oauth_client(request): # pylint: disable=unused-argument
idp = 'Google' if CLOUD == 'gcp' else 'Microsoft'
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Aside: I don't love that gcp is a magic string for Google, but it is what is for now.

Shouldn't you use IdentityProvider.GOOGLE.value and IdentityProvider.MICROSOFT.value for the idp?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes you're right, I'll use that.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bump, the literal strings are still here.

return json_response({'idp': idp, 'oauth2_client': request.app['hailctl_client_config']})


@routes.get('/roles')
@authenticated_devs_only
async def get_roles(request: web.Request, userdata: UserData) -> web.Response:
Expand Down Expand Up @@ -737,11 +743,49 @@ async def rest_logout(request: web.Request, userdata: UserData) -> web.Response:
return web.Response(status=200)


async def get_userinfo(request: web.Request, session_id: str) -> UserData:
async def get_userinfo(request: web.Request, auth_token: str) -> UserData:
flow_client: Flow = request.app['flow_client']
client_session = request.app['client_session']

userdata = await get_userinfo_from_hail_session_id(request, auth_token)
if userdata:
return userdata

hailctl_oauth_client = request.app['hailctl_client_config']
uid = await flow_client.get_identity_uid_from_access_token(
client_session, auth_token, oauth2_client=hailctl_oauth_client
)
if uid:
return await get_userinfo_from_login_id_or_hail_identity_id(request, uid)

raise web.HTTPUnauthorized()


async def get_userinfo_from_login_id_or_hail_identity_id(request, login_id):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This appears to only be used with a IdP's uid, right? I think I'm a bit confused on hail identity uid vs login id. The login id is the username from the email, right? And hail identity uid is the uid from the cloud provider?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

login_id is what we use to identify human users. In GCP it is currently an email, in Azure it is the actual sub of that user in AAD.

hail_identity is the email in GCP or common name in Azure of the robot identity assigned to the user. hail_identity_uid is the actual sub of the robot identity in both of those places.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This appears to only be used with a IdP's uid, right

This is not the case because we do not store UIDs for human users in GCP, only their email.

daniel-goldstein marked this conversation as resolved.
Show resolved Hide resolved
db = request.app['db']

users = [
x
async for x in db.select_and_fetchall(
'''
SELECT users.*
FROM users
WHERE (users.login_id = %s OR users.hail_identity_uid = %s) AND users.state = 'active'
''',
(login_id, login_id),
)
]

if len(users) != 1:
log.info('Unknown login id')
raise web.HTTPUnauthorized()
return users[0]


async def get_userinfo_from_hail_session_id(request, session_id):
daniel-goldstein marked this conversation as resolved.
Show resolved Hide resolved
# b64 encoding of 32-byte session ID is 44 bytes
if len(session_id) != 44:
log.info('Session id != 44 bytes')
raise web.HTTPUnauthorized()
return None

db = request.app['db']
users = [
Expand All @@ -758,8 +802,7 @@ async def get_userinfo(request: web.Request, session_id: str) -> UserData:
]

if len(users) != 1:
log.info(f'Unknown session id: {session_id}')
raise web.HTTPUnauthorized()
return None
return users[0]


Expand Down Expand Up @@ -801,7 +844,16 @@ async def on_startup(app):
await db.async_init(maxsize=50)
app['db'] = db
app['client_session'] = httpx.client_session()
app['flow_client'] = get_flow_client('/auth-oauth2-client-secret/client_secret.json')

credentials_file = '/auth-oauth2-client-secret/client_secret.json'
if CLOUD == 'azure':
app['flow_client'] = AzureFlow(credentials_file)
else:
assert CLOUD == 'gcp'
app['flow_client'] = GoogleFlow(credentials_file)
daniel-goldstein marked this conversation as resolved.
Show resolved Hide resolved

with open('/auth-oauth2-client-secret/hailctl_client_secret.json', 'r', encoding='utf-8') as f:
app['hailctl_client_config'] = json.loads(f.read())
danking marked this conversation as resolved.
Show resolved Hide resolved

kubernetes_asyncio.config.load_incluster_config()
app['k8s_client'] = kubernetes_asyncio.client.CoreV1Api()
Expand Down
116 changes: 0 additions & 116 deletions auth/auth/flow.py

This file was deleted.

75 changes: 0 additions & 75 deletions auth/pinned-requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -4,78 +4,3 @@
#
# pip-compile --output-file=hail/auth/pinned-requirements.txt hail/auth/requirements.txt
#
cachetools==5.3.1
daniel-goldstein marked this conversation as resolved.
Show resolved Hide resolved
# via
# -c hail/auth/../gear/pinned-requirements.txt
# -c hail/auth/../hail/python/pinned-requirements.txt
# google-auth
certifi==2023.7.22
# via
# -c hail/auth/../gear/pinned-requirements.txt
# -c hail/auth/../hail/python/dev/pinned-requirements.txt
# -c hail/auth/../hail/python/pinned-requirements.txt
# requests
charset-normalizer==3.2.0
# via
# -c hail/auth/../gear/pinned-requirements.txt
# -c hail/auth/../hail/python/dev/pinned-requirements.txt
# -c hail/auth/../hail/python/pinned-requirements.txt
# -c hail/auth/../web_common/pinned-requirements.txt
# requests
google-auth==2.22.0
# via
# -c hail/auth/../gear/pinned-requirements.txt
# -c hail/auth/../hail/python/pinned-requirements.txt
# google-auth-oauthlib
google-auth-oauthlib==0.8.0
# via -r hail/auth/requirements.txt
idna==3.4
# via
# -c hail/auth/../gear/pinned-requirements.txt
# -c hail/auth/../hail/python/dev/pinned-requirements.txt
# -c hail/auth/../hail/python/pinned-requirements.txt
# -c hail/auth/../web_common/pinned-requirements.txt
# requests
oauthlib==3.2.2
# via
# -c hail/auth/../hail/python/pinned-requirements.txt
# requests-oauthlib
pyasn1==0.5.0
# via
# -c hail/auth/../gear/pinned-requirements.txt
# -c hail/auth/../hail/python/pinned-requirements.txt
# pyasn1-modules
# rsa
pyasn1-modules==0.3.0
# via
# -c hail/auth/../gear/pinned-requirements.txt
# -c hail/auth/../hail/python/pinned-requirements.txt
# google-auth
requests==2.31.0
# via
# -c hail/auth/../gear/pinned-requirements.txt
# -c hail/auth/../hail/python/dev/pinned-requirements.txt
# -c hail/auth/../hail/python/pinned-requirements.txt
# requests-oauthlib
requests-oauthlib==1.3.1
# via
# -c hail/auth/../hail/python/pinned-requirements.txt
# google-auth-oauthlib
rsa==4.9
# via
# -c hail/auth/../gear/pinned-requirements.txt
# -c hail/auth/../hail/python/pinned-requirements.txt
# google-auth
six==1.16.0
# via
# -c hail/auth/../gear/pinned-requirements.txt
# -c hail/auth/../hail/python/dev/pinned-requirements.txt
# -c hail/auth/../hail/python/pinned-requirements.txt
# google-auth
urllib3==1.26.16
# via
# -c hail/auth/../gear/pinned-requirements.txt
# -c hail/auth/../hail/python/dev/pinned-requirements.txt
# -c hail/auth/../hail/python/pinned-requirements.txt
# google-auth
# requests
1 change: 0 additions & 1 deletion auth/requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,3 @@
-c ../hail/python/dev/pinned-requirements.txt
-c ../gear/pinned-requirements.txt
-c ../web_common/pinned-requirements.txt
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We basically don't need this at all at this point

google-auth-oauthlib>=0.5.2,<1
2 changes: 1 addition & 1 deletion batch/Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ include ../config.mk
build:
$(MAKE) -C .. batch-image batch-worker-image

JINJA_ENVIRONMENT = '{"code":{"sha":"$(shell git rev-parse --short=12 HEAD)"},"deploy":$(DEPLOY),"batch_image":{"image":"$(shell cat ../batch-image)"},"batch_worker_image":{"image":"$(shell cat ../batch-worker-image)"},"default_ns":{"name":"$(NAMESPACE)"},"batch_database":{"user_secret_name":"sql-batch-user-config"},"scope":"$(SCOPE)","global":{"docker_prefix":"$(DOCKER_PREFIX)"}}'
JINJA_ENVIRONMENT = '{"code":{"sha":"$(shell git rev-parse --short=12 HEAD)"},"deploy":$(DEPLOY),"batch_image":{"image":"$(shell cat ../batch-image)"},"batch_worker_image":{"image":"$(shell cat ../batch-worker-image)"},"default_ns":{"name":"$(NAMESPACE)"},"batch_database":{"user_secret_name":"sql-batch-user-config"},"scope":"$(SCOPE)","global":{"docker_prefix":"$(DOCKER_PREFIX)","cloud":"$(CLOUD)"}}'

.PHONY: deploy
deploy: build
Expand Down
4 changes: 4 additions & 0 deletions batch/batch/cloud/azure/driver/create_instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ def create_vm_config(
) -> dict:
_, cores = azure_machine_type_to_worker_type_and_cores(machine_type)

hail_azure_oauth_scope = os.environ['HAIL_AZURE_OAUTH_SCOPE']
region = instance_config.region_for(location)

if max_price is not None and not preemptible:
Expand Down Expand Up @@ -205,6 +206,7 @@ def create_vm_config(
DOCKER_ROOT_IMAGE=$(jq -r '.docker_root_image' userdata)
DOCKER_PREFIX=$(jq -r '.docker_prefix' userdata)
REGION=$(jq -r '.region' userdata)
HAIL_AZURE_OAUTH_SCOPE=$(jq -r '.hail_azure_oauth_scope' userdata)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This needs to make its way into some dev doc. What is it, why do we need it, where do we specify it, how much does it matter if you get it wrong.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll draft a dev doc alongside this. Happy to block this PR on such a doc but dismissing comments for now to get more feedback.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just a reminder of this before my approval.


INTERNAL_GATEWAY_IP=$(jq -r '.internal_ip' userdata)

Expand Down Expand Up @@ -258,6 +260,7 @@ def create_vm_config(
-e RESOURCE_GROUP=$RESOURCE_GROUP \
-e LOCATION=$LOCATION \
-e REGION=$REGION \
-e HAIL_AZURE_OAUTH_SCOPE=$HAIL_AZURE_OAUTH_SCOPE \
-e DOCKER_PREFIX=$DOCKER_PREFIX \
-e DOCKER_ROOT_IMAGE=$DOCKER_ROOT_IMAGE \
-e INSTANCE_CONFIG=$INSTANCE_CONFIG \
Expand Down Expand Up @@ -313,6 +316,7 @@ def create_vm_config(
'max_idle_time_msecs': max_idle_time_msecs,
'instance_config': base64.b64encode(json.dumps(instance_config.to_dict()).encode()).decode(),
'region': region,
'hail_azure_oauth_scope': hail_azure_oauth_scope,
}
user_data_str = base64.b64encode(json.dumps(user_data).encode('utf-8')).decode('utf-8')

Expand Down
6 changes: 6 additions & 0 deletions batch/batch/cloud/azure/worker/credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
import json
from typing import Dict

from hailtop.auth.auth import IdentityProvider

from ....worker.credentials import CloudUserCredentials


Expand All @@ -26,6 +28,10 @@ def password(self):
def mount_path(self):
return '/azure-credentials/key.json'

@property
def identity_json(self):
return {'idp': IdentityProvider.MICROSOFT.value}

def blobfuse_credentials(self, account: str, container: str) -> str:
# https://github.com/Azure/azure-storage-fuse
return f'''
Expand Down
Loading