|
| 1 | +import type { CloseEvent } from "@hocuspocus/common"; |
| 2 | +import { |
| 3 | + beforeHandleMessagePayload, |
| 4 | + Extension, |
| 5 | + IncomingMessage, |
| 6 | + MessageType, |
| 7 | +} from "@hocuspocus/server"; |
| 8 | +import * as syncProtocol from "y-protocols/sync"; |
| 9 | +import * as Y from "yjs"; |
| 10 | + |
| 11 | +/** |
| 12 | + * This extension rejects any changes to the restricted type. |
| 13 | + * |
| 14 | + * It does this by: |
| 15 | + * - extracting the yjsUpdate from the incoming message |
| 16 | + * - applying the update to the restricted type |
| 17 | + * - if the update is rejected, we throw an error and close the connection |
| 18 | + * - if the update is accepted, we do nothing |
| 19 | + */ |
| 20 | +export class RejectUnauthorized implements Extension { |
| 21 | + constructor(private readonly threadsMapKey: string) {} |
| 22 | + /** |
| 23 | + * Extract the yjsUpdate from the incoming message |
| 24 | + * @param message |
| 25 | + * @returns |
| 26 | + */ |
| 27 | + private getYUpdate(message: Uint8Array) { |
| 28 | + /** |
| 29 | + * The messages we are interested in are of the following format: |
| 30 | + * [docIdLength: number, ...docIdString: string, hocuspocusMessageType: number, ySyncMessageType: number, ...yjsUpdate: Uint8Array] |
| 31 | + * |
| 32 | + * We check that the hocuspocusMessageType is Sync and that the ySyncMessageType is messageYjsUpdate. |
| 33 | + */ |
| 34 | + const incomingMessage = new IncomingMessage(message); |
| 35 | + // Read the docID string, but don't use it |
| 36 | + incomingMessage.readVarString(); |
| 37 | + |
| 38 | + // Read the hocuspocusMessageType |
| 39 | + const hocuspocusMessageType = incomingMessage.readVarUint(); |
| 40 | + // If the hocuspocusMessageType is not Sync, we don't handle the message, since it is not an update |
| 41 | + if (hocuspocusMessageType !== MessageType.Sync) { |
| 42 | + return; |
| 43 | + } |
| 44 | + |
| 45 | + // Read the ySyncMessageType |
| 46 | + const ySyncMessageType = incomingMessage.readVarUint(); |
| 47 | + |
| 48 | + // If the ySyncMessageType is not messageYjsUpdate, we don't handle the message, since it is not an update |
| 49 | + if (ySyncMessageType !== syncProtocol.messageYjsUpdate) { |
| 50 | + // not an update |
| 51 | + return; |
| 52 | + } |
| 53 | + |
| 54 | + // Read the yjsUpdate |
| 55 | + const yUpdate = incomingMessage.readVarUint8Array(); |
| 56 | + |
| 57 | + return yUpdate; |
| 58 | + } |
| 59 | + |
| 60 | + /** |
| 61 | + * This function protects against changes to the restricted type. |
| 62 | + * It does this by: |
| 63 | + * - setting up an undo manager on the restricted type |
| 64 | + * - caching pending updates from the Ydoc to avoid certain attacks |
| 65 | + * - applying the received update and checking whether the restricted type has been changed |
| 66 | + * - catching errors that might try to circumvent the restrictions |
| 67 | + * - undoing changes on restricted types |
| 68 | + * - reapplying pending updates |
| 69 | + * |
| 70 | + * @param yUpdate The update to apply |
| 71 | + * @param ydoc The document that the update is being applied to |
| 72 | + * @param restrictedType The type that we want to protect |
| 73 | + * @returns true if the update was rejected, false otherwise |
| 74 | + */ |
| 75 | + private rollbackUpdateIfNeeded( |
| 76 | + yUpdate: Uint8Array, |
| 77 | + ydoc: Y.Doc, |
| 78 | + restrictedType: Y.AbstractType<any> |
| 79 | + ) { |
| 80 | + // don't handle changes of the local undo manager, which is used to undo invalid changes |
| 81 | + const um = new Y.UndoManager(restrictedType, { |
| 82 | + trackedOrigins: new Set(["remote change"]), |
| 83 | + }); |
| 84 | + const beforePendingDs = ydoc.store.pendingDs; |
| 85 | + const beforePendingStructs = ydoc.store.pendingStructs?.update; |
| 86 | + let didNeedToUndo = false; |
| 87 | + try { |
| 88 | + Y.applyUpdate(ydoc, yUpdate, "remote change"); |
| 89 | + } finally { |
| 90 | + while (um.undoStack.length) { |
| 91 | + um.undo(); |
| 92 | + didNeedToUndo = true; |
| 93 | + } |
| 94 | + um.destroy(); |
| 95 | + ydoc.store.pendingDs = beforePendingDs; |
| 96 | + ydoc.store.pendingStructs = null; |
| 97 | + if (beforePendingStructs) { |
| 98 | + Y.applyUpdateV2(ydoc, beforePendingStructs); |
| 99 | + } |
| 100 | + } |
| 101 | + |
| 102 | + return didNeedToUndo; |
| 103 | + } |
| 104 | + |
| 105 | + async beforeHandleMessage({ |
| 106 | + update, |
| 107 | + document: ydoc, |
| 108 | + }: beforeHandleMessagePayload) { |
| 109 | + const yUpdate = this.getYUpdate(update); |
| 110 | + |
| 111 | + if (!yUpdate) { |
| 112 | + return; |
| 113 | + } |
| 114 | + |
| 115 | + const protectedType = ydoc.getMap(this.threadsMapKey); |
| 116 | + const didNeedToUndo = this.rollbackUpdateIfNeeded( |
| 117 | + yUpdate, |
| 118 | + ydoc, |
| 119 | + protectedType |
| 120 | + ); |
| 121 | + |
| 122 | + if (didNeedToUndo) { |
| 123 | + // TODO, we can close their connection or just let them continue, since we've already undone their changes (and our changes are newer than theirs) |
| 124 | + const error = { |
| 125 | + reason: `Modification of a restricted type: ${this.threadsMapKey} was rejected`, |
| 126 | + } satisfies Partial<CloseEvent>; |
| 127 | + throw error; |
| 128 | + } |
| 129 | + } |
| 130 | +} |
0 commit comments