name: vector<u8>

The utf-8 encoded issuer string. E.g., b"https://www.facebook.com".

config_url: vector<u8>

The ut8-8 encoded OpenID configuration URL of the provider. E.g., b"https://www.facebook.com/.well-known/openid-configuration/".

providers: vector<jwks::OIDCProvider>

id: vector<u8>

payload: vector<u8>

kid: string::String

kty: string::String

alg: string::String

e: string::String

n: string::String

variant: copyable_any::Any

A JWK variant packed as an Any. Currently the variant type is one of the following. - RSA_JWK - UnsupportedJWK

issuer: vector<u8>

The utf-8 encoding of the issuer string (e.g., "https://www.facebook.com").

version: u64

A version number is needed by JWK consensus to dedup the updates. e.g, when on chain version = 5, multiple nodes can propose an update with version = 6. Bumped every time the JWKs for the current issuer is updated. The Rust authenticator only uses the latest version.

jwks: vector<jwks::JWK>

Vector of JWK's sorted by their unique ID (from get_jwk_id) in dictionary order.

entries: vector<jwks::ProviderJWKs>

Vector of ProviderJWKs sorted by ProviderJWKs::issuer in dictionary order.

jwks: jwks::AllProvidersJWKs

epoch: u64

jwks: jwks::AllProvidersJWKs

variant: copyable_any::Any

A Patch variant packed as an Any. Currently the variant type is one of the following. - PatchRemoveAll - PatchRemoveIssuer - PatchRemoveJWK - PatchUpsertJWK

dummy_field: bool

issuer: vector<u8>

issuer: vector<u8>

jwk_id: vector<u8>

issuer: vector<u8>

jwk: jwks::JWK

patches: vector<jwks::Patch>

jwks: jwks::AllProvidersJWKs

jwks: jwks::AllProvidersJWKs

public fun patch_federated_jwks(jwk_owner: &signer, patches: vector<Patch>) acquires FederatedJWKs {
    // Prevents accidental calls in 0x1::jwks that install federated JWKs at the Aptos framework address.
    assert!(!system_addresses::is_aptos_framework_address(signer::address_of(jwk_owner)),
        error::invalid_argument(EINSTALL_FEDERATED_JWKS_AT_APTOS_FRAMEWORK)
    );

    let jwk_addr = signer::address_of(jwk_owner);
    if (!exists<FederatedJWKs>(jwk_addr)) {
        move_to(jwk_owner, FederatedJWKs { jwks: AllProvidersJWKs { entries: vector[] } });
    };

    let fed_jwks = borrow_global_mut<FederatedJWKs>(jwk_addr);
    patches.for_each_ref(|obj|{
        let patch: &Patch = obj;
        apply_patch(&mut fed_jwks.jwks, *patch);
    });

    // TODO: Can we check the size more efficiently instead of serializing it via BCS?
    let num_bytes = bcs::to_bytes(fed_jwks).length();
    assert!(num_bytes < MAX_FEDERATED_JWKS_SIZE_BYTES, error::invalid_argument(EFEDERATED_JWKS_TOO_LARGE));
}

public entry fun update_federated_jwk_set(jwk_owner: &signer, iss: vector<u8>, kid_vec: vector<String>, alg_vec: vector<String>, e_vec: vector<String>, n_vec: vector<String>) acquires FederatedJWKs {
    assert!(!kid_vec.is_empty(), error::invalid_argument(EINVALID_FEDERATED_JWK_SET));
    let num_jwk = kid_vec.length<String>();
    assert!(alg_vec.length() == num_jwk , error::invalid_argument(EINVALID_FEDERATED_JWK_SET));
    assert!(e_vec.length() == num_jwk, error::invalid_argument(EINVALID_FEDERATED_JWK_SET));
    assert!(n_vec.length() == num_jwk, error::invalid_argument(EINVALID_FEDERATED_JWK_SET));

    let remove_all_patch = new_patch_remove_all();
    let patches = vector[remove_all_patch];
    while (!kid_vec.is_empty()) {
        let kid = kid_vec.pop_back();
        let alg = alg_vec.pop_back();
        let e = e_vec.pop_back();
        let n = n_vec.pop_back();
        let jwk = new_rsa_jwk(kid, alg, e, n);
        let patch = new_patch_upsert_jwk(iss, jwk);
        patches.push_back(patch)
    };
    patch_federated_jwks(jwk_owner, patches);
}

