fix: add distributed protection for external auth refresh (#27479)

Before refreshing, the caller has to put a lease on the row to prevent
any other replicas from trying to refresh using the same token. Callers
unable to get a lease enter a holding pattern waiting for the results.

Also use this same lease query to re-read the link. This more or less
functions the same as the re-read we used to have after a refresh
failure, except it also allows us to avoid making the failed request in
the first place. The downside is that we have to refetch the link every
time even if the caller just fetched it, but since we have to make a
database call to set the lease anyway...
This commit is contained in:
Asher
2026-08-19 13:00:31 -08:00
committed by GitHub
parent 24b83dc727
commit 7ba2e6e22a
18 changed files with 1272 additions and 603 deletions
+54 -20
View File
@@ -732,6 +732,60 @@ CREATE TYPE workspace_transition AS ENUM (
'delete'
);
CREATE TABLE external_auth_links (
provider_id text NOT NULL,
user_id uuid NOT NULL,
created_at timestamp with time zone NOT NULL,
updated_at timestamp with time zone NOT NULL,
oauth_access_token text NOT NULL,
oauth_refresh_token text NOT NULL,
oauth_expiry timestamp with time zone NOT NULL,
oauth_access_token_key_id text,
oauth_refresh_token_key_id text,
oauth_extra jsonb,
oauth_refresh_failure_reason text DEFAULT ''::text NOT NULL,
refresh_lease_expires_at timestamp with time zone
);
COMMENT ON COLUMN external_auth_links.oauth_access_token_key_id IS 'The ID of the key used to encrypt the OAuth access token. If this is NULL, the access token is not encrypted';
COMMENT ON COLUMN external_auth_links.oauth_refresh_token_key_id IS 'The ID of the key used to encrypt the OAuth refresh token. If this is NULL, the refresh token is not encrypted';
COMMENT ON COLUMN external_auth_links.oauth_refresh_failure_reason IS 'This error means the refresh token is invalid. Cached so we can avoid calling the external provider again for the same error.';
COMMENT ON COLUMN external_auth_links.refresh_lease_expires_at IS 'Indicates a replica is refreshing the token; prevents concurrent refreshes.';
CREATE FUNCTION acquire_external_auth_link_refresh_lease(arg_provider_id text, arg_user_id uuid, timeout_ms bigint) RETURNS SETOF external_auth_links
LANGUAGE plpgsql
AS $$
DECLARE r external_auth_links;
BEGIN
UPDATE external_auth_links
SET
refresh_lease_expires_at = NOW() + (timeout_ms || ' ms')::interval
WHERE
provider_id = arg_provider_id
AND user_id = arg_user_id
AND (refresh_lease_expires_at IS NULL OR refresh_lease_expires_at < NOW())
RETURNING * INTO r;
-- Got the lease, return the one row.
IF FOUND THEN
RETURN NEXT r;
RETURN;
END IF;
-- Differentiate between unable to get the lease and the row being gone.
IF EXISTS (SELECT 1 FROM external_auth_links WHERE provider_id = arg_provider_id AND user_id = arg_user_id) THEN
RAISE EXCEPTION 'row is currently leased by another replica'
USING ERRCODE = 'check_violation',
CONSTRAINT = 'external_auth_link_active_lease';
END IF;
-- Row is gone, return nothing.
RETURN;
END;
$$;
COMMENT ON FUNCTION acquire_external_auth_link_refresh_lease(arg_provider_id text, arg_user_id uuid, timeout_ms bigint) IS 'Acquire a lease on the external auth link and return the row. If there is already an active lease, an exception is raised.';
CREATE FUNCTION aggregate_usage_event() RETURNS trigger
LANGUAGE plpgsql
AS $$
@@ -2333,26 +2387,6 @@ COMMENT ON COLUMN dbcrypt_keys.revoked_at IS 'The time at which the key was revo
COMMENT ON COLUMN dbcrypt_keys.test IS 'A column used to test the encryption.';
CREATE TABLE external_auth_links (
provider_id text NOT NULL,
user_id uuid NOT NULL,
created_at timestamp with time zone NOT NULL,
updated_at timestamp with time zone NOT NULL,
oauth_access_token text NOT NULL,
oauth_refresh_token text NOT NULL,
oauth_expiry timestamp with time zone NOT NULL,
oauth_access_token_key_id text,
oauth_refresh_token_key_id text,
oauth_extra jsonb,
oauth_refresh_failure_reason text DEFAULT ''::text NOT NULL
);
COMMENT ON COLUMN external_auth_links.oauth_access_token_key_id IS 'The ID of the key used to encrypt the OAuth access token. If this is NULL, the access token is not encrypted';
COMMENT ON COLUMN external_auth_links.oauth_refresh_token_key_id IS 'The ID of the key used to encrypt the OAuth refresh token. If this is NULL, the refresh token is not encrypted';
COMMENT ON COLUMN external_auth_links.oauth_refresh_failure_reason IS 'This error means the refresh token is invalid. Cached so we can avoid calling the external provider again for the same error.';
CREATE TABLE files (
hash character varying(64) NOT NULL,
created_at timestamp with time zone NOT NULL,