Feat: Enable flytekit to authenticate with proxy in front of FlyteAdmin

[fg91]

TL;DR

Part of an effort to integrate Flyte with GCP Identity Aware Proxy (IAP).
See tracking issue for details and motivation: flyteorg/flyte#3965

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested (Deployed Flyte with IAP in sandbox cluster according to deployment guide in the readme added in this PR. Tested that FlyteRemote can talk to admin through the IAP, e.g. to retrieve executions, start new ones etc.)
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

• Add a new option called proxyCommand to the platform config in which users can specify an external command which, if specified, is used to generate ID tokens for a proxy in front of Flyte.
  • This idea is derived from the existing command (external command auth type) which is used to generate tokens for Flyte itself.
• The ID tokens generated by this proxyCommand are added as "proxy-authorization" header to every request (http, gRPC) that flytekit's client makes to flyteadmin, including the initial unauthenticated requests of the auth flow with flyteadmin. These additional ID tokens in the "proxy-authorization" header allow clients to interact with a flyteadmin that is protected by GCP Identity Aware Proxy (or in theory similar services from other providers)
• Refactor the AuthorizationClient which currently can only be used for a pkce auth flow (despite its doc string claiming it works for a general OAuth 2.0 flow) to actually work for both pkce and default OAuth 2.0.
• Add a plugin which provides a CLI that generates ID tokens for GCP Identity Aware Proxy
  • This plugin re-uses the AuthorizationClient in flytekit to perform an OAuth 2.0 flow with accounts.google.com.
  • The command is to be used as proxyCommand in the platform config for Flyte deployments protected with IAP
• Add a step-by-step guide on how to deploy Flyte protected with GCP IAP

Tracking Issue

https://github.com/flyteorg/flyte/issues/

Follow-up issue

NA
OR
https://github.com/flyteorg/flyte/issues/

[fg91]

TODO: do we need the legacy config entry or do we only support this in yaml config?

[wild-endeavor]

i think this is the correct declaration. Let's keep it. But what is this? what does this do?

[fg91]

Added an explanation in this PR's description.

[fg91]

When trying to reach flyteadmin behind GCP Identity Aware Proxy and when omitting the "proxy-authorization" header, we get back:

<_InactiveRpcError of RPC that terminated with:
        status = StatusCode.UNKNOWN
        details = "Stream removed"
        debug_error_string = "UNKNOWN:Error received from peer ipv4:<my-ip>:443 {grpc_message:"Stream removed", grpc_status:2, created_time:"2023-xxx"}"

[fg91]

@wild-endeavor I merged #1795 into this one as requested 👍

[fg91]

The order of mocking vs instantiation of the authenticator matters now.
I added the more explicit check assert authn._creds.access_token == "abc" because this would fail in case somebody in the future reverses the order again (instantiate the authenticator first, then mock).

[fg91]

Let's be more specific here. When adapting the magic mocks, I noticed that here one easily gets a MagicMock object for t which then also is not None but is not "abc".

This makes sure that in the future no one accidentally breaks the mocking.

[fg91]

Thanks for taking a look @wild-endeavor

I added a second subcommand to flyte-iap to generate ID tokens for service accounts:

~ flyte-iap --help 
Usage: flyte-iap [OPTIONS] COMMAND [ARGS]...

  Generate ID tokens for GCP Identity Aware Proxy (IAP).

Options:
  --help  Show this message and exit.

Commands:
  generate-service-account-id-token                                         # <- new
                                  Generate a service account ID token...
  generate-user-id-token          Generate a user account ID token for...

The subcommand uses either

• a service account key json file or
• contacts the GCP metadata server to obtain a token when executed in Compute Engine etc.

To generate a token for service accounts, in contrast to generating one for user accounts, no browser needs to be opened interactively. The goal of this subcommand is that CICD pipelines can talk to flyteadmin through IAP. (@corleyma said they will require this.)

With that, flyte-iap is complete.

flytectl still doesn't understand the proxyCommand in the config but I will work on this next:

❯ flytectl config init
Error:

1 error(s) decoding:

* '' has invalid keys: proxycommand
ERRO[0000]

1 error(s) decoding:

* '' has invalid keys: proxycommand  src="main.go:13"

[fg91]

i think this looks good. cc @EngHabu to take a look though.

Tagging you as reviewer, @EngHabu, as suggested by @wild-endeavor 🙏

