Skip to content

Commit

Permalink
feat: await fetchRootKey calls while making API calls
Browse files Browse the repository at this point in the history
move fetchRootKey calls from constructor to API calls

BREAKING CHANGE:
  • Loading branch information
nathanosdev committed Jan 13, 2025
1 parent 6b65198 commit 68839a6
Show file tree
Hide file tree
Showing 7 changed files with 66 additions and 35 deletions.
5 changes: 1 addition & 4 deletions packages/agent/src/actor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -531,9 +531,6 @@ function _createActorMethod(
const ecid = effectiveCanisterId !== undefined ? Principal.from(effectiveCanisterId) : cid;
const arg = IDL.encode(func.argTypes, args);

if (agent.rootKey == null)
throw new AgentError('Agent root key not initialized before making call');

const { requestId, response, requestDetails } = await agent.call(cid, {
methodName,
arg,
Expand All @@ -545,7 +542,7 @@ function _createActorMethod(
const cert = (response.body as v3ResponseBody).certificate;
certificate = await Certificate.create({
certificate: bufFromBufLike(cert),
rootKey: agent.rootKey,
rootKey: await agent.getRootKey(),
canisterId: Principal.from(canisterId),
blsVerify,
});
Expand Down
7 changes: 6 additions & 1 deletion packages/agent/src/agent/api.ts
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,12 @@ export interface SubmitResponse {
* An Agent able to make calls and queries to a Replica.
*/
export interface Agent {
readonly rootKey: ArrayBuffer | null;
/**
* Returns the root key of the network. This is used to verify the authenticity of
* responses from the network.
*/
getRootKey(): Promise<ArrayBuffer>;

/**
* Returns the principal ID associated with this agent (by default). It only shows
* the principal of the default identity in the agent, which is the principal used
Expand Down
3 changes: 1 addition & 2 deletions packages/agent/src/agent/http/http.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -800,10 +800,9 @@ test('retry requests that fail due to a network failure', async () => {
const agent = new HttpAgent({
host: HTTP_AGENT_HOST,
fetch: mockFetch,
rootKey: new Uint8Array(32),
});

agent.rootKey = new Uint8Array(32);

try {
await agent.call(Principal.managementCanister(), {
methodName: 'test',
Expand Down
56 changes: 38 additions & 18 deletions packages/agent/src/agent/http/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,14 @@ export interface HttpAgentOptions {
* Alternate root key to use for verifying certificates. If not provided, the default IC root key will be used.
*/
rootKey?: ArrayBuffer;

/**
* Whether the agent should fetch the root key from the network. Defaults to false.
*
* WARNING!!! Do not enable this in production environments,
* as it can be used as an attack vector by malicious nodes.
*/
shouldFetchRootKey?: boolean;
}

function getDefaultFetch(): typeof fetch {
Expand Down Expand Up @@ -245,19 +253,21 @@ other computations so that this class can stay as simple as possible while
allowing extensions.
*/
export class HttpAgent implements Agent {
public rootKey: ArrayBuffer;
#identity: Promise<Identity> | null;
readonly #fetch: typeof fetch;
readonly #fetchOptions?: Record<string, unknown>;
readonly #callOptions?: Record<string, unknown>;
#timeDiffMsecs = 0;
readonly host: URL;
readonly #credentials: string | undefined;
#rootKeyFetched = false;
readonly #retryTimes; // Retry requests N times before erroring by default
#backoffStrategy: BackoffStrategyFactory;
readonly #maxIngressExpiryInMinutes: number;

#rootKey: ArrayBuffer;
readonly #shouldFetchRootKey: boolean;
#fetchRootKeyPromise: Promise<ArrayBuffer> | null = null;

// Public signature to help with type checking.
public readonly _isAgent = true;
public config: HttpAgentOptions = {};
Expand Down Expand Up @@ -288,7 +298,11 @@ export class HttpAgent implements Agent {
this.#fetch = options.fetch || getDefaultFetch() || fetch.bind(global);
this.#fetchOptions = options.fetchOptions;
this.#callOptions = options.callOptions;
this.rootKey = options.rootKey ? options.rootKey : fromHex(IC_ROOT_KEY);

this.#rootKey = options.rootKey ?? fromHex(IC_ROOT_KEY);
this.#shouldFetchRootKey = options.shouldFetchRootKey ?? false;
// kick off the fetchRootKey process asynchronously, if needed
(async () => await this.fetchRootKey())();

const host = determineHost(options.host);
this.host = new URL(host);
Expand Down Expand Up @@ -354,16 +368,9 @@ export class HttpAgent implements Agent {
return new this({ ...options });
}

public static async create(
options: HttpAgentOptions & { shouldFetchRootKey?: boolean } = {
shouldFetchRootKey: false,
},
): Promise<HttpAgent> {
public static async create(options: HttpAgentOptions = {}): Promise<HttpAgent> {
const agent = HttpAgent.createSync(options);
const initPromises: Promise<ArrayBuffer | void>[] = [agent.syncTime()];
if (agent.host.toString() !== 'https://icp-api.io' && options.shouldFetchRootKey) {
initPromises.push(agent.fetchRootKey());
}
await Promise.all(initPromises);
return agent;
}
Expand Down Expand Up @@ -437,7 +444,7 @@ export class HttpAgent implements Agent {
): Promise<SubmitResponse> {
// TODO - restore this value
const callSync = options.callSync ?? true;
const id = await (identity !== undefined ? await identity : await this.#identity);
const id = identity !== undefined ? await identity : await this.#identity;
if (!id) {
throw new IdentityInvalidError(
"This identity has expired due this application's security policy. Please refresh your authentication.",
Expand Down Expand Up @@ -1125,14 +1132,27 @@ export class HttpAgent implements Agent {
return cbor.decode(await response.arrayBuffer());
}

async getRootKey(): Promise<ArrayBuffer> {
return this.fetchRootKey();
}

public async fetchRootKey(): Promise<ArrayBuffer> {
if (!this.#rootKeyFetched) {
const status = await this.status();
// Hex-encoded version of the replica root key
this.rootKey = (status as JsonObject & { root_key: ArrayBuffer }).root_key;
this.#rootKeyFetched = true;
if (!this.#shouldFetchRootKey) {
return this.#rootKey;
}

if (this.#fetchRootKeyPromise === null) {
this.#fetchRootKeyPromise = this.#fetchRootKey();
}
return this.rootKey;

return await this.#fetchRootKeyPromise;
}

async #fetchRootKey(): Promise<ArrayBuffer> {
const status = await this.status();
// Hex-encoded version of the replica root key
this.#rootKey = (status as JsonObject & { root_key: ArrayBuffer }).root_key;
return this.#rootKey;
}

public invalidateIdentity(): void {
Expand Down
19 changes: 15 additions & 4 deletions packages/agent/src/agent/proxy.ts
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,8 @@ export class ProxyStubAgent {
export class ProxyAgent implements Agent {
private _nextId = 0;
private _pendingCalls = new Map<number, [(resolve: any) => void, (reject: any) => void]>();
public rootKey = null;

#fetchRootKeyPromise: Promise<ArrayBuffer> | null = null;

constructor(private _backend: (msg: ProxyMessage) => void) {}

Expand Down Expand Up @@ -241,10 +242,20 @@ export class ProxyAgent implements Agent {
});
}

async getRootKey(): Promise<ArrayBuffer> {
return this.fetchRootKey();
}

public async fetchRootKey(): Promise<ArrayBuffer> {
if (this.#fetchRootKeyPromise === null) {
this.#fetchRootKeyPromise = this.#fetchRootKey();
}

return await this.#fetchRootKeyPromise;
}

async #fetchRootKey(): Promise<ArrayBuffer> {
// Hex-encoded version of the replica root key
const rootKey = ((await this.status()) as any).root_key;
this.rootKey = rootKey;
return rootKey;
return ((await this.status()) as any).root_key;
}
}
8 changes: 4 additions & 4 deletions packages/agent/src/canisterStatus/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -147,13 +147,13 @@ export const request = async (options: {
});
const cert = await Certificate.create({
certificate: response.certificate,
rootKey: agent.rootKey,
rootKey: await agent.getRootKey(),
canisterId: canisterId,
});

const lookup = (cert: Certificate, path: Path) => {
const lookup = async (cert: Certificate, path: Path) => {
if (path === 'subnet') {
const data = fetchNodeKeys(response.certificate, canisterId, agent.rootKey);
const data = fetchNodeKeys(response.certificate, canisterId, await agent.getRootKey());
return {
path: path,
data,
Expand All @@ -167,7 +167,7 @@ export const request = async (options: {
};

// must pass in the rootKey if we have no delegation
const { path, data } = lookup(cert, uniquePaths[index]);
const { path, data } = await lookup(cert, uniquePaths[index]);
if (!data) {
// Typically, the cert lookup will throw
console.warn(`Expected to find result for path ${path}, but instead found nothing.`);
Expand Down
3 changes: 1 addition & 2 deletions packages/agent/src/polling/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -92,10 +92,9 @@ export async function pollForResponse(
}

const state = await agent.readState(canisterId, { paths: [path] }, undefined, currentRequest);
if (agent.rootKey == null) throw new Error('Agent root key not initialized before polling');
const cert = await Certificate.create({
certificate: state.certificate,
rootKey: agent.rootKey,
rootKey: await agent.getRootKey(),
canisterId: canisterId,
blsVerify,
});
Expand Down

0 comments on commit 68839a6

Please sign in to comment.