Skip to content

Commit

Permalink
🚧 Privy signature for smart wallet
Browse files Browse the repository at this point in the history
  • Loading branch information
@KONFeature
KONFeature committed Dec 18, 2024
1 parent f8662a3 commit f6a0d23
Show file tree
Hide file tree
Showing 7 changed files with 153 additions and 98 deletions.
25 changes: 15 additions & 10 deletions packages/backend-elysia/src/domain/auth/routes/wallet.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,11 +59,12 @@ export const walletAuthRoutes = new Elysia({ prefix: "/wallet" })
"/privyLogin",
async ({
// Request
body: { expectedChallenge, signature, wallet, walletId, ssoId },
body: { expectedChallenge, signature, wallet, ssoId },
// Response
error,
// Context
client,
getEcdsaWalletAddress,
walletJwt,
generateSdkJwt,
resolveSsoSession,
Expand All @@ -81,33 +82,38 @@ export const walletAuthRoutes = new Elysia({ prefix: "/wallet" })
return error(404, "Invalid signature");
}

const authenticatorId = `privy-${walletId}` as const;
const authenticatorId = `privy-${wallet}` as const;

// Get the wallet address
const walletAddress = await getEcdsaWalletAddress({
ecdsaAddress: wallet,
});

// Create the token and set the cookie
const token = await walletJwt.sign({
address: wallet,
address: walletAddress,
authenticatorId,
publicKey: wallet,
sub: wallet,
sub: walletAddress,
iat: Date.now(),
transports: undefined,
});

// Finally, generate a JWT token for the SDK
const sdkJwt = await generateSdkJwt({ wallet });
const sdkJwt = await generateSdkJwt({ wallet: walletAddress });

// If all good, mark the sso as done
if (ssoId) {
await resolveSsoSession({
id: ssoId,
wallet,
wallet: walletAddress,
authenticatorId,
});
}

return {
token,
address: wallet,
address: walletAddress,
authenticatorId,
publicKey: wallet,
sdkJwt,
Expand All @@ -118,7 +124,6 @@ export const walletAuthRoutes = new Elysia({ prefix: "/wallet" })
body: t.Object({
expectedChallenge: t.String(),
wallet: t.Address(),
walletId: t.String(),
signature: t.Hex(),
// potential sso id
ssoId: t.Optional(t.Hex()),
Expand Down Expand Up @@ -224,7 +229,7 @@ export const walletAuthRoutes = new Elysia({ prefix: "/wallet" })
generateSdkJwt,
authenticatorRepository,
walletJwt,
getAuthenticatorWalletAddress,
getWebAuthnWalletAddress,
parseCompressedWebAuthNResponse,
resolveSsoSession,
}) => {
Expand Down Expand Up @@ -268,7 +273,7 @@ export const walletAuthRoutes = new Elysia({ prefix: "/wallet" })
// Get the wallet address
const walletAddress =
previousWallet ??
(await getAuthenticatorWalletAddress({
(await getWebAuthnWalletAddress({
authenticatorId: credential.id,
pubKey: publicKey,
}));
Expand Down
27 changes: 23 additions & 4 deletions packages/backend-elysia/src/domain/auth/services/WebAuthNService.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ import {
import {WebAuthN, kernelAddresses, KernelWallet} from "@frak-labs/app-essentials";
import { Elysia } from "elysia";
import { getSenderAddress } from "permissionless/actions";
import { type Hex, concatHex, keccak256, toHex } from "viem";
import { type Address, type Hex, concatHex, keccak256, toHex } from "viem";
import { entryPoint06Address } from "viem/account-abstraction";
import { AuthenticatorRepository } from "../repositories/AuthenticatorRepository";

Expand All @@ -34,7 +34,7 @@ export const webAuthNService = new Elysia({ name: "Service.webAuthN" })
/**
* Get a wallet address from an authenticator
*/
async function getAuthenticatorWalletAddress({
async function getWebAuthnWalletAddress({
authenticatorId,
pubKey,
}: { authenticatorId: string; pubKey: { x: Hex; y: Hex } }) {
Expand All @@ -52,6 +52,24 @@ export const webAuthNService = new Elysia({ name: "Service.webAuthN" })
});
}

/**
* Get a wallet address from an authenticator
*/
async function getEcdsaWalletAddress({
ecdsaAddress,
}: { ecdsaAddress: Address }) {
// Compute base stuff to fetch the smart wallet address
const initCode = KernelWallet.getFallbackWalletInitCode({
ecdsaAddress,
});

// Get the sender address based on the init code
return getSenderAddress(client, {
initCode: concatHex([kernelAddresses.factory, initCode]),
entryPointAddress: entryPoint06Address,
});
}

/**
* Check if a signature is valid for a given wallet
*/
Expand All @@ -76,7 +94,7 @@ export const webAuthNService = new Elysia({ name: "Service.webAuthN" })
}

// Check if the address match the signature provided
const walletAddress = await getAuthenticatorWalletAddress({
const walletAddress = await getWebAuthnWalletAddress({
authenticatorId: signature.id,
pubKey: authenticator.publicKey,
});
Expand Down Expand Up @@ -122,7 +140,8 @@ export const webAuthNService = new Elysia({ name: "Service.webAuthN" })
isValidWebAuthNSignature,
parseCompressedWebAuthNResponse,
authenticatorRepository,
getAuthenticatorWalletAddress,
getWebAuthnWalletAddress,
getEcdsaWalletAddress,
};
})
.as("plugin");
35 changes: 28 additions & 7 deletions packages/wallet/app/context/wallet/smartWallet/connector.ts
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
import { currentChain } from "@/context/blockchain/provider";
import { getSmartAccountProvider } from "@/context/wallet/smartWallet/provider";
import type { SmartAccountV06 } from "@/context/wallet/smartWallet/utils";
import type { ConnectedWallet } from "@privy-io/react-auth";
import type { Transport } from "viem";
import type { PrivyInterface } from "@privy-io/react-auth";
import type { Hex, Transport } from "viem";
import { createConnector } from "wagmi";

smartAccountConnector.type = "frakSmartAccountConnector" as const;
Expand All @@ -26,10 +26,17 @@ export function smartAccountConnector<
// The current provider
let provider: Provider | undefined;

// The privy message signer (null if not ready)
let signViaPrivy: PrivyInterface["signMessage"] | undefined = undefined;

// Create the wagmi connector itself
return createConnector<
Provider,
{ onPrivyWalletsUpdate: (args: { wallets: ConnectedWallet[] }) => void }
{
onPrivyInterfaceUpdate: (
args: PrivyInterface["signMessage"]
) => void;
}
>((config) => ({
id: "frak-wallet-connector",
name: "Frak Smart Account",
Expand Down Expand Up @@ -146,6 +153,21 @@ export function smartAccountConnector<
chainId: config.chains[0].id,
});
},

async signViaPrivy(message, address) {
if (!signViaPrivy) {
throw new Error("Privy not ready yet");
}
return (await signViaPrivy(
message,
{
title: "Action confirmation",
description:
"By signing the following hash, you will authorize the current frak action",
},
address
)) as Hex;
},
});
}
return provider;
Expand All @@ -159,10 +181,9 @@ export function smartAccountConnector<
onDisconnect() {
config.emitter.emit("disconnect");
},
// When the list of privy wallets change
onPrivyWalletsUpdate({ wallets }: { wallets: ConnectedWallet[] }) {
// todo: Do some stuff with it
console.log("Wagmi provider wallets update", { wallets });

onPrivyInterfaceUpdate(privySignMsg: PrivyInterface["signMessage"]) {
signViaPrivy = privySignMsg;
},
}));
}
114 changes: 63 additions & 51 deletions packages/wallet/app/context/wallet/smartWallet/provider.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import {
} from "@/context/blockchain/aa-provider";
import { currentChain, currentViemClient } from "@/context/blockchain/provider";
import { getSignOptions } from "@/context/wallet/action/signOptions";
import { frakFallbackWalletSmartAccount } from "@/context/wallet/smartWallet/FrakFallbackSmartWallet";
import { frakWalletSmartAccount } from "@/context/wallet/smartWallet/FrakSmartWallet";
import type { SmartAccountV06 } from "@/context/wallet/smartWallet/utils";
import { parseWebAuthNAuthentication } from "@/context/wallet/smartWallet/webAuthN";
Expand All @@ -19,17 +20,24 @@ import {
createSmartAccountClient,
} from "permissionless";
import { getUserOperationGasPrice } from "permissionless/actions/pimlico";
import type { Transport } from "viem";
import type { Address, Hex, Transport } from "viem";
import type { SmartAccount } from "viem/account-abstraction";

