Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Encrypt and Decrypt support #39

Merged
merged 4 commits into from
Apr 3, 2023
Merged

Encrypt and Decrypt support #39

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a7dc2eb
Initial commit
neekolas Apr 3, 2023
fceab4a
Fix failing test
neekolas Apr 3, 2023
4781a67
Switch to a RefCell
neekolas Apr 3, 2023
efa7486
Use regular borrow when possible
neekolas Apr 3, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 54 additions & 4 deletions bindings/wasm/crate/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ pub fn create_outbound_session(
sending_handle: &str,
receiving_handle: &str,
message: &str,
) -> Result<String, JsValue> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure I follow the need for the Box here.

My understanding was that JsValues through WASM_BINDGEN get copied between the memory space?

Heap Allocating here seems unnecessary at a high level.

) -> Result<Box<[JsValue]>, JsValue> {
// Look up both handles in INSTANCE_MAP
let instances = INSTANCE_MAP.lock().unwrap();
let mut sending_instance = instances
Expand All @@ -53,8 +53,13 @@ pub fn create_outbound_session(
.get(receiving_handle)
.unwrap()
.set(receiving_instance);

match result {
Ok((_, ciphertext_json)) => Ok(ciphertext_json),
Ok((session_id, ciphertext_json)) => Ok(vec![
JsValue::from_str(&session_id),
JsValue::from_str(&ciphertext_json),
]
Comment on lines +52 to +55
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using the Vec here changes the memory structure. Given the size cannot be known at compile time.

Imagine using a tuple would be easier to use and would avoid the need for boxing of the values.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unfortunately, wasm-bindgen doesn't appear to support tuples as return values. rustwasm/wasm-bindgen#122

.into_boxed_slice()),
Err(e) => Err(JsValue::from_str(&e.to_string())),
}
}
Expand All @@ -64,7 +69,7 @@ pub fn create_inbound_session(
sending_handle: &str,
receiving_handle: &str,
message: &str,
) -> Result<String, JsValue> {
) -> Result<Box<[JsValue]>, JsValue> {
// Look up both handles in INSTANCE_MAP
let instances = INSTANCE_MAP.lock().unwrap();
let sending_instance = instances
Expand All @@ -90,7 +95,52 @@ pub fn create_inbound_session(
.set(receiving_instance);

match result {
Ok((_, plaintext)) => Ok(plaintext),
Ok((session_id, ciphertext_json)) => Ok(vec![
JsValue::from_str(&session_id),
JsValue::from_str(&ciphertext_json),
]
.into_boxed_slice()),
Err(e) => Err(JsValue::from_str(&e.to_string())),
}
}

#[wasm_bindgen]
pub fn encrypt_message(
sending_handle: &str,
session_id: &str,
message: &str,
) -> Result<String, JsValue> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These types look both to be Strings, only the error gets converted to a wasmbind:JsValue. Shouldn't they be the same type?

let instances = INSTANCE_MAP.lock().unwrap();
let mut instance = instances
.get(sending_handle)
.ok_or("sending_handle not found")?
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like this conversion from an Option to Result.

You could have used a match, but this is much cleaner

.take();

let result = instance.encrypt_message_serialized(session_id, message);
// Do we actually need to do this?
instances.get(sending_handle).unwrap().set(instance);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You don't have to do this here. Instead of using take() you should be able to use get_mut() on the cell (so .get(sending_handle).ok_or(...)?.get_mut() approximately.

This pattern comes from my own lack of Rust borrow-checker knowledge when implementing the inbound/outbound session functions. I couldn't take mutable references of two INSTANCE_MAP values in the same scope so ended up doing it via Cell in a weird way.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In fact, now that we're using .public_account, I think we can do away with this pattern in those existing functions too. I might tackle that in a separate PR.

Copy link
Contributor Author

@neekolas neekolas Apr 3, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The type checker doesn't seem to like that, but it recommends that I replace the cell with a RefCell. RefCell then has a method borrow_mut. If I use that, everything is happy and I don't have to manually set the cell to updated values.

Does have some scary language about concurrent usage leading to a panic, but I think that's fine for now given that all our code is synchronous.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've committed that change here

Copy link
Contributor

@michaelx11 michaelx11 Apr 3, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good, anything works for now as long as borrow/type checker likes it. We can take a holistic look at the map/cell strategy in a subsequent PR.

Actually, on second look: RefCell seems much better than the take/set strategy. Nice!

neekolas marked this conversation as resolved.
Show resolved Hide resolved

match result {
Ok(ciphertext_json) => Ok(ciphertext_json),
Err(e) => Err(JsValue::from_str(&e.to_string())),
}
Comment on lines +108 to +111
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems fine to me, however there are more idomatic ways to use a Match statement.

  1. Pretty sure that RHS Ok() is not needed. Values are implicitly Ok
  2. result is an 'Result' which you are trying to map the error. You can use result.map_error(|e| JsValue::from_str(&e.to_string()))

}

#[wasm_bindgen]
pub fn decrypt_message(
handle: &str,
session_id: &str,
ciphertext: &str,
) -> Result<String, JsValue> {
let instances = INSTANCE_MAP.lock().unwrap();
let mut instance = instances.get(handle).ok_or("handle not found")?.take();

let result = instance.decrypt_message_serialized(session_id, ciphertext);
// Do we actually need to do this?
neekolas marked this conversation as resolved.
Show resolved Hide resolved
instances.get(handle).unwrap().set(instance);

match result {
Ok(plaintext) => Ok(plaintext),
Err(e) => Err(JsValue::from_str(&e.to_string())),
}
}
Expand Down
83 changes: 70 additions & 13 deletions bindings/wasm/src/xmtpv3.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ import init, {
new_voodoo_instance,
create_outbound_session,
create_inbound_session,
decrypt_message,
encrypt_message,
e2e_selftest,
} from "./pkg/libxmtp.js";

Expand All @@ -20,6 +22,11 @@ export const setWasmInit = (arg: () => InitInput) => {

let initialized: Promise<void> | undefined = undefined;

type SessionResult = {
sessionId: string;
payload: string;
};

export class VoodooInstance {
// Handle to the voodooinstance object in the Wasm module
public handle: string = "";
Expand All @@ -31,16 +38,37 @@ export class VoodooInstance {
this.handle = handle;
}

createOutboundSession(other: VoodooInstance, msg: string): Promise<string> {
return new Promise((resolve, reject) => {
resolve(this.wasmModule.createOutboundSession(this.handle, other.handle, msg));
});
async createOutboundSession(
other: VoodooInstance,
msg: string
): Promise<SessionResult> {
const [sessionId, payload] = this.wasmModule.createOutboundSession(
this.handle,
other.handle,
msg
);

return { sessionId, payload };
}

async createInboundSession(
other: VoodooInstance,
msg: string
): Promise<SessionResult> {
const [sessionId, payload] = this.wasmModule.createInboundSession(
other.handle,
this.handle,
msg
);
return { sessionId, payload };
}

async encryptMessage(sessionId: string, msg: string): Promise<string> {
return this.wasmModule.encryptMessage(this.handle, sessionId, msg);
}

createInboundSession(other: VoodooInstance, msg: string): Promise<string> {
return new Promise((resolve, reject) => {
resolve(this.wasmModule.createInboundSession(other.handle, this.handle, msg));
});
async decryptMessage(sessionId: string, ciphertext: string): Promise<string> {
return this.wasmModule.decryptMessage(this.handle, sessionId, ciphertext);
}
}

Expand All @@ -54,7 +82,6 @@ export class XMTPv3 {
}
}


// Manages the Wasm module, which loads a singleton version of our Rust code
export class XMTPWasm {
private constructor() {}
Expand All @@ -69,12 +96,42 @@ export class XMTPWasm {
return new VoodooInstance(this, handle);
}

createOutboundSession(sendHandle: string, receiveHandle: string, msg: string): string {
return create_outbound_session(sendHandle, receiveHandle, msg);
createOutboundSession(
sendHandle: string,
receiveHandle: string,
msg: string
): [string, string] {
return create_outbound_session(sendHandle, receiveHandle, msg) as [
string,
string
];
}

createInboundSession(
sendHandle: string,
receiveHandle: string,
msg: string
): [string, string] {
return create_inbound_session(sendHandle, receiveHandle, msg) as [
string,
string
];
}

encryptMessage(
sendHandle: string,
sessionId: string,
message: string
): string {
return encrypt_message(sendHandle, sessionId, message);
}

createInboundSession(sendHandle: string, receiveHandle: string, msg: string): string {
return create_inbound_session(sendHandle, receiveHandle, msg);
decryptMessage(
handle: string,
sessionId: string,
ciphertext: string
): string {
return decrypt_message(handle, sessionId, ciphertext);
}

/**
Expand Down
45 changes: 37 additions & 8 deletions bindings/wasm/tests/wasmpkg.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,19 +12,48 @@ it("can instantiate", async () => {

it("can run self test", async () => {
const xmtp = await XMTPWasm.initialize();
const xmtpv3 = await xmtp.getXMTPv3();
const res = await xmtpv3.selfTest();
const xmtpv3 = xmtp.getXMTPv3();
const res = xmtpv3.selfTest();
expect(res).toBe(true);
});

it("can do a simple conversation", async () => {
const wasm = await XMTPWasm.initialize();
const alice = await wasm.newVoodooInstance();
const bob = await wasm.newVoodooInstance();
const alice = wasm.newVoodooInstance();
const bob = wasm.newVoodooInstance();

const outboundJson = await alice.createOutboundSession(bob, "hello there");
const { sessionId, payload } = await alice.createOutboundSession(
bob,
"hello there"
);
expect(typeof sessionId).toBe("string");
// Unused, but test JSON parseable
const _ = JSON.parse(outboundJson);
const inboundPlaintext = await bob.createInboundSession(alice, outboundJson);
expect(inboundPlaintext).toBe("hello there");
const _ = JSON.parse(payload);
const { payload: inboundPayload } = await bob.createInboundSession(
alice,
payload
);
expect(inboundPayload).toBe("hello there");
});

it("can send a message", async () => {
const wasm = await XMTPWasm.initialize();
const alice = wasm.newVoodooInstance();
const bob = wasm.newVoodooInstance();

const { sessionId, payload } = await alice.createOutboundSession(
bob,
"hello there"
);
await bob.createInboundSession(alice, payload);
expect(typeof sessionId).toBe("string");

const msg = "hello there";
const encrypted = await alice.encryptMessage(sessionId, msg);
expect(typeof encrypted).toBe("string");

// Alice can't decrypt her own message. Does work for Bob though
// const decrypted = await alice.decryptMessage(sessionId, encrypted);
const decrypted = await bob.decryptMessage(sessionId, encrypted);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

beautiful

expect(decrypted).toBe(msg);
});