jwks - [mainnet]

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);
    vector::for_each_ref(&patches, |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 = vector::length(&bcs::to_bytes(fed_jwks));
    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!(!vector::is_empty(&kid_vec), error::invalid_argument(EINVALID_FEDERATED_JWK_SET));
    let num_jwk = vector::length<String>(&kid_vec);
    assert!(vector::length(&alg_vec) == num_jwk , error::invalid_argument(EINVALID_FEDERATED_JWK_SET));
    assert!(vector::length(&e_vec) == num_jwk, error::invalid_argument(EINVALID_FEDERATED_JWK_SET));
    assert!(vector::length(&n_vec) == num_jwk, error::invalid_argument(EINVALID_FEDERATED_JWK_SET));

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

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

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);
    vector::push_back(&mut provider_set.providers, 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);
    vector::push_back(&mut provider_set.providers, 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[] } });
}

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

    if (name_exists) {
        let old_provider = vector::swap_remove(&mut provider_set.providers, 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()) {
        vector::for_each(provider_jwks_vec, |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 = cur_issuer_jwks.version + 1;
            upsert_provider_jwks(&mut observed_jwks.jwks, cur_issuer_jwks);
        });
    } else {
        vector::for_each(provider_jwks_vec, |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);
    vector::for_each_ref(&patches.patches, |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) = vector::find(&jwks.entries, |obj| {
        let provider_jwks: &ProviderJWKs = obj;
        issuer == provider_jwks.issuer
    });

    if (issuer_found) {
        try_get_jwk_by_id(vector::borrow(&jwks.entries, 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) = vector::find(&provider_jwks.jwks, |obj|{
        let jwk: &JWK = obj;
        jwk_id == get_jwk_id(jwk)
    });

    if (jwk_id_found) {
        option::some(*vector::borrow(&provider_jwks.jwks, index))
    } else {
        option::none()
    }
}

fun get_jwk_id(jwk: &JWK): vector<u8> {
    let variant_type_name = *string::bytes(copyable_any::type_name(&jwk.variant));
    if (variant_type_name == b"0x1::jwks::RSA_JWK") {
        let rsa = copyable_any::unpack<RSA_JWK>(jwk.variant);
        *string::bytes(&rsa.kid)
    } else if (variant_type_name == b"0x1::jwks::UnsupportedJWK") {
        let unsupported = copyable_any::unpack<UnsupportedJWK>(jwk.variant);
        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 = vector::length(&jwks.entries);
    while (index < num_entries) {
        let cur_entry = vector::borrow(&jwks.entries, index);
        let comparison = compare_u8_vector(provider_jwks.issuer, cur_entry.issuer);
        if (is_greater_than(&comparison)) {
            index = index + 1;
        } else {
            found = is_equal(&comparison);
            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 = vector::borrow_mut(&mut jwks.entries, index);
        let old_entry = option::some(*entry);
        *entry = provider_jwks;
        old_entry
    } else {
        vector::insert(&mut jwks.entries, index, provider_jwks);
        option::none()
    };

    ret
}

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

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

    ret
}

fun upsert_jwk(set: &mut ProviderJWKs, jwk: JWK): Option<JWK> {
    let found = false;
    let index = 0;
    let num_entries = vector::length(&set.jwks);
    while (index < num_entries) {
        let cur_entry = vector::borrow(&set.jwks, index);
        let comparison = compare_u8_vector(get_jwk_id(&jwk), get_jwk_id(cur_entry));
        if (is_greater_than(&comparison)) {
            index = index + 1;
        } else {
            found = is_equal(&comparison);
            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 = vector::borrow_mut(&mut set.jwks, index);
        let old_entry = option::some(*entry);
        *entry = jwk;
        old_entry
    } else {
        vector::insert(&mut set.jwks, index, jwk);
        option::none()
    };

    ret
}

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

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

    ret
}

fun apply_patch(jwks: &mut AllProvidersJWKs, patch: Patch) {
    let variant_type_name = *string::bytes(copyable_any::type_name(&patch.variant));
    if (variant_type_name == b"0x1::jwks::PatchRemoveAll") {
        jwks.entries = vector[];
    } else if (variant_type_name == b"0x1::jwks::PatchRemoveIssuer") {
        let cmd = copyable_any::unpack<PatchRemoveIssuer>(patch.variant);
        remove_issuer(jwks, cmd.issuer);
    } else if (variant_type_name == b"0x1::jwks::PatchRemoveJWK") {
        let cmd = copyable_any::unpack<PatchRemoveJWK>(patch.variant);
        // 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 (option::is_some(&existing_jwk_set)) {
            let jwk_set = option::extract(&mut existing_jwk_set);
            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 = copyable_any::unpack<PatchUpsertJWK>(patch.variant);
        // 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 (option::is_some(&existing_jwk_set)) {
            option::extract(&mut existing_jwk_set)
        } 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))
    }
}