Skip to content

Commit

Permalink
Merge pull request #58 from dwhiffing/encryption-context
Browse files Browse the repository at this point in the history
Encryption context
  • Loading branch information
dwhiffing authored Sep 13, 2024
2 parents cd57412 + 39a02e5 commit e6071d7
Show file tree
Hide file tree
Showing 13 changed files with 293 additions and 20 deletions.
13 changes: 9 additions & 4 deletions app/api/execute/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ import { NextRequest, NextResponse } from "next/server";
import { getToolsFromContracts } from "@/utils/generateToolFromABI";
import { routeBodySchema } from "./schemas";
import { contractCollection } from "@/utils/collections";
import { hashPin } from "@/utils/crypt";

export const runtime = "nodejs";

Expand All @@ -21,7 +22,9 @@ export async function POST(req: NextRequest) {
return NextResponse.json({ error: errorMessages }, { status: 400 });
}

const { toolCall, didToken } = result.data;
const { toolCall, didToken, pin } = result.data;

const encryptionContext = await hashPin(pin);

// parse contractAddress from toolCall.name; Should be in format `${contractKey}_${functionName}_${overload function index}``
const contractKey = parseInt(toolCall.name.split("_").at(0) as string, 10);
Expand All @@ -40,9 +43,11 @@ export async function POST(req: NextRequest) {
}

