@replit/river 0.23.13 → 0.23.15

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

@@ -59,6 +59,21 @@ var ControlMessageHandshakeRequestSchema = import_typebox.Type.Object({
   type: import_typebox.Type.Literal("HANDSHAKE_REQ"),
   protocolVersion: import_typebox.Type.String(),
   sessionId: import_typebox.Type.String(),
+  /**
+   * Specifies what the server's expected session state (from the pov of the client). This can be
+   * used by the server to know whether this is a new or a reestablished connection, and whether it
+   * is compatible with what it already has.
+   */
+  expectedSessionState: import_typebox.Type.Optional(
+    import_typebox.Type.Object({
+      /**
+       * reconnect is set to true if the client explicitly wants to reestablish an existing
+       * connection.
+       */
+      reconnect: import_typebox.Type.Boolean(),
+      nextExpectedSeq: import_typebox.Type.Integer()
+    })
+  ),
   metadata: import_typebox.Type.Optional(import_typebox.Type.Unknown())
 });
 var ControlMessageHandshakeResponseSchema = import_typebox.Type.Object({
@@ -83,7 +98,12 @@ var ControlMessagePayloadSchema = import_typebox.Type.Union([
 var OpaqueTransportMessageSchema = TransportMessageSchema(
   import_typebox.Type.Unknown()
 );
-function handshakeResponseMessage(from, to, status) {
+var SESSION_STATE_MISMATCH = "session state mismatch";
+function handshakeResponseMessage({
+  from,
+  to,
+  status
+}) {
   return {
     id: (0, import_nanoid.nanoid)(),
     from,
@@ -106,7 +126,7 @@ function isAck(controlFlag) {
 var import_api = require("@opentelemetry/api");
 
 // package.json
-var version = "0.23.13";
+var version = "0.23.15";
 
 // tracing/index.ts
 function createSessionTelemetryInfo(session, propagationCtx) {
@@ -411,9 +431,17 @@ var Session = class {
   get connected() {
     return this.connection !== void 0;
   }
+  get nextExpectedAck() {
+    return this.seq;
+  }
   get nextExpectedSeq() {
     return this.ack;
   }
+  // This is only used in tests to make the session misbehave.
+  /* @internal */
+  advanceAckForTesting(by) {
+    this.ack += by;
+  }
   constructMsg(partialMsg) {
     const msg = {
       ...partialMsg,
@@ -738,6 +766,20 @@ var Transport = class {
     if (this.log) {
       session.bindLogger(this.log);
     }
+    const currentSession = this.sessions.get(session.to);
+    if (currentSession) {
+      this.log?.warn(
+        `session ${session.id} from ${session.to} surreptitiously replacing ${currentSession.id}`,
+        {
+          ...currentSession.loggingMetadata,
+          tags: ["invariant-violation"]
+        }
+      );
+      this.deleteSession({
+        session: currentSession,
+        closeHandshakingConnection: false
+      });
+    }
     this.sessions.set(session.to, session);
     this.eventDispatcher.dispatchEvent("sessionStatus", {
       status: "connect",
@@ -745,6 +787,49 @@ var Transport = class {
     });
     return session;
   }
+  createNewSession({
+    to,
+    conn,
+    sessionId,
+    propagationCtx
+  }) {
+    let session = this.sessions.get(to);
+    if (session !== void 0) {
+      this.log?.info(
+        `session for ${to} already exists, replacing it with a new session as requested`,
+        session.loggingMetadata
+      );
+      this.deleteSession({
+        session,
+        closeHandshakingConnection: false
+      });
+      session = void 0;
+    }
+    session = this.createSession(to, conn, propagationCtx);
+    session.advertisedSessionId = sessionId;
+    this.log?.info(`created new session for ${to}`, session.loggingMetadata);
+    return session;
+  }
+  getExistingSession({
+    to,
+    sessionId,
+    nextExpectedSeq
+  }) {
+    const session = this.sessions.get(to);
+    if (
+      // reject this request if there was no previous session to replace
+      session === void 0 || // or if both parties do not agree about the next expected sequence number
+      session.nextExpectedAck < nextExpectedSeq || // or if both parties do not agree on the advertised session id
+      session.advertisedSessionId !== sessionId
+    ) {
+      return false;
+    }
+    this.log?.info(
+      `reused existing session for ${to}`,
+      session.loggingMetadata
+    );
+    return session;
+  }
   getOrCreateSession({
     to,
     conn,
@@ -793,6 +878,17 @@ var Transport = class {
     }
     session.close();
     session.telemetry.span.end();
+    const currentSession = this.sessions.get(session.to);
+    if (currentSession && currentSession.id !== session.id) {
+      this.log?.warn(
+        `session ${session.id} disconnect from ${session.to}, mismatch with ${currentSession.id}`,
+        {
+          ...session.loggingMetadata,
+          tags: ["invariant-violation"]
+        }
+      );
+      return;
+    }
     this.sessions.delete(session.to);
     this.log?.info(
       `session ${session.id} disconnect from ${session.to}`,
@@ -1133,9 +1229,13 @@ var ServerTransport = class extends Transport {
           message: "malformed handshake meta"
         });
         const reason = "received malformed handshake metadata";
-        const responseMsg = handshakeResponseMessage(this.clientId, from, {
-          ok: false,
-          reason
+        const responseMsg = handshakeResponseMessage({
+          from: this.clientId,
+          to: from,
+          status: {
+            ok: false,
+            reason
+          }
         });
         conn.send(this.codec.toBuffer(responseMsg));
         this.log?.warn(`received malformed handshake metadata from ${from}`, {
@@ -1159,9 +1259,13 @@ var ServerTransport = class extends Transport {
           code: import_api4.SpanStatusCode.ERROR,
           message: reason
         });
-        const responseMsg = handshakeResponseMessage(this.clientId, from, {
-          ok: false,
-          reason
+        const responseMsg = handshakeResponseMessage({
+          from: this.clientId,
+          to: from,
+          status: {
+            ok: false,
+            reason
+          }
         });
         conn.send(this.codec.toBuffer(responseMsg));
         this.log?.warn(`rejected handshake from ${from}`, {
@@ -1193,9 +1297,13 @@ var ServerTransport = class extends Transport {
         message: "invalid handshake request"
       });
       const reason = "received invalid handshake msg";
-      const responseMsg2 = handshakeResponseMessage(this.clientId, parsed.from, {
-        ok: false,
-        reason
+      const responseMsg2 = handshakeResponseMessage({
+        from: this.clientId,
+        to: parsed.from,
+        status: {
+          ok: false,
+          reason
+        }
       });
       conn.send(this.codec.toBuffer(responseMsg2));
       this.log?.warn(reason, {
@@ -1221,9 +1329,13 @@ var ServerTransport = class extends Transport {
         message: "incorrect protocol version"
       });
       const reason = `incorrect version (got: ${gotVersion} wanted ${PROTOCOL_VERSION})`;
-      const responseMsg2 = handshakeResponseMessage(this.clientId, parsed.from, {
-        ok: false,
-        reason
+      const responseMsg2 = handshakeResponseMessage({
+        from: this.clientId,
+        to: parsed.from,
+        status: {
+          ok: false,
+          reason
+        }
       });
       conn.send(this.codec.toBuffer(responseMsg2));
       this.log?.warn(
@@ -1243,20 +1355,67 @@ var ServerTransport = class extends Transport {
     if (parsedMetadata === false) {
       return false;
     }
-    const { session, isTransparentReconnect } = this.getOrCreateSession({
-      to: parsed.from,
-      conn,
-      sessionId: parsed.payload.sessionId,
-      propagationCtx: parsed.tracing
-    });
+    let session;
+    let isTransparentReconnect;
+    if (!parsed.payload.expectedSessionState) {
+      ({ session, isTransparentReconnect } = this.getOrCreateSession({
+        to: parsed.from,
+        conn,
+        sessionId: parsed.payload.sessionId,
+        propagationCtx: parsed.tracing
+      }));
+    } else if (parsed.payload.expectedSessionState.reconnect) {
+      const existingSession = this.getExistingSession({
+        to: parsed.from,
+        sessionId: parsed.payload.sessionId,
+        nextExpectedSeq: parsed.payload.expectedSessionState.nextExpectedSeq
+      });
+      if (existingSession === false) {
+        conn.telemetry?.span.setStatus({
+          code: import_api4.SpanStatusCode.ERROR,
+          message: SESSION_STATE_MISMATCH
+        });
+        const reason = SESSION_STATE_MISMATCH;
+        const responseMsg2 = handshakeResponseMessage({
+          from: this.clientId,
+          to: parsed.from,
+          status: {
+            ok: false,
+            reason
+          }
+        });
+        conn.send(this.codec.toBuffer(responseMsg2));
+        this.log?.warn(
+          `'received handshake msg with incompatible existing session state: ${parsed.payload.sessionId}`,
+          { ...conn.loggingMetadata, clientId: this.clientId }
+        );
+        this.protocolError(ProtocolError.HandshakeFailed, reason);
+        return false;
+      }
+      session = existingSession;
+      isTransparentReconnect = false;
+    } else {
+      const createdSession = this.createNewSession({
+        to: parsed.from,
+        conn,
+        sessionId: parsed.payload.sessionId,
+        propagationCtx: parsed.tracing
+      });
+      session = createdSession;
+      isTransparentReconnect = false;
+    }
     this.sessionHandshakeMetadata.set(session, parsedMetadata);
     this.log?.debug(
       `handshake from ${parsed.from} ok, responding with handshake success`,
       conn.loggingMetadata
     );
-    const responseMsg = handshakeResponseMessage(this.clientId, parsed.from, {
-      ok: true,
-      sessionId: session.id
+    const responseMsg = handshakeResponseMessage({
+      from: this.clientId,
+      to: parsed.from,
+      status: {
+        ok: true,
+        sessionId: session.id
+      }
     });
     conn.send(this.codec.toBuffer(responseMsg));
     this.onConnect(conn, session, isTransparentReconnect);