1 | import { Context } from './context'
|
2 | import { MaybePromise } from './composer'
|
3 | import { MiddlewareFn } from './middleware'
|
4 |
|
5 | export interface SessionStore<T> {
|
6 | get: (name: string) => MaybePromise<T | undefined>
|
7 | set: (name: string, value: T) => MaybePromise<void>
|
8 | delete: (name: string) => MaybePromise<void>
|
9 | }
|
10 |
|
11 | interface SessionOptions<S extends object> {
|
12 | getSessionKey?: (ctx: Context) => Promise<string | undefined>
|
13 | store?: SessionStore<S>
|
14 | }
|
15 |
|
16 | export interface SessionContext<S extends object> extends Context {
|
17 | session?: S
|
18 | }
|
19 |
|
20 |
|
21 |
|
22 |
|
23 |
|
24 |
|
25 |
|
26 |
|
27 |
|
28 |
|
29 |
|
30 |
|
31 |
|
32 |
|
33 |
|
34 |
|
35 | export function session<S extends object>(
|
36 | options?: SessionOptions<S>
|
37 | ): MiddlewareFn<SessionContext<S>> {
|
38 | const getSessionKey = options?.getSessionKey ?? defaultGetSessionKey
|
39 | const store = options?.store ?? new MemorySessionStore()
|
40 | return async (ctx, next) => {
|
41 | const key = await getSessionKey(ctx)
|
42 | if (key == null) {
|
43 | return await next()
|
44 | }
|
45 | ctx.session = await store.get(key)
|
46 | await next()
|
47 | if (ctx.session == null) {
|
48 | await store.delete(key)
|
49 | } else {
|
50 | await store.set(key, ctx.session)
|
51 | }
|
52 | }
|
53 | }
|
54 |
|
55 | async function defaultGetSessionKey(ctx: Context): Promise<string | undefined> {
|
56 | const fromId = ctx.from?.id
|
57 | const chatId = ctx.chat?.id
|
58 | if (fromId == null || chatId == null) {
|
59 | return undefined
|
60 | }
|
61 | return `${fromId}:${chatId}`
|
62 | }
|
63 |
|
64 |
|
65 | export class MemorySessionStore<T> implements SessionStore<T> {
|
66 | private readonly store = new Map<string, { session: T; expires: number }>()
|
67 |
|
68 | constructor(private readonly ttl = Infinity) {}
|
69 |
|
70 | get(name: string): T | undefined {
|
71 | const entry = this.store.get(name)
|
72 | if (entry == null) {
|
73 | return undefined
|
74 | } else if (entry.expires < Date.now()) {
|
75 | this.delete(name)
|
76 | return undefined
|
77 | }
|
78 | return entry.session
|
79 | }
|
80 |
|
81 | set(name: string, value: T): void {
|
82 | const now = Date.now()
|
83 | this.store.set(name, { session: value, expires: now + this.ttl })
|
84 | }
|
85 |
|
86 | delete(name: string): void {
|
87 | this.store.delete(name)
|
88 | }
|
89 | }
|
90 |
|
91 | export function isSessionContext<S extends object>(
|
92 | ctx: Context
|
93 | ): ctx is SessionContext<S> {
|
94 | return 'session' in ctx
|
95 | }
|