From 19003cd40e4e2c4f420d069d484ddf4ef1a96bfa Mon Sep 17 00:00:00 2001 From: Bereket Engida Date: Sun, 13 Oct 2024 16:06:17 +0300 Subject: [PATCH] feat: use secondary storage as primary for sessions --- .../src/adapters/kysely-adapter/dialect.ts | 1 + .../__snapshots__/adapter.kysley.test.ts.snap | 12 ++- .../test/adapter.kysley.test.ts | 11 +-- .../src/api/routes/session.test.ts | 45 ++++++++++- packages/better-auth/src/db/get-tables.ts | 22 ++++-- .../better-auth/src/db/internal-adapter.ts | 75 +++++++++++++++++-- .../better-auth/src/integrations/next-js.ts | 14 +++- packages/better-auth/src/types/options.ts | 12 +++ 8 files changed, 164 insertions(+), 28 deletions(-) diff --git a/packages/better-auth/src/adapters/kysely-adapter/dialect.ts b/packages/better-auth/src/adapters/kysely-adapter/dialect.ts index f239af6f7b..7580c7b83b 100644 --- a/packages/better-auth/src/adapters/kysely-adapter/dialect.ts +++ b/packages/better-auth/src/adapters/kysely-adapter/dialect.ts @@ -45,6 +45,7 @@ function getDatabaseType( export const createKyselyAdapter = async (config: BetterAuthOptions) => { const db = config.database; + if ("db" in db) { return { kysely: db.db, diff --git a/packages/better-auth/src/adapters/kysely-adapter/test/__snapshots__/adapter.kysley.test.ts.snap b/packages/better-auth/src/adapters/kysely-adapter/test/__snapshots__/adapter.kysley.test.ts.snap index c030b5ba55..888120be87 100644 --- a/packages/better-auth/src/adapters/kysely-adapter/test/__snapshots__/adapter.kysley.test.ts.snap +++ b/packages/better-auth/src/adapters/kysely-adapter/test/__snapshots__/adapter.kysley.test.ts.snap @@ -1,11 +1,17 @@ // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html exports[`adapter test > should create schema > __snapshots__/adapter.drizzle 1`] = ` -"create table "session" ("id" text primary key, "expiresAt" date not null, "ipAddress" text, "userAgent" text, "userId" text not null references "user" ("id")); +"create table "user" ("id" text primary key, "name" text not null, "email" text not null unique, "emailVerified" boolean not null, "image" text, "createdAt" date not null, "updatedAt" date not null, "twoFactorEnabled" boolean, "twoFactorSecret" text, "twoFactorBackupCodes" text); + +create table "session" ("id" text primary key, "expiresAt" date not null, "ipAddress" text, "userAgent" text, "userId" text not null references "user" ("id"), "activeOrganizationId" text); create table "account" ("id" text primary key, "accountId" text not null, "providerId" text not null, "userId" text not null references "user" ("id"), "accessToken" text, "refreshToken" text, "idToken" text, "expiresAt" date, "password" text); -create table "user" ("id" text primary key, "name" text not null, "email" text not null unique, "emailVerified" boolean not null, "image" text, "createdAt" date not null, "updatedAt" date not null); +create table "verification" ("id" text primary key, "identifier" text not null, "value" text not null, "expiresAt" date not null); -create table "verification" ("id" text primary key, "identifier" text not null, "value" text not null, "expiresAt" date not null)" +create table "organization" ("id" text primary key, "name" text not null, "slug" text not null unique, "logo" text, "createdAt" date not null, "metadata" text); + +create table "member" ("id" text primary key, "organizationId" text not null references "organization" ("id"), "userId" text not null, "email" text not null, "role" text not null, "createdAt" date not null); + +create table "invitation" ("id" text primary key, "organizationId" text not null references "organization" ("id"), "email" text not null, "role" text, "status" text not null, "expiresAt" date not null, "inviterId" text not null references "user" ("id"))" `; diff --git a/packages/better-auth/src/adapters/kysely-adapter/test/adapter.kysley.test.ts b/packages/better-auth/src/adapters/kysely-adapter/test/adapter.kysley.test.ts index 83b303fbb7..ecc3e162bf 100644 --- a/packages/better-auth/src/adapters/kysely-adapter/test/adapter.kysley.test.ts +++ b/packages/better-auth/src/adapters/kysely-adapter/test/adapter.kysley.test.ts @@ -8,6 +8,7 @@ import Database from "better-sqlite3"; import { kyselyAdapter } from ".."; import { Kysely, SqliteDialect } from "kysely"; import { getTestInstance } from "../../../test-utils/test-instance"; +import { organization, twoFactor } from "../../../plugins"; describe("adapter test", async () => { const database = new Database(path.join(__dirname, "test.db")); @@ -39,6 +40,7 @@ describe("adapter test", async () => { it("should create schema", async () => { const res = await adapter.createSchema!({ database: new Database(path.join(__dirname, "test-2.db")), + plugins: [organization(), twoFactor()], }); expect(res.code).toMatchSnapshot("__snapshots__/adapter.drizzle"); }); @@ -46,13 +48,4 @@ describe("adapter test", async () => { await runAdapterTest({ adapter, }); - - it("should support kysely instance", async () => { - const { auth } = await getTestInstance({ - database: { - db, - type: "sqlite", - }, - }); - }); }); diff --git a/packages/better-auth/src/api/routes/session.test.ts b/packages/better-auth/src/api/routes/session.test.ts index 7beae13a73..44fa8ed677 100644 --- a/packages/better-auth/src/api/routes/session.test.ts +++ b/packages/better-auth/src/api/routes/session.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it, vi } from "vitest"; import { getTestInstance } from "../../test-utils/test-instance"; import { parseSetCookieHeader } from "../../cookies"; import { getDate } from "../../utils/date"; +import type { Session } from "../../types"; describe("session", async () => { const { client, testUser, sessionSetter } = await getTestInstance(); @@ -134,7 +135,6 @@ describe("session", async () => { onSuccess(context) { const header = context.response.headers.get("set-cookie"); const cookies = parseSetCookieHeader(header || ""); - expect(cookies.get("better-auth.session_token")).toMatchObject({ value: expect.any(String), "max-age": (60 * 60 * 24 * 7).toString(), @@ -244,7 +244,7 @@ describe("session", async () => { describe("session storage", async () => { let store = new Map(); - const { client, signInWithTestUser } = await getTestInstance({ + const { client, signInWithTestUser, db } = await getTestInstance({ secondaryStorage: { set(key, value, ttl) { store.set(key, value); @@ -290,4 +290,45 @@ describe("session storage", async () => { }, }); }); + + it("should only store session in database if option is set", async () => { + await signInWithTestUser(); + const sessionCount = await db.findMany({ + model: "session", + }); + expect(sessionCount.length).toBe(0); + const { db: db2, signInWithTestUser: signInWithTestUser2 } = + await getTestInstance({ + session: { + storeSessionInDatabase: false, + }, + }); + await signInWithTestUser2(); + const sessionCount2 = await db2.findMany({ + model: "session", + }); + expect(sessionCount2.length).toBe(2); + }); + + it("should revoke session", async () => { + const { headers } = await signInWithTestUser(); + const session = await client.session({ + fetchOptions: { + headers, + }, + }); + expect(session.data).not.toBeNull(); + await client.user.revokeSession({ + fetchOptions: { + headers, + }, + id: session.data?.session?.id || "", + }); + const revokedSession = await client.session({ + fetchOptions: { + headers, + }, + }); + expect(revokedSession.data).toBeNull(); + }); }); diff --git a/packages/better-auth/src/db/get-tables.ts b/packages/better-auth/src/db/get-tables.ts index 68449a8beb..e288ee61fb 100644 --- a/packages/better-auth/src/db/get-tables.ts +++ b/packages/better-auth/src/db/get-tables.ts @@ -4,9 +4,22 @@ import type { BetterAuthOptions } from "../types"; export type BetterAuthDbSchema = Record< string, { + /** + * The name of the table in the database + */ tableName: string; + /** + * The fields of the table + */ fields: Record; + /** + * Whether to disable migrations for this table + * @default false + */ disableMigrations?: boolean; + /** + * The order of the table + */ order?: number; } >; @@ -57,8 +70,6 @@ export const getAuthTables = ( } satisfies BetterAuthDbSchema; const { user, session, account, ...pluginTables } = pluginSchema || {}; - const accountFields = options.account?.fields; - const userFields = options.user?.fields; return { user: { tableName: options.user?.modelName || "user", @@ -100,7 +111,7 @@ export const getAuthTables = ( ...user?.fields, ...options.user?.additionalFields, }, - order: 0, + order: 1, }, session: { tableName: options.session?.modelName || "session", @@ -133,7 +144,7 @@ export const getAuthTables = ( ...session?.fields, ...options.session?.additionalFields, }, - order: 1, + order: 2, }, account: { tableName: options.account?.modelName || "account", @@ -185,7 +196,7 @@ export const getAuthTables = ( }, ...account?.fields, }, - order: 2, + order: 3, }, verification: { tableName: options.verification?.modelName || "verification", @@ -206,6 +217,7 @@ export const getAuthTables = ( fieldName: options.verification?.fields?.expiresAt || "expiresAt", }, }, + order: 4, }, ...pluginTables, ...(shouldAddRateLimitTable ? rateLimitTable : {}), diff --git a/packages/better-auth/src/db/internal-adapter.ts b/packages/better-auth/src/db/internal-adapter.ts index 0e4a4cbf85..7d4d89c77c 100644 --- a/packages/better-auth/src/db/internal-adapter.ts +++ b/packages/better-auth/src/db/internal-adapter.ts @@ -128,22 +128,27 @@ export const createInternalAdapter = ( userAgent: headers?.get("user-agent") || "", ...override, }; - const session = await createWithHooks(data, "session"); - if (secondaryStorage && session) { + + if (secondaryStorage) { const user = await adapter.findOne({ model: tables.user.tableName, where: [{ field: "id", value: userId }], }); secondaryStorage.set( - session.id, + data.id, JSON.stringify({ - session, + session: data, user, }), sessionExpiration, ); + if (options.session?.storeSessionInDatabase) { + await createWithHooks(data, "session"); + } + } else { + await createWithHooks(data, "session"); } - return session; + return data; }, findSession: async (sessionId: string) => { if (secondaryStorage) { @@ -199,6 +204,40 @@ export const createInternalAdapter = ( }; }, updateSession: async (sessionId: string, session: Partial) => { + if (secondaryStorage) { + const currentSession = await secondaryStorage.get(sessionId); + let updatedSession: Session | null = null; + if (currentSession) { + const parsedSession = JSON.parse(currentSession) as { + session: Session; + user: User; + }; + updatedSession = { + ...parsedSession.session, + ...session, + }; + await secondaryStorage.set( + sessionId, + JSON.stringify({ + session: updatedSession, + user: parsedSession.user, + }), + parsedSession.session.expiresAt + ? new Date(parsedSession.session.expiresAt).getTime() + : undefined, + ); + } else { + return null; + } + if (options.session?.storeSessionInDatabase) { + await updateWithHooks( + session, + [{ field: "id", value: sessionId }], + "session", + ); + } + return updatedSession; + } const updatedSession = await updateWithHooks( session, [{ field: "id", value: sessionId }], @@ -209,6 +248,18 @@ export const createInternalAdapter = ( deleteSession: async (id: string) => { if (secondaryStorage) { await secondaryStorage.delete(id); + if (options.session?.storeSessionInDatabase) { + await adapter.delete({ + model: tables.session.tableName, + where: [ + { + field: "id", + value: id, + }, + ], + }); + } + return; } await adapter.delete({ model: tables.session.tableName, @@ -226,7 +277,7 @@ export const createInternalAdapter = ( model: tables.session.tableName, where: [ { - field: "userId", + field: tables.session.fields.userId.fieldName || "userId", value: userId, }, ], @@ -234,6 +285,18 @@ export const createInternalAdapter = ( for (const session of sessions) { await secondaryStorage.delete(session.id); } + if (options.session?.storeSessionInDatabase) { + await adapter.delete({ + model: tables.session.tableName, + where: [ + { + field: tables.session.fields.userId.fieldName || "userId", + value: userId, + }, + ], + }); + } + return; } await adapter.delete({ model: tables.session.tableName, diff --git a/packages/better-auth/src/integrations/next-js.ts b/packages/better-auth/src/integrations/next-js.ts index 9a3cab9c74..6995cc89da 100644 --- a/packages/better-auth/src/integrations/next-js.ts +++ b/packages/better-auth/src/integrations/next-js.ts @@ -18,8 +18,19 @@ export function toNextJsHandler(auth: Auth | Auth["handler"]) { * If not, it redirects to the redirectTo URL. */ export function authMiddleware(options: { + /** + * The base path of the auth API + * @default "/api/auth" + */ basePath?: string; + /** + * The URL to redirect to if the user is not authenticated + * @default "/" + */ redirectTo?: string; + /** + * A custom redirect function + */ customRedirect?: ( session: { user: User; @@ -29,9 +40,6 @@ export function authMiddleware(options: { ) => Promise; }) { return async (request: NextRequest) => { - // if (request.method !== "GET") { - // return NextResponse.next(); - // } const url = new URL(request.url).origin; const basePath = options?.basePath || "/api/auth"; const fullURL = `${url}${basePath}/session`; diff --git a/packages/better-auth/src/types/options.ts b/packages/better-auth/src/types/options.ts index 54b25bd255..947235e28c 100644 --- a/packages/better-auth/src/types/options.ts +++ b/packages/better-auth/src/types/options.ts @@ -192,6 +192,18 @@ export interface BetterAuthOptions { additionalFields?: { [key: string]: FieldAttribute; }; + /** + * By default if secondary storage is provided + * the session is stored in the secondary storage. + * + * Set this to true to store the session in the database + * as well. + * + * Reads are always done from the secondary storage. + * + * @default false + */ + storeSessionInDatabase?: boolean; }; account?: { modelName?: string;