public fun get_patched_jwk(issuer: vector<u8>, jwk_id: vector<u8>): JWK acquires PatchedJWKs {
    try_get_patched_jwk(issuer, jwk_id).extract()
}

public fun try_get_patched_jwk(issuer: vector<u8>, jwk_id: vector<u8>): Option<JWK> acquires PatchedJWKs {
    let jwks = &borrow_global<PatchedJWKs>(@aptos_framework).jwks;
    try_get_jwk_by_issuer(jwks, issuer, jwk_id)
}

public fun upsert_oidc_provider(fx: &signer, name: vector<u8>, config_url: vector<u8>): Option<vector<u8>> acquires SupportedOIDCProviders {
    system_addresses::assert_aptos_framework(fx);
    chain_status::assert_genesis();

    let provider_set = borrow_global_mut<SupportedOIDCProviders>(@aptos_framework);

    let old_config_url= remove_oidc_provider_internal(provider_set, name);
    provider_set.providers.push_back(OIDCProvider { name, config_url });
    old_config_url
}

public fun upsert_oidc_provider_for_next_epoch(fx: &signer, name: vector<u8>, config_url: vector<u8>): Option<vector<u8>> acquires SupportedOIDCProviders {
    system_addresses::assert_aptos_framework(fx);

    let provider_set = if (config_buffer::does_exist<SupportedOIDCProviders>()) {
        config_buffer::extract_v2<SupportedOIDCProviders>()
    } else {
        *borrow_global<SupportedOIDCProviders>(@aptos_framework)
    };

    let old_config_url = remove_oidc_provider_internal(&mut provider_set, name);
    provider_set.providers.push_back(OIDCProvider { name, config_url });
    config_buffer::upsert(provider_set);
    old_config_url
}

public fun remove_oidc_provider(fx: &signer, name: vector<u8>): Option<vector<u8>> acquires SupportedOIDCProviders {
    system_addresses::assert_aptos_framework(fx);
    chain_status::assert_genesis();

    let provider_set = borrow_global_mut<SupportedOIDCProviders>(@aptos_framework);
    remove_oidc_provider_internal(provider_set, name)
}

public fun remove_oidc_provider_for_next_epoch(fx: &signer, name: vector<u8>): Option<vector<u8>> acquires SupportedOIDCProviders {
    system_addresses::assert_aptos_framework(fx);

    let provider_set = if (config_buffer::does_exist<SupportedOIDCProviders>()) {
        config_buffer::extract_v2<SupportedOIDCProviders>()
    } else {
        *borrow_global<SupportedOIDCProviders>(@aptos_framework)
    };
    let ret = remove_oidc_provider_internal(&mut provider_set, name);
    config_buffer::upsert(provider_set);
    ret
}

public(friend) fun on_new_epoch(framework: &signer) acquires SupportedOIDCProviders {
    system_addresses::assert_aptos_framework(framework);
    if (config_buffer::does_exist<SupportedOIDCProviders>()) {
        let new_config = config_buffer::extract_v2<SupportedOIDCProviders>();
        if (exists<SupportedOIDCProviders>(@aptos_framework)) {
            *borrow_global_mut<SupportedOIDCProviders>(@aptos_framework) = new_config;
        } else {
            move_to(framework, new_config);
        }
    }
}

public fun set_patches(fx: &signer, patches: vector<Patch>) acquires Patches, PatchedJWKs, ObservedJWKs {
    system_addresses::assert_aptos_framework(fx);
    borrow_global_mut<Patches>(@aptos_framework).patches = patches;
    regenerate_patched_jwks();
}

public fun new_patch_remove_all(): Patch {
    Patch {
        variant: copyable_any::pack(PatchRemoveAll {}),
    }
}

