From bddbbb684dcbc76233b518816b4c5e381edac81a Mon Sep 17 00:00:00 2001
From: lifegpc <g1710431395@gmail.com>
Date: Sat, 15 Jul 2023 15:38:15 +0800
Subject: [PATCH] Add support for refresh token

---
 db.ts                     | 39 ++++++++++++++++++++++++++++++++++-----
 routes/api/_middleware.ts | 17 +++++++++++++++++
 routes/api/token.ts       |  2 +-
 3 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/db.ts b/db.ts
index 9f9dfcd..d6fdc46 100644
--- a/db.ts
+++ b/db.ts
@@ -152,12 +152,16 @@ export type Token = {
     uid: number;
     token: string;
     expired: Date;
+    http_only: boolean;
+    secure: boolean;
 };
 type TokenRaw = {
     id: number;
     uid: number;
     token: string;
     expired: string;
+    http_only: number;
+    secure: number;
 };
 const ALL_TABLES = [
     "version",
@@ -247,7 +251,9 @@ const TOKEN_TABLE = `CREATE TABLE token (
     id INTEGER PRIMARY KEY AUTOINCREMENT,
     uid INT,
     token TEXT,
-    expired TEXT
+    expired TEXT,
+    http_only BOOLEAN,
+    secure BOOLEAN
 );`;
 
 export class EhDb {
@@ -259,7 +265,7 @@ export class EhDb {
     #lock_file: string | undefined;
     #dblock_file: string | undefined;
     #_tags: Map<string, number> | undefined;
-    readonly version = parse_ver("1.0.0-8");
+    readonly version = parse_ver("1.0.0-9");
     constructor(base_path: string) {
         const db_path = join(base_path, "data.db");
         sure_dir_sync(base_path);
@@ -385,6 +391,11 @@ export class EhDb {
                 this.db.execute("DROP TABLE token;");
                 this.db.execute(TOKEN_TABLE);
             }
+            if (compare_ver(v, parse_ver("1.0.0-9")) === -1) {
+                this.db.execute("ALTER TABLE token ADD http_only BOOLEAN;");
+                this.db.execute("ALTER TABLE token ADD secure BOOLEAN;");
+                this.db.execute("UPDATE token SET http_only = 1, secure = 0;");
+            }
             this.#write_version();
             if (need_optimize) this.optimize();
         }
@@ -585,14 +596,19 @@ export class EhDb {
             )[0];
         });
     }
-    add_token(uid: number, added: number): Token {
+    add_token(
+        uid: number,
+        added: number,
+        http_only: boolean,
+        secure: boolean,
+    ): Token {
         let token = randomstring();
         while (this.get_token(token)) {
             token = randomstring();
         }
         this.db.query(
-            "INSERT INTO token (uid, token, expired) VALUES (?, ?, ?);",
-            [uid, token, new Date(added + 2592000000)],
+            "INSERT INTO token (uid, token, expired, http_only, secure) VALUES (?, ?, ?, ?, ?);",
+            [uid, token, new Date(added + 2592000000), http_only, secure],
         );
         const t = this.get_token(token);
         if (!t) throw Error("Failed to add token.");
@@ -776,8 +792,12 @@ export class EhDb {
     convert_token(m: TokenRaw[]) {
         return m.map((m) => {
             const e = new Date(m.expired);
+            const h = m.http_only !== 0;
+            const s = m.secure !== 0;
             const t = <Token> <unknown> m;
             t.expired = e;
+            t.http_only = h;
+            t.secure = s;
             return t;
         });
     }
@@ -1102,4 +1122,13 @@ export class EhDb {
             ]);
         });
     }
+    update_token(token: string, added: number): Token {
+        this.db.query(
+            "UPDATE token SET expired = ? WHERE token = ?;",
+            [new Date(added + 2592000000), token],
+        );
+        const t = this.get_token(token);
+        if (!t) throw Error("Failed to update token.");
+        return t;
+    }
 }
diff --git a/routes/api/_middleware.ts b/routes/api/_middleware.ts
index f7b04f7..e0f7a18 100644
--- a/routes/api/_middleware.ts
+++ b/routes/api/_middleware.ts
@@ -2,16 +2,19 @@ import { MiddlewareHandlerContext } from "$fresh/server.ts";
 import { get_task_manager } from "../../server.ts";
 import { parse_cookies } from "../../server/cookies.ts";
 import { return_error } from "../../server/utils.ts";
+import type { Token } from "../../db.ts";
 
 function handle_auth(req: Request, ctx: MiddlewareHandlerContext) {
     if (req.method === "OPTIONS") return true;
     const m = get_task_manager();
     if (m.db.get_user_count() === 0) return true;
     const u = new URL(req.url);
+    let is_from_cookie = false;
     let token: string | null | undefined = req.headers.get("X-TOKEN");
     const cookies = parse_cookies(req.headers.get("Cookie"));
     if (!token) {
         token = cookies.get("token");
+        is_from_cookie = true;
     }
     const check = () => {
         if (u.pathname === "/api/token" && req.method === "PUT") return true;
@@ -28,6 +31,11 @@ function handle_auth(req: Request, ctx: MiddlewareHandlerContext) {
         return check();
     }
     ctx.state.user = user;
+    if (is_from_cookie) {
+        if (t.expired.getTime() - 2505600000 < now) {
+            ctx.state.new_token = m.db.update_token(t.token, now);
+        }
+    }
     return true;
 }
 
@@ -54,6 +62,15 @@ export async function handler(req: Request, ctx: MiddlewareHandlerContext) {
         if (origin) {
             headers.set("Access-Control-Allow-Origin", "*");
         }
+        if (ctx.state.new_token) {
+            const t = <Token> ctx.state.new_token;
+            headers.append(
+                "Set-Cookie",
+                `token=${t.token}; Expires=${t.expired.toUTCString()}${
+                    t.http_only ? "; HttpOnly" : ""
+                }${t.secure ? "; Secure" : ""}`,
+            );
+        }
         return new Response(res.body, {
             status: res.status,
             headers: headers,
diff --git a/routes/api/token.ts b/routes/api/token.ts
index f67d76b..5b10b11 100644
--- a/routes/api/token.ts
+++ b/routes/api/token.ts
@@ -69,7 +69,7 @@ export const handler: Handlers = {
         if (!isEqual(pa, password)) {
             return return_error(4, USER_PASSWORD_ERROR);
         }
-        const token = m.db.add_token(u.id, now);
+        const token = m.db.add_token(u.id, now, http_only, secure);
         const headers: HeadersInit = {};
         if (set_cookie) {
             headers["Set-Cookie"] =