[EngHabu]

I'm less sure about reusing this command for flytectl. Requiring python dependency isn't something I would want to require in CI/CD and other automated systems

[fg91]

Hi Haytham, I really appreciate the time you’ve invested in looking at this PR, I'm sure you have a busy schedule 🙏

I'm less sure about reusing this command for flytectl. Requiring python dependency isn't something I would want to require in CI/CD and other automated systems

In our case, we will run pyflyte register/package in CICD as well but I agree with your concern that this shouldn't be required. I documented 2 additional ways to obtain an id token for service accounts using 1) GCP's gcloud sdk and 2) a bash script from the identity aware proxy documentation. This way, users can choose to not use python.
2437b84

[wild-endeavor]

tested this with pkce, cc, and device flow btw, all works

[fg91]

tested this with pkce, cc, and device flow btw, all works

Thanks for testing this @wild-endeavor! 🙏
I also implemented the same logic for flytectl/the admin client in idl. Everything works, just need to polish a bit and adapt/add tests. Will ping you there as well then.

[jtyberg]

I'm trying to get all Flyte clients to pass the additional proxy-authorization header with each request.

I can use the flytekit-identity-aware-proxy plugin (with proxyCommand in the Flyte config) successfully from a FlyteRemote script, but I'm seeing this error when trying to use flytectl.

bin/flytectl get projects                                                 

Error: 

1 error(s) decoding:

* '' has invalid keys: proxycommand
ERRO[0000] 

1 error(s) decoding:

* '' has invalid keys: proxycommand  src="main.go:13"

It seems the issue is that flytekit and flytectl are stuck at different versions of flyteidl? (this PR never made it in).

[fg91]

I'm trying to get all Flyte clients to pass the additional proxy-authorization header with each request.

I can use the flytekit-identity-aware-proxy plugin (with proxyCommand in the Flyte config) successfully from a FlyteRemote script, but I'm seeing this error when trying to use flytectl.

bin/flytectl get projects                                                 



Error: 



1 error(s) decoding:



* '' has invalid keys: proxycommand

ERRO[0000] 



1 error(s) decoding:



* '' has invalid keys: proxycommand  src="main.go:13"

It seems the issue is that flytekit and flytectl are stuck at different versions of flyteidl? (this PR never made it in).

The most recent versions of flytectl (released after the move to the monorepo) use the flyteidl admin client which supports proxy auth, try upgrading to e.g. >= 0.8.18 please.

[jtyberg]

The most recent versions of flytectl (released after the move to the monorepo) use the flyteidl admin client which supports proxy auth, try upgrading to e.g. >= 0.8.18 please.

I apologize, I should have posted my flytectl version, which is the v0.8.18 that you recommend.

Installed using curl -sL https://ctl.flyte.org/install | bash.

bin/flytectl version

 A new release of flytectl is available: 0.8.18 → v0.8.22 
To upgrade, run: flytectl upgrade 
https://github.com/flyteorg/flytectl/releases/tag/v0.8.22 

{
  "App": "flytectl",
  "Build": "0a0cbce",
  "Version": "0.8.18",
  "BuildTime": "2024-06-03 10:50:36.469871 -0400 EDT m=+0.045492501"
}%                                                                                                                      

v0.8.18 still shows the error when I add proxyCommand to the flytectl config.

bin/flytectl version

Error: 

1 error(s) decoding:

* '' has invalid keys: proxycommand

[fg91]

I apologize, I should have posted my flytectl version, which is the v0.8.18 that you recommend.

Sorry, you are right about v0.8.18 (was on my phone). Can you please try v0.8.23? The install script pointed to the archived flytectl repo until this fix a few days ago.

[jtyberg]

I apologize, I should have posted my flytectl version, which is the v0.8.18 that you recommend.

Sorry, you are right about v0.8.18 (was on my phone). Can you please try v0.8.23? The install script pointed to the archived flytectl repo until this fix a few days ago.

Ah, v0.8.23 did it. I can now use the proxyCommand field in my Flyte config.yaml.

bin/flytectl get projects                    

 ----- ------ ----------------- 
| ID  | NAME | DESCRIPTION     |
 ----- ------ ----------------- 
| myproject | myproject  | myproject description |
 ----- ------ ----------------- 
1 rows

Thanks for this! (And apologies for opening an old thread).