public fun new_patch_remove_issuer(issuer: vector<u8>): Patch {
    Patch {
        variant: copyable_any::pack(PatchRemoveIssuer { issuer }),
    }
}

public fun new_patch_remove_jwk(issuer: vector<u8>, jwk_id: vector<u8>): Patch {
    Patch {
        variant: copyable_any::pack(PatchRemoveJWK { issuer, jwk_id })
    }
}

public fun new_patch_upsert_jwk(issuer: vector<u8>, jwk: JWK): Patch {
    Patch {
        variant: copyable_any::pack(PatchUpsertJWK { issuer, jwk })
    }
}

public fun new_rsa_jwk(kid: String, alg: String, e: String, n: String): JWK {
    JWK {
        variant: copyable_any::pack(RSA_JWK {
            kid,
            kty: utf8(b"RSA"),
            e,
            n,
            alg,
        }),
    }
}

public fun new_unsupported_jwk(id: vector<u8>, payload: vector<u8>): JWK {
    JWK {
        variant: copyable_any::pack(UnsupportedJWK { id, payload })
    }
}

public fun initialize(fx: &signer) {
    system_addresses::assert_aptos_framework(fx);
    move_to(fx, SupportedOIDCProviders { providers: vector[] });
    move_to(fx, ObservedJWKs { jwks: AllProvidersJWKs { entries: vector[] } });
    move_to(fx, Patches { patches: vector[] });
    move_to(fx, PatchedJWKs { jwks: AllProvidersJWKs { entries: vector[] } });
}

public fun initialize_with_defaults(fx: &signer, providers: vector<OIDCProvider>, patches: vector<Patch>) {
    system_addresses::assert_aptos_framework(fx);
    move_to(fx, SupportedOIDCProviders { providers });
    move_to(fx, ObservedJWKs { jwks: AllProvidersJWKs { entries: vector[] } });
    move_to(fx, Patches { patches });
    move_to(fx, PatchedJWKs { jwks: AllProvidersJWKs { entries: vector[] } });
    regenerate_patched_jwks();
}

fun remove_oidc_provider_internal(provider_set: &mut SupportedOIDCProviders, name: vector<u8>): Option<vector<u8>> {
    let (name_exists, idx) = provider_set.providers.find(|obj| {
        let provider: &OIDCProvider = obj;
        provider.name == name
    });

    if (name_exists) {
        let old_provider = provider_set.providers.swap_remove(idx);
        option::some(old_provider.config_url)
    } else {
        option::none()
    }
}

public fun upsert_into_observed_jwks(fx: &signer, provider_jwks_vec: vector<ProviderJWKs>) acquires ObservedJWKs, PatchedJWKs, Patches {
    system_addresses::assert_aptos_framework(fx);
    let observed_jwks = borrow_global_mut<ObservedJWKs>(@aptos_framework);

    if (features::is_jwk_consensus_per_key_mode_enabled()) {
        provider_jwks_vec.for_each(|proposed_provider_jwks|{
            let maybe_cur_issuer_jwks = remove_issuer(&mut observed_jwks.jwks, proposed_provider_jwks.issuer);
            let cur_issuer_jwks = if (option::is_some(&maybe_cur_issuer_jwks)) {
                option::extract(&mut maybe_cur_issuer_jwks)
            } else {
                ProviderJWKs {
                    issuer: proposed_provider_jwks.issuer,
                    version: 0,
                    jwks: vector[],
                }
            };
            assert!(cur_issuer_jwks.version + 1 == proposed_provider_jwks.version, error::invalid_argument(EUNEXPECTED_VERSION));
            vector::for_each(proposed_provider_jwks.jwks, |jwk|{
                let variant_type_name = *string::bytes(copyable_any::type_name(&jwk.variant));
                let is_delete = if (variant_type_name == b"0x1::jwks::UnsupportedJWK") {
                    let repr = copyable_any::unpack<UnsupportedJWK>(jwk.variant);
                    &repr.payload == &DELETE_COMMAND_INDICATOR
                } else {
                    false
                };
                if (is_delete) {
                    remove_jwk(&mut cur_issuer_jwks, get_jwk_id(&jwk));
                } else {
                    upsert_jwk(&mut cur_issuer_jwks, jwk);
                }
            });
            cur_issuer_jwks.version += 1;
            upsert_provider_jwks(&mut observed_jwks.jwks, cur_issuer_jwks);
        });
    } else {
        provider_jwks_vec.for_each(|provider_jwks| {
            upsert_provider_jwks(&mut observed_jwks.jwks, provider_jwks);
        });
    };

    let epoch = reconfiguration::current_epoch();
    emit(ObservedJWKsUpdated { epoch, jwks: observed_jwks.jwks });
    regenerate_patched_jwks();
}