/**
* Properties
*/
type SmartAccountProvierParameters = {
type SmartAccountProviderParameters = {
/**
* Method when the account has changed
*/
onAccountChanged: (newWallet?: WebAuthNWallet | PrivyWallet) => void;

/**
* Method used to sign aa message via privy
* @param data
* @param address
*/
signViaPrivy: (data: Hex, address: Address) => Promise<Hex>;
};

/**
Expand All @@ -42,7 +50,7 @@ type SmartAccountProvierParameters = {
export function getSmartAccountProvider<
transport extends Transport = Transport,
account extends SmartAccountV06 = SmartAccountV06,
>({ onAccountChanged }: SmartAccountProvierParameters) {
>({ onAccountChanged, signViaPrivy }: SmartAccountProviderParameters) {
console.log("Building a new smart account provider");
// A few types shortcut
type ConnectorClient = SmartAccountClient<
Expand Down Expand Up @@ -103,24 +111,11 @@ export function getSmartAccountProvider<
return undefined;
}

const privyWallet =
currentWebAuthNWallet.authenticatorId.startsWith("privy-")
? (currentWebAuthNWallet as PrivyWallet)
: undefined;
const webauthnWallet = privyWallet
? undefined
: (currentWebAuthNWallet as WebAuthNWallet);

// Otherwise, build it
if (privyWallet) {
// todo: Custom for privy here
return undefined;
}
targetSmartAccount = webauthnWallet
? await buildSmartAccount({
wallet: webauthnWallet,
})
: undefined;
targetSmartAccount = await buildSmartAccount({
wallet: currentWebAuthNWallet,
signViaPrivy,
});

// Save the new one
currentSmartAccountClient = targetSmartAccount;
Expand Down Expand Up @@ -149,39 +144,56 @@ async function buildSmartAccount<
account extends SmartAccountV06 = SmartAccountV06,
>({
wallet,
}: { wallet: WebAuthNWallet }): Promise<
signViaPrivy,
}: {
wallet: WebAuthNWallet | PrivyWallet;
signViaPrivy: (data: Hex, address: Address) => Promise<Hex>;
}): Promise<
SmartAccountClient<transport, typeof currentChain, SmartAccount<account>>
> {
// Get the smart wallet client
const smartAccount = await frakWalletSmartAccount(currentViemClient, {
authenticatorId: wallet.authenticatorId,
signerPubKey: wallet.publicKey,
signatureProvider: async (message) => {
// Get the signature options from server
const options = await getSignOptions({
authenticatorId: wallet.authenticatorId,
toSign: message,
});

// Start the client authentication
const authenticationResponse = await startAuthentication({
optionsJSON: options,
});

// Store that in our last webauthn action atom
jotaiStore.set(lastWebAuthNActionAtom, {
wallet: wallet.address,
signature: authenticationResponse,
msg: options.challenge,
});

// Store this shit somewhere

// Perform the verification of the signature
return parseWebAuthNAuthentication(authenticationResponse);
},
preDeterminedAccountAddress: wallet.address,
});
let smartAccount: SmartAccountV06;
// Get the webauthn smart wallet client
if (typeof wallet.publicKey === "object") {
// That's a webauthn wallet
smartAccount = await frakWalletSmartAccount(currentViemClient, {
authenticatorId: wallet.authenticatorId,
signerPubKey: wallet.publicKey,
signatureProvider: async (message) => {
// Get the signature options from server
const options = await getSignOptions({
authenticatorId: wallet.authenticatorId,
toSign: message,
});

// Start the client authentication
const authenticationResponse = await startAuthentication({
optionsJSON: options,
});

// Store that in our last webauthn action atom
jotaiStore.set(lastWebAuthNActionAtom, {
wallet: wallet.address,
signature: authenticationResponse,
msg: options.challenge,
});

// Store this shit somewhere

// Perform the verification of the signature
return parseWebAuthNAuthentication(authenticationResponse);
},
preDeterminedAccountAddress: wallet.address,
});
} else {
// That's a privy wallet
smartAccount = await frakFallbackWalletSmartAccount(currentViemClient, {
ecdsaAddress: wallet.publicKey,
preDeterminedAccountAddress: wallet.address,
signatureProvider({ hash }) {
return signViaPrivy(hash, wallet.publicKey);
},
});
}

// Get the bundler and paymaster clients
const pimlicoTransport = getPimlicoTransport();
Expand Down
Loading

0 comments on commit f6a0d23

Please sign in to comment.