From 9cd08ec0eb1ca9d0b84c9ea5a6a98932b23439e3 Mon Sep 17 00:00:00 2001
From: Sattvik Chakravarthy <sattvik@gmail.com>
Date: Wed, 13 Sep 2023 15:36:07 +0530
Subject: [PATCH] fix: password reset changes

---
 .../java/io/supertokens/inmemorydb/Start.java |  2 +-
 .../queries/EmailPasswordQueries.java         | 42 ++++++++++++++-----
 2 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/src/main/java/io/supertokens/inmemorydb/Start.java b/src/main/java/io/supertokens/inmemorydb/Start.java
index 91f3daa1d..10b365bfa 100644
--- a/src/main/java/io/supertokens/inmemorydb/Start.java
+++ b/src/main/java/io/supertokens/inmemorydb/Start.java
@@ -777,7 +777,7 @@ public void addPasswordResetToken(AppIdentifier appIdentifier, PasswordResetToke
             throws StorageQueryException, UnknownUserIdException, DuplicatePasswordResetTokenException {
         try {
             EmailPasswordQueries.addPasswordResetToken(this, appIdentifier, passwordResetTokenInfo.userId,
-                    passwordResetTokenInfo.token, passwordResetTokenInfo.tokenExpiry);
+                    passwordResetTokenInfo.token, passwordResetTokenInfo.tokenExpiry, passwordResetTokenInfo.email);
         } catch (SQLException e) {
             if (e instanceof SQLiteException) {
                 String serverMessage = e.getMessage();
diff --git a/src/main/java/io/supertokens/inmemorydb/queries/EmailPasswordQueries.java b/src/main/java/io/supertokens/inmemorydb/queries/EmailPasswordQueries.java
index 003a100b3..360186f63 100644
--- a/src/main/java/io/supertokens/inmemorydb/queries/EmailPasswordQueries.java
+++ b/src/main/java/io/supertokens/inmemorydb/queries/EmailPasswordQueries.java
@@ -81,7 +81,7 @@ static String getQueryToCreatePasswordResetTokensTable(Start start) {
                 + "token VARCHAR(128) NOT NULL UNIQUE,"
                 + "token_expiry BIGINT UNSIGNED NOT NULL,"
                 + "PRIMARY KEY (app_id, user_id, token),"
-                + "FOREIGN KEY (app_id, user_id) REFERENCES " + Config.getConfig(start).getEmailPasswordUsersTable()
+                + "FOREIGN KEY (app_id, user_id) REFERENCES " + Config.getConfig(start).getAppIdToUserIdTable()
                 + " (app_id, user_id) ON DELETE CASCADE ON UPDATE CASCADE"
                 + ");";
     }
@@ -215,17 +215,30 @@ public static PasswordResetTokenInfo getPasswordResetTokenInfo(Start start, AppI
     }
 
     public static void addPasswordResetToken(Start start, AppIdentifier appIdentifier, String userId, String tokenHash,
-                                             long expiry)
+                                             long expiry, String email)
             throws SQLException, StorageQueryException {
-        String QUERY = "INSERT INTO " + getConfig(start).getPasswordResetTokensTable()
-                + "(app_id, user_id, token, token_expiry)" + " VALUES(?, ?, ?, ?)";
+        if (email != null) {
+            String QUERY = "INSERT INTO " + getConfig(start).getPasswordResetTokensTable()
+                    + "(app_id, user_id, token, token_expiry, email)" + " VALUES(?, ?, ?, ?, ?)";
 
-        update(start, QUERY, pst -> {
-            pst.setString(1, appIdentifier.getAppId());
-            pst.setString(2, userId);
-            pst.setString(3, tokenHash);
-            pst.setLong(4, expiry);
-        });
+            update(start, QUERY, pst -> {
+                pst.setString(1, appIdentifier.getAppId());
+                pst.setString(2, userId);
+                pst.setString(3, tokenHash);
+                pst.setLong(4, expiry);
+                pst.setString(5, email);
+            });
+        } else {
+            String QUERY = "INSERT INTO " + getConfig(start).getPasswordResetTokensTable()
+                    + "(app_id, user_id, token, token_expiry)" + " VALUES(?, ?, ?, ?)";
+
+            update(start, QUERY, pst -> {
+                pst.setString(1, appIdentifier.getAppId());
+                pst.setString(2, userId);
+                pst.setString(3, tokenHash);
+                pst.setLong(4, expiry);
+            });
+        }
     }
 
     public static AuthRecipeUserInfo signUp(Start start, TenantIdentifier tenantIdentifier, String userId, String email,
@@ -322,6 +335,15 @@ public static void deleteUser_Transaction(Connection sqlCon, Start start, AppIde
                     pst.setString(2, userId);
                 });
             }
+
+            {
+                String QUERY = "DELETE FROM " + getConfig(start).getPasswordResetTokensTable()
+                        + " WHERE app_id = ? AND user_id = ?";
+                update(sqlCon, QUERY, pst -> {
+                    pst.setString(1, appIdentifier.getAppId());
+                    pst.setString(2, userId);
+                });
+            }
         }
     }