public fun remove_issuer_from_observed_jwks(fx: &signer, issuer: vector<u8>): Option<ProviderJWKs> acquires ObservedJWKs, PatchedJWKs, Patches {
    system_addresses::assert_aptos_framework(fx);
    let observed_jwks = borrow_global_mut<ObservedJWKs>(@aptos_framework);
    let old_value = remove_issuer(&mut observed_jwks.jwks, issuer);

    let epoch = reconfiguration::current_epoch();
    emit(ObservedJWKsUpdated { epoch, jwks: observed_jwks.jwks });
    regenerate_patched_jwks();

    old_value
}

fun regenerate_patched_jwks() acquires PatchedJWKs, Patches, ObservedJWKs {
    let jwks = borrow_global<ObservedJWKs>(@aptos_framework).jwks;
    let patches = borrow_global<Patches>(@aptos_framework);
    patches.patches.for_each_ref(|obj|{
        let patch: &Patch = obj;
        apply_patch(&mut jwks, *patch);
    });
    *borrow_global_mut<PatchedJWKs>(@aptos_framework) = PatchedJWKs { jwks };
}

fun try_get_jwk_by_issuer(jwks: &AllProvidersJWKs, issuer: vector<u8>, jwk_id: vector<u8>): Option<JWK> {
    let (issuer_found, index) = jwks.entries.find(|obj| {
        let provider_jwks: &ProviderJWKs = obj;
        issuer == provider_jwks.issuer
    });

    if (issuer_found) {
        try_get_jwk_by_id(jwks.entries.borrow(index), jwk_id)
    } else {
        option::none()
    }
}

fun try_get_jwk_by_id(provider_jwks: &ProviderJWKs, jwk_id: vector<u8>): Option<JWK> {
    let (jwk_id_found, index) = provider_jwks.jwks.find(|obj|{
        let jwk: &JWK = obj;
        jwk_id == get_jwk_id(jwk)
    });

    if (jwk_id_found) {
        option::some(provider_jwks.jwks[index])
    } else {
        option::none()
    }
}

fun get_jwk_id(jwk: &JWK): vector<u8> {
    let variant_type_name = *jwk.variant.type_name().bytes();
    if (variant_type_name == b"0x1::jwks::RSA_JWK") {
        let rsa = jwk.variant.unpack<RSA_JWK>();
        *rsa.kid.bytes()
    } else if (variant_type_name == b"0x1::jwks::UnsupportedJWK") {
        let unsupported = jwk.variant.unpack<UnsupportedJWK>();
        unsupported.id
    } else {
        abort(error::invalid_argument(EUNKNOWN_JWK_VARIANT))
    }
}

fun upsert_provider_jwks(jwks: &mut AllProvidersJWKs, provider_jwks: ProviderJWKs): Option<ProviderJWKs> {
    // NOTE: Using a linear-time search here because we do not expect too many providers.
    let found = false;
    let index = 0;
    let num_entries = jwks.entries.length();
    while (index < num_entries) {
        let cur_entry = jwks.entries.borrow(index);
        let comparison = compare_u8_vector(provider_jwks.issuer, cur_entry.issuer);
        if (comparison.is_greater_than()) {
            index += 1;
        } else {
            found = comparison.is_equal();
            break
        }
    };

    // Now if `found == true`, `index` points to the JWK we want to update/remove; otherwise, `index` points to
    // where we want to insert.
    let ret = if (found) {
        let entry = jwks.entries.borrow_mut(index);
        let old_entry = option::some(*entry);
        *entry = provider_jwks;
        old_entry
    } else {
        jwks.entries.insert(index, provider_jwks);
        option::none()
    };

    ret
}

