feat: add user locking to isolation reservation table

CHANGELOG.md

@@ -7,6 +7,10 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+## [9.3.1]
+
+- Bump version
+
 ## [9.3.0]
 
 - Adds SAML support
@@ -44,6 +48,7 @@ CREATE TABLE IF NOT EXISTS saml_relay_state (
     state TEXT NOT NULL,
     redirect_uri TEXT NOT NULL,
     created_at BIGINT NOT NULL,
+    expires_at BIGINT NOT NULL,
     CONSTRAINT saml_relay_state_pkey PRIMARY KEY(app_id, tenant_id, relay_state),
     CONSTRAINT saml_relay_state_app_id_fkey FOREIGN KEY(app_id) REFERENCES apps (app_id) ON DELETE CASCADE,
     CONSTRAINT saml_relay_state_tenant_id_fkey FOREIGN KEY(app_id, tenant_id) REFERENCES tenants (app_id, tenant_id) ON DELETE CASCADE
@@ -59,6 +64,7 @@ CREATE TABLE IF NOT EXISTS saml_claims (
     code VARCHAR(256) NOT NULL,
     claims TEXT NOT NULL,
     created_at BIGINT NOT NULL,
+    expires_at BIGINT NOT NULL,
     CONSTRAINT saml_claims_pkey PRIMARY KEY(app_id, tenant_id, code),
     CONSTRAINT saml_claims_app_id_fkey FOREIGN KEY(app_id) REFERENCES apps (app_id) ON DELETE CASCADE,
     CONSTRAINT saml_claims_tenant_id_fkey FOREIGN KEY(app_id, tenant_id) REFERENCES tenants (app_id, tenant_id) ON DELETE CASCADE

build.gradle

@@ -128,4 +128,4 @@ tasks.withType(Test).configureEach {
             }
         }
     }
-}
+}

implementationDependencies.json

@@ -6,16 +6,26 @@
 		"name":"jackson-dataformat-yaml 2.18.2",
 		"src":"https://repo.maven.apache.org/maven2/com/fasterxml/jackson/dataformat/jackson-dataformat-yaml/2.18.2/jackson-dataformat-yaml-2.18.2-sources.jar"
 	},
-	{
-		"jar":"https://repo.maven.apache.org/maven2/org/yaml/snakeyaml/2.3/snakeyaml-2.3.jar",
-		"name":"snakeyaml 2.3",
-		"src":"https://repo.maven.apache.org/maven2/org/yaml/snakeyaml/2.3/snakeyaml-2.3-sources.jar"
-	},
 	{
 		"jar":"https://repo.maven.apache.org/maven2/com/fasterxml/jackson/core/jackson-databind/2.18.2/jackson-databind-2.18.2.jar",
 		"name":"jackson-databind 2.18.2",
 		"src":"https://repo.maven.apache.org/maven2/com/fasterxml/jackson/core/jackson-databind/2.18.2/jackson-databind-2.18.2-sources.jar"
 	},
+	{
+		"jar":"https://repo.maven.apache.org/maven2/com/fasterxml/jackson/core/jackson-core/2.18.2/jackson-core-2.18.2.jar",
+		"name":"jackson-core 2.18.2",
+		"src":"https://repo.maven.apache.org/maven2/com/fasterxml/jackson/core/jackson-core/2.18.2/jackson-core-2.18.2-sources.jar"
+	},
+	{
+		"jar":"https://repo.maven.apache.org/maven2/com/fasterxml/jackson/core/jackson-annotations/2.18.2/jackson-annotations-2.18.2.jar",
+		"name":"jackson-annotations 2.18.2",
+		"src":"https://repo.maven.apache.org/maven2/com/fasterxml/jackson/core/jackson-annotations/2.18.2/jackson-annotations-2.18.2-sources.jar"
+	},
+	{
+		"jar":"https://repo.maven.apache.org/maven2/org/yaml/snakeyaml/2.3/snakeyaml-2.3.jar",
+		"name":"snakeyaml 2.3",
+		"src":"https://repo.maven.apache.org/maven2/org/yaml/snakeyaml/2.3/snakeyaml-2.3-sources.jar"
+	},
 	{
 		"jar":"https://repo.maven.apache.org/maven2/ch/qos/logback/logback-classic/1.5.13/logback-classic-1.5.13.jar",
 		"name":"logback-classic 1.5.13",

src/main/java/io/supertokens/storage/postgresql/LockedUserImpl.java

@@ -0,0 +1,104 @@
+/*
+ *    Copyright (c) 2024, VRAI Labs and/or its affiliates. All rights reserved.
+ *
+ *    This software is licensed under the Apache License, Version 2.0 (the
+ *    "License") as published by the Apache Software Foundation.
+ *
+ *    You may not use this file except in compliance with the License. You may
+ *    obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    Unless required by applicable law or agreed to in writing, software
+ *    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *    License for the specific language governing permissions and limitations
+ *    under the License.
+ */
+
+package io.supertokens.storage.postgresql;
+
+import io.supertokens.pluginInterface.useridmapping.LockedUser;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import java.sql.Connection;
+import java.lang.ref.WeakReference;
+
+/**
+ * PostgreSQL implementation of LockedUser.
+ * Tracks the connection to validate the lock is still active.
+ */
+public class LockedUserImpl implements LockedUser {
+
+    @Nonnull
+    private final String recipeUserId;
+
+    @Nonnull
+    private final String recipeId;
+
+    @Nullable
+    private final String primaryUserId;
+
+    // WeakReference so we don't prevent connection from being garbage collected
+    private final WeakReference<Connection> connectionRef;
+
+    public LockedUserImpl(@Nonnull String recipeUserId, @Nonnull String recipeId,
+                          @Nullable String primaryUserId, @Nonnull Connection connection) {
+        this.recipeUserId = recipeUserId;
+        this.recipeId = recipeId;
+        this.primaryUserId = primaryUserId;
+        this.connectionRef = new WeakReference<>(connection);
+    }
+
+    @Override
+    @Nonnull
+    public String getRecipeUserId() {
+        return recipeUserId;
+    }
+
+    @Override
+    @Nonnull
+    public String getRecipeId() {
+        return recipeId;
+    }
+
+    @Override
+    @Nullable
+    public String getPrimaryUserId() {
+        return primaryUserId;
+    }
+
+    @Override
+    public boolean isValidForConnection(Object connection) {
+        Connection originalCon = connectionRef.get();
+        if (originalCon == null) {
+            return false;
+        }
+        // Check that the provided connection is the same instance as the one used to acquire the lock
+        if (originalCon != connection) {
+            return false;
+        }
+        try {
+            return !originalCon.isClosed();
+        } catch (Exception e) {
+            return false;
+        }
+    }
+
+    @Override
+    public String toString() {
+        Connection con = connectionRef.get();
+        boolean connectionAlive = false;
+        try {
+            connectionAlive = con != null && !con.isClosed();
+        } catch (Exception ignored) {
+        }
+        return "LockedUser{" +
+               "recipeUserId='" + recipeUserId + '\'' +
+               ", recipeId='" + recipeId + '\'' +
+               ", primaryUserId='" + primaryUserId + '\'' +
+               ", isLinked=" + isLinked() +
+               ", isPrimary=" + isPrimary() +
+               ", connectionAlive=" + connectionAlive +
+               '}';
+    }
+}