Skip to content

Commit

Permalink
Handle concurrent auth requests by same user
Browse files Browse the repository at this point in the history
  • Loading branch information
@cephalization
cephalization committed Oct 29, 2023
1 parent e2e8827 commit 7b6b604
Show file tree
Hide file tree
Showing 4 changed files with 73 additions and 56 deletions.
4 changes: 4 additions & 0 deletions apps/messages/src/agent.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@ const connectionRequestBodySchema = valibot.object({
topicId: valibot.number(),
/** The project that the room represents */
projectId: valibot.number(),
/** Key that must be sent with first party connection */
key: valibot.string(),
/** Temporary token that must be sent with first party connection */
token: valibot.string(),
/** temporary */
Expand Down Expand Up @@ -81,11 +83,13 @@ export default class Agent implements Party.Server {
switch (result.output.action) {
case 'connect': {
const tempToken = result.output.token;
const tempKey = result.output.key;
const socket = new PartySocket({
host: result.output.host,
room: result.output.id,
id: this.id,
query: () => ({
key: tempKey,
token: tempToken
})
});
Expand Down
28 changes: 9 additions & 19 deletions apps/messages/src/server.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ export default class Server implements Party.Server {
try {
invariant(req.method.toLocaleLowerCase() === 'post', 'must be a post request');
const body = await req.json();
const { userId } = parse(object({ userId: string() }), body);
const { userId, roomId } = parse(object({ userId: string(), roomId: string() }), body);

// validate user based on bearer token
const valid = await this.auth.validateSession(
Expand All @@ -91,8 +91,9 @@ export default class Server implements Party.Server {
}

// create token
const key = `${userId}-${roomId}`;
const result = await this.context.Queries.TemporaryTokens.createToken(
userId,
key,
generateRandomString(16)
);

Expand All @@ -101,7 +102,7 @@ export default class Server implements Party.Server {
}

// return token
return new Response(JSON.stringify({ token: result.value }), {
return new Response(JSON.stringify({ token: result.value, key }), {
status: 200
});
} catch {
Expand All @@ -114,28 +115,16 @@ export default class Server implements Party.Server {
static async onBeforeConnect(req: Party.Request, lobby: Party.Lobby) {
// if key in url === key in storage, assume agent
// and let them in
const keyInUrl = new URL(req.url).searchParams.get('token');
const tokenInUrl = new URL(req.url).searchParams.get('token');
const keyInUrl = new URL(req.url).searchParams.get('key');

if (!keyInUrl) {
if (!tokenInUrl || !keyInUrl) {
return new Response('Unauthorized', { status: 401 });
}

const { Queries } = Server.getNewDbClient(lobby.env.DATABASE_URL as string);
const agentKey = `agent-${lobby.id}`;
const valid = await Queries.TemporaryTokens.validateAndConsumeToken(agentKey, keyInUrl);

if (valid) {
console.log('AGENT VALID');
return req;
}

// otherwise, assume user and validate key as bearer token
const userInUrl = new URL(req.url).searchParams.get('user_id');
if (!userInUrl) {
return new Response('Unauthorized', { status: 401 });
}

const userValid = await Queries.TemporaryTokens.validateAndConsumeToken(userInUrl, keyInUrl);
const userValid = await Queries.TemporaryTokens.validateAndConsumeToken(keyInUrl, tokenInUrl);

if (!userValid) {
return new Response('Unauthorized', { status: 401 });
Expand Down Expand Up @@ -167,6 +156,7 @@ export default class Server implements Party.Server {
body: JSON.stringify({
action: 'connect',
id: this.party.id,
key: agentKey,
topicId: this.context.topicId,
projectId: this.context.projectId,
host: this.party.env.PARTY_HOST,
Expand Down
90 changes: 56 additions & 34 deletions apps/web/src/lib/stores/messages.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,49 +38,67 @@ export function createMessagesStore({
topicId,
callbacks
}: CreateMessagesStoreArgs) {
let socket: PartySocket | undefined;
let client: ReturnType<typeof createPartyClient<SafePartyEvents, SafePartyResponses>>;
const state: {
socket: PartySocket | undefined;
client: ReturnType<typeof createPartyClient<SafePartyEvents, SafePartyResponses>> | undefined;
abort?: AbortController;
} = {
socket: undefined,
client: undefined
};

const { subscribe } = readable({ ...initialState }, (set, update) => {
if (!partyOptions || !browser) return;

new Promise(() => {
const go = async () => {
try {
const authResponse = await fetch(`/api/party-auth`, {
method: 'POST',
body: JSON.stringify({
projectId,
topicId
})
});
try {
const abort = new AbortController();
state.abort = abort; // set the abort controller on the state object
// ensure that the cleanup function always has access to the abort controller
state.abort?.signal.addEventListener('abort', () => {
state.socket?.close();
state.client?.unsubscribe();
});
fetch(`/api/party-auth`, {
method: 'POST',
body: JSON.stringify({
projectId,
topicId
}),
signal: abort.signal
})
.then(async (authResponse) => {
if (!authResponse.ok) {
throw new Error('Failed to authenticate');
}

const body = await authResponse.json();
const response = parse(
object({
key: optional(string()),
token: optional(string()),
error: optional(string())
}),
body
);

console.log('response', response);

if (response.error) {
throw new Error(response.error);
}

if (!response.token) {
if (!response.token || !response.key) {
throw new Error('No token');
}

socket = new PartySocket({
state.socket = new PartySocket({
...partyOptions,
query: { ...partyOptions.query, token: response.token, user_id: partyOptions.id }
query: { ...partyOptions.query, token: response.token, key: response.key }
});
client = createPartyClient<SafePartyEvents, SafePartyResponses>(socket);
invariant(socket, 'socket should be defined');
invariant(client, 'client should be defined');
state.client = createPartyClient<SafePartyEvents, SafePartyResponses>(state.socket);
invariant(state.socket, 'socket should be defined');
invariant(state.client, 'client should be defined');

client.on('Init', (e) => {
state.client.on('Init', (e) => {
const newState = {
messages: e.messages,
activeUserIds: new Set(e.userIds)
Expand All @@ -89,7 +107,7 @@ export function createMessagesStore({
callbacks?.Init?.();
});

client.on('SetMessages', (e) => {
state.client.on('SetMessages', (e) => {
update((state) => {
const newState = {
...state,
Expand Down Expand Up @@ -120,11 +138,11 @@ export function createMessagesStore({
{ trailing: true, leading: true }
);

client.on('MessageEdited', (e) => {
state.client.on('MessageEdited', (e) => {
updateMessages(e.message);
});

client.on('UserJoined', (e) => {
state.client.on('UserJoined', (e) => {
update((state) => {
const newState = {
...state,
Expand All @@ -135,7 +153,7 @@ export function createMessagesStore({
callbacks?.UserJoined?.();
});

client.on('UserLeft', (e) => {
state.client.on('UserLeft', (e) => {
update((state) => {
const newState = {
...state,
Expand All @@ -145,22 +163,26 @@ export function createMessagesStore({
});
callbacks?.UserLeft?.();
});
} catch (e) {
update((v) => ({ ...v, error: 'Could not connect to server' }));
}
};

go();
});
})
.catch(() => {
// noop
// the user rejected the request by switching pages
});
} catch (e) {
update((v) => ({ ...v, error: 'Could not connect to server' }));
}

return () => socket?.close();
return () => {
// cleanup
state.abort?.abort();
};
});

return {
subscribe,
addMessage: (message: OptimisticMessage) => {
invariant(client, 'client should be defined');
client.send({ type: 'addMessage', ...message });
invariant(state.client, 'client should be defined');
state.client.send({ type: 'addMessage', ...message });
}
};
}
7 changes: 4 additions & 3 deletions apps/web/src/routes/api/party-auth/+server.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,13 @@ export const POST: RequestHandler = async ({ locals, request }) => {
const body = await request.json();
const { projectId, topicId } = body;

const url = `${PARTY_HOST}/party/${projectId}-${topicId}`;
console.log(url);
const roomId = `${projectId}-${topicId}`;
const url = `${PARTY_HOST}/party/${roomId}`;
const authResponse = await fetch(url, {
method: 'POST',
body: JSON.stringify({
userId: session.user.userId
userId: session.user.userId,
roomId
}),
headers: {
Authorization: `Bearer ${session.sessionId}`
Expand Down

0 comments on commit 7b6b604

Please sign in to comment.