fun remove_issuer(jwks: &mut AllProvidersJWKs, issuer: vector<u8>): Option<ProviderJWKs> {
    let (found, index) = jwks.entries.find(|obj| {
        let provider_jwk_set: &ProviderJWKs = obj;
        provider_jwk_set.issuer == issuer
    });

    let ret = if (found) {
        option::some(jwks.entries.remove(index))
    } else {
        option::none()
    };

    ret
}

fun upsert_jwk(set: &mut ProviderJWKs, jwk: JWK): Option<JWK> {
    let found = false;
    let index = 0;
    let num_entries = set.jwks.length();
    while (index < num_entries) {
        let cur_entry = set.jwks.borrow(index);
        let comparison = compare_u8_vector(get_jwk_id(&jwk), get_jwk_id(cur_entry));
        if (comparison.is_greater_than()) {
            index += 1;
        } else {
            found = comparison.is_equal();
            break
        }
    };

    // Now if `found == true`, `index` points to the JWK we want to update/remove; otherwise, `index` points to
    // where we want to insert.
    let ret = if (found) {
        let entry = set.jwks.borrow_mut(index);
        let old_entry = option::some(*entry);
        *entry = jwk;
        old_entry
    } else {
        set.jwks.insert(index, jwk);
        option::none()
    };

    ret
}

fun remove_jwk(jwks: &mut ProviderJWKs, jwk_id: vector<u8>): Option<JWK> {
    let (found, index) = jwks.jwks.find(|obj| {
        let jwk: &JWK = obj;
        jwk_id == get_jwk_id(jwk)
    });

    let ret = if (found) {
        option::some(jwks.jwks.remove(index))
    } else {
        option::none()
    };

    ret
}

fun apply_patch(jwks: &mut AllProvidersJWKs, patch: Patch) {
    let variant_type_name = *patch.variant.type_name().bytes();
    if (variant_type_name == b"0x1::jwks::PatchRemoveAll") {
        jwks.entries = vector[];
    } else if (variant_type_name == b"0x1::jwks::PatchRemoveIssuer") {
        let cmd = patch.variant.unpack<PatchRemoveIssuer>();
        remove_issuer(jwks, cmd.issuer);
    } else if (variant_type_name == b"0x1::jwks::PatchRemoveJWK") {
        let cmd = patch.variant.unpack<PatchRemoveJWK>();
        // TODO: This is inefficient: we remove the issuer, modify its JWKs & and reinsert the updated issuer. Why
        // not just update it in place?
        let existing_jwk_set = remove_issuer(jwks, cmd.issuer);
        if (existing_jwk_set.is_some()) {
            let jwk_set = existing_jwk_set.extract();
            remove_jwk(&mut jwk_set, cmd.jwk_id);
            upsert_provider_jwks(jwks, jwk_set);
        };
    } else if (variant_type_name == b"0x1::jwks::PatchUpsertJWK") {
        let cmd = patch.variant.unpack<PatchUpsertJWK>();
        // TODO: This is inefficient: we remove the issuer, modify its JWKs & and reinsert the updated issuer. Why
        // not just update it in place?
        let existing_jwk_set = remove_issuer(jwks, cmd.issuer);
        let jwk_set = if (existing_jwk_set.is_some()) {
            existing_jwk_set.extract()
        } else {
            ProviderJWKs {
                version: 0,
                issuer: cmd.issuer,
                jwks: vector[],
            }
        };
        upsert_jwk(&mut jwk_set, cmd.jwk);
        upsert_provider_jwks(jwks, jwk_set);
    } else {
        abort(std::error::invalid_argument(EUNKNOWN_PATCH_VARIANT))
    }
}