@replit/river 0.22.0 → 0.23.1

package/dist/transport/impls/ws/server.cjs

@@ -146,7 +146,7 @@ var import_nanoid2 = require("nanoid");
 var import_api = require("@opentelemetry/api");
 
 // package.json
-var version = "0.22.0";
+var version = "0.23.1";
 
 // tracing/index.ts
 function createSessionTelemetryInfo(session, propagationCtx) {
@@ -216,12 +216,6 @@ var Session = class {
    * for this session.
    */
   advertisedSessionId;
-  /**
-   * The metadata for this session, as parsed from the handshake.
-   *
-   * Will only ever be populated on the server side.
-   */
-  metadata;
   /**
    * Number of messages we've sent along this session (excluding handshake and acks)
    */
@@ -379,8 +373,8 @@ var Session = class {
   replaceWithNewConnection(newConn) {
     this.closeStaleConnection(newConn);
     this.cancelGrace();
-    this.connection = newConn;
     this.sendBufferedMessages(newConn);
+    this.connection = newConn;
   }
   beginGrace(cb) {
     log?.info(
@@ -808,17 +802,29 @@ var ServerTransport = class extends Transport {
    * The options for this transport.
    */
   options;
+  /**
+   * Optional handshake options for the server.
+   */
+  handshakeExtensions;
+  /**
+   * A map of session handshake data for each session.
+   */
+  sessionHandshakeMetadata;
   constructor(clientId, providedOptions) {
     super(clientId, providedOptions);
     this.options = {
       ...defaultServerTransportOptions,
       ...providedOptions
     };
+    this.sessionHandshakeMetadata = /* @__PURE__ */ new WeakMap();
     log?.info(`initiated server transport`, {
       clientId: this.clientId,
       protocolVersion: PROTOCOL_VERSION
     });
   }
+  extendHandshake(options) {
+    this.handshakeExtensions = options;
+  }
   handleConnection(conn) {
     if (this.state !== "open")
       return;
@@ -901,6 +907,52 @@ var ServerTransport = class extends Transport {
       );
     });
   }
+  async validateHandshakeMetadata(conn, session, rawMetadata, from) {
+    let parsedMetadata = {};
+    if (this.handshakeExtensions) {
+      if (!import_value.Value.Check(this.handshakeExtensions.schema, rawMetadata)) {
+        conn.telemetry?.span.setStatus({
+          code: import_api3.SpanStatusCode.ERROR,
+          message: "malformed handshake meta"
+        });
+        const reason = "received malformed handshake metadata";
+        const responseMsg = handshakeResponseMessage(this.clientId, from, {
+          ok: false,
+          reason
+        });
+        conn.send(this.codec.toBuffer(responseMsg));
+        log?.warn(`received malformed handshake metadata from ${from}`, {
+          clientId: this.clientId,
+          connId: conn.id
+        });
+        this.protocolError(ProtocolError.HandshakeFailed, reason);
+        return false;
+      }
+      parsedMetadata = await this.handshakeExtensions.validate(
+        rawMetadata,
+        session
+      );
+      if (parsedMetadata === false) {
+        const reason = "rejected by handshake handler";
+        conn.telemetry?.span.setStatus({
+          code: import_api3.SpanStatusCode.ERROR,
+          message: reason
+        });
+        const responseMsg = handshakeResponseMessage(this.clientId, from, {
+          ok: false,
+          reason
+        });
+        conn.send(this.codec.toBuffer(responseMsg));
+        log?.warn(`rejected handshake from ${from}`, {
+          clientId: this.clientId,
+          connId: conn.id
+        });
+        this.protocolError(ProtocolError.HandshakeFailed, reason);
+        return false;
+      }
+    }
+    return parsedMetadata;
+  }
   async receiveHandshakeRequestMessage(data, conn) {
     const parsed = this.parseMsg(data);
     if (!parsed) {
@@ -956,86 +1008,23 @@ var ServerTransport = class extends Transport {
       this.protocolError(ProtocolError.HandshakeFailed, reason);
       return false;
     }
+    const oldSession = this.sessions.get(parsed.from);
+    const parsedMetadata = await this.validateHandshakeMetadata(
+      conn,
+      oldSession,
+      parsed.payload.metadata,
+      parsed.from
+    );
+    if (parsedMetadata === false) {
+      return false;
+    }
     const { session, isReconnect } = this.getOrCreateSession(
       parsed.from,
       conn,
       parsed.payload.sessionId,
       parsed.tracing
     );
-    let handshakeMetadata;
-    if (this.options.handshake) {
-      if (!import_value.Value.Check(
-        this.options.handshake.requestSchema,
-        parsed.payload.metadata
-      )) {
-        conn.telemetry?.span.setStatus({
-          code: import_api3.SpanStatusCode.ERROR,
-          message: "malformed handshake meta"
-        });
-        const reason = "received malformed handshake metadata";
-        const responseMsg2 = handshakeResponseMessage(
-          this.clientId,
-          parsed.from,
-          { ok: false, reason }
-        );
-        conn.send(this.codec.toBuffer(responseMsg2));
-        log?.warn(`received malformed handshake metadata from ${parsed.from}`, {
-          clientId: this.clientId,
-          connId: conn.id
-        });
-        this.protocolError(ProtocolError.HandshakeFailed, reason);
-        this.deleteSession(session);
-        return false;
-      }
-      const parsedMetadata = await this.options.handshake.parse(
-        parsed.payload.metadata,
-        session,
-        isReconnect
-      );
-      if (parsedMetadata === false) {
-        conn.telemetry?.span.setStatus({
-          code: import_api3.SpanStatusCode.ERROR,
-          message: "rejected by handshake handler"
-        });
-        const reason = "rejected by handshake handler";
-        const responseMsg2 = handshakeResponseMessage(
-          this.clientId,
-          parsed.from,
-          { ok: false, reason }
-        );
-        conn.send(this.codec.toBuffer(responseMsg2));
-        log?.warn(`rejected handshake from ${parsed.from}`, {
-          clientId: this.clientId,
-          connId: conn.id
-        });
-        this.protocolError(ProtocolError.HandshakeFailed, reason);
-        this.deleteSession(session);
-        return false;
-      }
-      if (!import_value.Value.Check(this.options.handshake.parsedSchema, parsedMetadata)) {
-        conn.telemetry?.span.setStatus({
-          code: import_api3.SpanStatusCode.ERROR,
-          message: "malformed handshake meta"
-        });
-        const reason = "failed to parse handshake metadata";
-        const responseMsg2 = handshakeResponseMessage(
-          this.clientId,
-          parsed.from,
-          { ok: false, reason }
-        );
-        conn.send(this.codec.toBuffer(responseMsg2));
-        log?.error(`failed to parse handshake metadata`, {
-          clientId: this.clientId,
-          connId: conn.id
-        });
-        this.protocolError(ProtocolError.HandshakeFailed, reason);
-        this.deleteSession(session);
-        return false;
-      }
-      handshakeMetadata = parsedMetadata;
-    }
-    handshakeMetadata ??= {};
-    session.metadata = handshakeMetadata;
+    this.sessionHandshakeMetadata.set(session, parsedMetadata);
     log?.debug(
       `handshake from ${parsed.from} ok, responding with handshake success`,
       { clientId: this.clientId, connId: conn.id }