@replit/river 0.22.0 → 0.23.1

package/dist/transport/index.cjs

@@ -171,7 +171,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 getPropagationContext(ctx) {
@@ -250,12 +250,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)
    */
@@ -413,8 +407,8 @@ var Session = class {
   replaceWithNewConnection(newConn) {
     this.closeStaleConnection(newConn);
     this.cancelGrace();
-    this.connection = newConn;
     this.sendBufferedMessages(newConn);
+    this.connection = newConn;
   }
   beginGrace(cb) {
     log?.info(
@@ -925,6 +919,10 @@ var ClientTransport = class extends Transport {
    * tests or a special case where you don't want to reconnect.
    */
   reconnectOnConnectionDrop = true;
+  /**
+   * Optional handshake options for this client.
+   */
+  handshakeExtensions;
   constructor(clientId, providedOptions) {
     super(clientId, providedOptions);
     this.options = {
@@ -934,6 +932,9 @@ var ClientTransport = class extends Transport {
     this.inflightConnectionPromises = /* @__PURE__ */ new Map();
     this.retryBudget = new LeakyBucketRateLimit(this.options);
   }
+  extendHandshake(options) {
+    this.handshakeExtensions = options;
+  }
   handleConnection(conn, to) {
     if (this.state !== "open")
       return;
@@ -1157,10 +1158,10 @@ var ClientTransport = class extends Transport {
   }
   async sendHandshake(to, conn) {
     let metadata;
-    if (this.options.handshake) {
-      metadata = await this.options.handshake.get();
-      if (!import_value.Value.Check(this.options.handshake.schema, metadata)) {
-        log?.error(`handshake metadata did not match schema`, {
+    if (this.handshakeExtensions) {
+      metadata = await this.handshakeExtensions.construct();
+      if (!import_value.Value.Check(this.handshakeExtensions.schema, metadata)) {
+        log?.error(`constructed handshake metadata did not match schema`, {
           clientId: this.clientId,
           connectedTo: to,
           tags: ["invariant-violation"]
@@ -1176,7 +1177,7 @@ var ClientTransport = class extends Transport {
         return false;
       }
     }
-    const { session } = this.getOrCreateSession(to, conn);
+    const { session } = this.getOrCreateSession(to);
     const requestMsg = handshakeRequestMessage(
       this.clientId,
       to,
@@ -1201,17 +1202,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;
@@ -1294,6 +1307,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) {
@@ -1349,86 +1408,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 }