From c6c2bdbe4aafc445d5be341083a9cc29d5123794 Mon Sep 17 00:00:00 2001
From: hailin <hailin@gdzx.xyz>
Date: Tue, 17 Jun 2025 22:45:20 +0800
Subject: [PATCH] .

---
 license/crypto.go | 22 +++++++++++++++++-----
 main.go           |  1 +
 storage/db.go     | 28 +++++++++++++++++++++++++++-
 3 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/license/crypto.go b/license/crypto.go
index 9dcdda2..a99ac0d 100644
--- a/license/crypto.go
+++ b/license/crypto.go
@@ -5,19 +5,31 @@ import (
 	"crypto/elliptic"
 	"crypto/rand"
 	"crypto/sha256"
+	"crypto/x509"
 	"encoding/asn1"
 	"encoding/base64"
+	"license-server/storage"
 	"math/big"
 )
 
 var privateKey *ecdsa.PrivateKey
 
-func init() {
-	var err error
-	privateKey, err = ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
-	if err != nil {
-		panic(err)
+// 显式从 main 传入 DB，保证先建库再读写
+func InitCrypto(db storage.Database) {
+	// 1. 读取
+	if pem, ok := db.GetPrivateKey(); ok {
+		if der, err := base64.StdEncoding.DecodeString(pem); err == nil {
+			if key, err := x509.ParseECPrivateKey(der); err == nil {
+				privateKey = key
+				return
+			}
+		}
 	}
+	// 2. 不存在或解析失败 → 重新生成并落库
+	key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	der, _ := x509.MarshalECPrivateKey(key)
+	db.SavePrivateKey(base64.StdEncoding.EncodeToString(der))
+	privateKey = key
 }
 
 // ecdsaSignature 是 ASN.1 编码中使用的结构体
diff --git a/main.go b/main.go
index a468d3a..6cce885 100644
--- a/main.go
+++ b/main.go
@@ -9,6 +9,7 @@ import (
 
 func main() {
 	db := storage.InitDB()
+	license.InitCrypto(db)
 	app := fiber.New()
 
 	app.Post("/api/license/generate", license.GenerateLicenseHandler(db))
diff --git a/storage/db.go b/storage/db.go
index ba1ce68..249fdb5 100644
--- a/storage/db.go
+++ b/storage/db.go
@@ -2,6 +2,8 @@ package storage
 
 import (
 	"database/sql"
+	"os"
+	"path/filepath"
 
 	_ "github.com/mattn/go-sqlite3"
 )
@@ -11,7 +13,13 @@ type Database struct {
 }
 
 func InitDB() Database {
-	db, _ := sql.Open("sqlite3", "./license.db")
+	const dbPath = "/root/database/license.db"
+
+	// 确保目录存在
+	_ = os.MkdirAll(filepath.Dir(dbPath), 0700)
+
+	db, _ := sql.Open("sqlite3", dbPath)
+
 	db.Exec(`CREATE TABLE IF NOT EXISTS activations (
         id INTEGER PRIMARY KEY AUTOINCREMENT,
         machine_id TEXT UNIQUE,
@@ -23,6 +31,13 @@ func InitDB() Database {
         machine_id TEXT UNIQUE,
         license TEXT
     )`)
+
+	/* ➊ 额外表：持久化私钥（仅一行） */
+	db.Exec(`CREATE TABLE IF NOT EXISTS private (
+        id INTEGER PRIMARY KEY,     
+        key TEXT                    
+    )`)
+
 	return Database{db}
 }
 
@@ -49,3 +64,14 @@ func (d Database) SaveGenerated(machineID string, licenseText string) {
 	d.db.Exec("INSERT OR REPLACE INTO generations(machine_id, license) VALUES (?, ?)",
 		machineID, licenseText)
 }
+
+/* ===== 新增两个工具方法 ===== */
+func (d Database) GetPrivateKey() (string, bool) {
+	row := d.db.QueryRow("SELECT key FROM private WHERE id = 1")
+	var key string
+	return key, row.Scan(&key) == nil
+}
+
+func (d Database) SavePrivateKey(key string) {
+	d.db.Exec("INSERT OR REPLACE INTO private(id, key) VALUES (1, ?)", key)
+}