try {
const tool = getToolsFromContracts([contract], didToken).find(
(t) => t.name === toolCall.name,
);
const tool = getToolsFromContracts(
[contract],
didToken,
encryptionContext,
).find((t) => t.name === toolCall.name);

if (!tool) {
return NextResponse.json(
Expand Down
6 changes: 6 additions & 0 deletions app/api/execute/schemas.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,4 +17,10 @@ export const routeBodySchema = z.object({
invalid_type_error: "didToken must be a string",
})
.min(1, "didToken cannot be an empty string"),
pin: z
.string({
required_error: "pin is required",
invalid_type_error: "pin must be a string",
})
.min(1, "pin cannot be an empty string"),
});
21 changes: 17 additions & 4 deletions app/api/wallet/route.ts
Original file line number Diff line number Diff line change
@@ -1,21 +1,34 @@
import { getWalletUUIDandAccessKey } from "@/utils/tee";
import { Magic } from "@magic-sdk/admin";
import { NextRequest, NextResponse } from "next/server";
import { hashPin } from "@/utils/crypt";

export const runtime = "nodejs";
const magic = await Magic.init(process.env.MAGIC_SECRET_KEY);

export async function GET(req: NextRequest) {
try {
console.log(req.nextUrl.searchParams.get("didToken"));
const didToken = req.nextUrl.searchParams.get("didToken") ?? "";
const pin = req.nextUrl.searchParams.get("pin");

if (!didToken) throw new Error("TOKEN missing");

const userMetadata = await magic.users.getMetadataByToken(didToken);
const publicAddress = userMetadata.publicAddress ?? "";
const encryptionContext = pin ? await hashPin(pin) : undefined;

const result = await getWalletUUIDandAccessKey(publicAddress);

return NextResponse.json({ wallet_address: result.wallet_address });
try {
const result = await getWalletUUIDandAccessKey(
publicAddress,
encryptionContext,
);
return NextResponse.json({ wallet_address: result.wallet_address });
} catch (e) {
return NextResponse.json(
{ error: "Wallet does not exist and no PIN was provided" },
{ status: 400 },
);
}
} catch (e: any) {
return NextResponse.json({ error: e.message }, { status: e.status ?? 500 });
}
Expand Down
7 changes: 7 additions & 0 deletions components/ChatMessageBubble.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import { ToolArgsTable } from "./ToolArgsTable";
import { useContracts } from "@/utils/useContracts";
import { CHAINS } from "@/constants";
import { IContract } from "@/types";
import { usePinInput } from "./PinInput";

type IToolCall = {
name: string;
Expand Down Expand Up @@ -117,6 +118,7 @@ export function ToolCallMessageBubble(props: { message: Message }) {
const { colorClassName, alignmentClassName, icon } = getStyleForRole(
props.message.role,
);
const { getPin, pinInput } = usePinInput();

let content: { text: string; toolCall?: IToolCall } = {
text: "",
Expand All @@ -133,11 +135,14 @@ export function ToolCallMessageBubble(props: { message: Message }) {

setLoading(true);
try {
const pin = await getPin();
if (!pin) throw new Error("Invalid PIN");
const resp = await fetch("/api/execute", {
method: "POST",
body: JSON.stringify({
toolCall,
didToken,
pin,
disabledContractKeys: disabledKeys,
}),
});
Expand Down Expand Up @@ -227,6 +232,8 @@ export function ToolCallMessageBubble(props: { message: Message }) {
{renderContent}
</div>
</div>

{pinInput}
</>
);
}
Expand Down
19 changes: 17 additions & 2 deletions components/MagicProvider.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import { Magic } from "magic-sdk";
import { Web3 } from "web3";
import { createContext, useContext, useEffect, useMemo, useState } from "react";
import { usePinInput } from "./PinInput";

// Create and export the context
export const MagicContext = createContext<{
Expand Down Expand Up @@ -39,6 +40,13 @@ const MagicProvider = ({ children }: any) => {
const [address, setAddress] = useState<string | null>(null);
const [didToken, setDidToken] = useState<string | null>(null);

const { getPin, pinInput } = usePinInput({
title: "Enter your TEE Wallet PIN",
description:
"You will be asked to enter this value whenever you try to execute a transaction",
allowCancel: false,
});

useEffect(() => {
if (process.env.NEXT_PUBLIC_MAGIC_API_KEY) {
const magic = new Magic(process.env.NEXT_PUBLIC_MAGIC_API_KEY || "", {
Expand Down Expand Up @@ -69,14 +77,19 @@ const MagicProvider = ({ children }: any) => {
let didToken = await magic.user.getIdToken();
setDidToken(didToken);

const response = await fetch(`/api/wallet?didToken=${didToken}`);
let response = await fetch(`/api/wallet?didToken=${didToken}`);
if (!response.ok) {
const pin = await getPin();
response = await fetch(`/api/wallet?didToken=${didToken}&pin=${pin}`);
}

const json = await response.json();
setTEEWalletAddress(json.wallet_address);
}
setIsLoading(false);
};
checkIfLoggedIn();
}, [magic]);
}, [magic, getPin]);

const handleLogin = async () => {
if (!magic) return;
Expand All @@ -101,6 +114,7 @@ const MagicProvider = ({ children }: any) => {
setIsLoggedIn(false);
setDidToken(null);
setAddress(null);
setTEEWalletAddress(null);
};

const value = useMemo(() => {
Expand All @@ -124,6 +138,7 @@ const MagicProvider = ({ children }: any) => {
}}
>
{children}
{pinInput}
</MagicContext.Provider>
);
};
Expand Down
109 changes: 109 additions & 0 deletions components/PinInput.tsx
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
import {
InputOTP,
InputOTPGroup,
InputOTPSlot,
} from "@/components/ui/input-otp";
import {
AlertDialog,
AlertDialogAction,
AlertDialogContent,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { useCallback, useEffect, useRef, useState } from "react";
import {
AlertDialogCancel,
AlertDialogDescription,
} from "@radix-ui/react-alert-dialog";

export function PinInput(props: {
open: boolean;
title: string;
description: string;
onSubmit?: (s: string) => void;
onCancel?: () => void;
pinLength?: number;
}) {
const [value, setValue] = useState("");
const length = props.pinLength ?? 4;

useEffect(() => {
if (props.open) setValue("");
}, [props.open]);

return (
<AlertDialog open={props.open}>
<AlertDialogContent className={props.description ? "w-[24rem]" : "w-60"}>
<form
className="flex flex-col gap-6"
onSubmit={(e) => {
e.preventDefault();
if (value.length === length) props.onSubmit?.(value);
}}
>
<AlertDialogHeader className="items-center gap-2">
<AlertDialogTitle>{props.title}</AlertDialogTitle>
{props.description && (
<AlertDialogDescription className="!mt-0 text-center opacity-60">
{props.description}
</AlertDialogDescription>
)}
<InputOTP
autoFocus={props.open}
maxLength={length}
value={value}
onChange={setValue}
>
<InputOTPGroup>
{Array.from({ length }, (_, i) => (
<InputOTPSlot key={i} index={i} />
))}
</InputOTPGroup>
</InputOTP>
</AlertDialogHeader>

<AlertDialogFooter className="gap-4">
{props.onCancel && (
<AlertDialogCancel onClick={() => props.onCancel?.()}>
Cancel
</AlertDialogCancel>
)}
<AlertDialogAction disabled={value.length < length} type="submit">
Submit
</AlertDialogAction>
</AlertDialogFooter>
</form>
</AlertDialogContent>
</AlertDialog>
);
}

export const usePinInput = ({
title = "Enter PIN",
description = "",
allowCancel = true,
} = {}) => {
const [isPinOpen, setIsPinOpen] = useState(false);
const pinPromiseRef = useRef<(s?: string) => void>();
const pinInput = (
<PinInput
open={isPinOpen}
title={title}
description={description}
onCancel={allowCancel ? pinPromiseRef.current : undefined}
onSubmit={pinPromiseRef.current}
/>
);

const getPin = useCallback(async () => {
setIsPinOpen(true);
const pin = await new Promise((resolve) => {
pinPromiseRef.current = resolve;
});
setIsPinOpen(false);
return pin;
}, []);

return { pinInput, getPin };
};
71 changes: 71 additions & 0 deletions components/ui/input-otp.tsx
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
"use client"

import * as React from "react"
import { OTPInput, OTPInputContext } from "input-otp"
import { Dot } from "lucide-react"

import { cn } from "@/lib/utils"

const InputOTP = React.forwardRef<
React.ElementRef<typeof OTPInput>,
React.ComponentPropsWithoutRef<typeof OTPInput>
>(({ className, containerClassName, ...props }, ref) => (
<OTPInput
ref={ref}
containerClassName={cn(
"flex items-center gap-2 has-[:disabled]:opacity-50",
containerClassName
)}
className={cn("disabled:cursor-not-allowed", className)}
{...props}
/>
))
InputOTP.displayName = "InputOTP"

const InputOTPGroup = React.forwardRef<
React.ElementRef<"div">,
React.ComponentPropsWithoutRef<"div">
>(({ className, ...props }, ref) => (
<div ref={ref} className={cn("flex items-center", className)} {...props} />
))
InputOTPGroup.displayName = "InputOTPGroup"

const InputOTPSlot = React.forwardRef<
React.ElementRef<"div">,
React.ComponentPropsWithoutRef<"div"> & { index: number }
>(({ index, className, ...props }, ref) => {
const inputOTPContext = React.useContext(OTPInputContext)
const { char, hasFakeCaret, isActive } = inputOTPContext.slots[index]

return (
<div
ref={ref}
className={cn(
"relative flex h-10 w-10 items-center justify-center border-y border-r border-input text-sm transition-all first:rounded-l-md first:border-l last:rounded-r-md",
isActive && "z-10 ring-2 ring-ring ring-offset-background",
className
)}
{...props}
>
{char}
{hasFakeCaret && (
<div className="pointer-events-none absolute inset-0 flex items-center justify-center">
<div className="h-4 w-px animate-caret-blink bg-foreground duration-1000" />
</div>
)}
</div>
)
})
InputOTPSlot.displayName = "InputOTPSlot"

const InputOTPSeparator = React.forwardRef<
React.ElementRef<"div">,
React.ComponentPropsWithoutRef<"div">
>(({ ...props }, ref) => (
<div ref={ref} role="separator" {...props}>
<Dot />
</div>
))
InputOTPSeparator.displayName = "InputOTPSeparator"

export { InputOTP, InputOTPGroup, InputOTPSlot, InputOTPSeparator }
10 changes: 10 additions & 0 deletions package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
"eslint-config-next": "13.4.12",
"ethers": "^6.13.2",
"etherscan-api": "^10.3.0",
"input-otp": "^1.2.4",
"langchain": "^0.2.12",
"lucide-react": "^0.428.0",
"magic-sdk": "^28.5.0",
Expand Down
Loading

0 comments on commit e6071d7

Please sign in to comment.