first up

store/cache.go

@@ -0,0 +1,54 @@
+package store
+
+import (
+	"log"
+	"math/rand"
+
+	"github.com/Sakurasan/to"
+	"github.com/patrickmn/go-cache"
+)
+
+var (
+	KeysCache *cache.Cache
+	AuthCache *cache.Cache
+)
+
+func init() {
+	KeysCache = cache.New(cache.NoExpiration, cache.NoExpiration)
+	AuthCache = cache.New(cache.NoExpiration, cache.NoExpiration)
+}
+
+func LoadKeysCache() {
+	keys, err := GetAllKeys()
+	if err != nil {
+		log.Println(err)
+		return
+	}
+	for _, key := range keys {
+		KeysCache.Set(key.Key, true, cache.NoExpiration)
+	}
+}
+
+func FromKeyCacheRandomItem() string {
+	items := KeysCache.Items()
+	idx := rand.Intn(len(items))
+	item := items[to.String(idx)]
+	return item.Object.(string)
+}
+
+func LoadAuthCache() {
+	users, err := GetAllUsers()
+	if err != nil {
+		log.Println(err)
+		return
+	}
+	for _, user := range users {
+		AuthCache.Set(user.Token, true, cache.NoExpiration)
+	}
+}
+
+func IsExistAuthCache(auth string) bool {
+	items := AuthCache.Items()
+	_, ok := items[auth]
+	return ok
+}

store/db.go

@@ -0,0 +1,33 @@
+package store
+
+import (
+	"log"
+	"os"
+
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+)
+
+var db *gorm.DB
+
+func init() {
+	if _, err := os.Stat("db"); os.IsNotExist(err) {
+		errDir := os.MkdirAll("db", 0755)
+		if errDir != nil {
+			log.Fatalln("Error creating directory:", err)
+		}
+	}
+	var err error
+	db, err = gorm.Open(sqlite.Open("./db/cat.db"), &gorm.Config{})
+	if err != nil {
+		panic("failed to connect database")
+	}
+
+	// 自动迁移 User 结构体
+	err = db.AutoMigrate(&User{}, &Key{})
+	if err != nil {
+		panic(err)
+	}
+	LoadKeysCache()
+	LoadAuthCache()
+}

store/keydb.go

@@ -0,0 +1,64 @@
+package store
+
+import "time"
+
+type Key struct {
+	ID        uint      `gorm:"primarykey" json:"id,omitempty"`
+	Key       string    `gorm:"unique;not null" json:"key,omitempty"`
+	Name      string    `gorm:"unique;not null" json:"name,omitempty"`
+	UserId    string    `json:"-,omitempty"`
+	CreatedAt time.Time `json:"createdAt,omitempty"`
+	UpdatedAt time.Time `json:"updatedAt,omitempty"`
+}
+
+func GetKeyrByName(name string) (*Key, error) {
+	var key Key
+	result := db.First(&key, "name = ?", name)
+	if result.Error != nil {
+		return nil, result.Error
+	}
+	return &key, nil
+}
+
+func GetAllKeys() ([]Key, error) {
+	var keys []Key
+	if err := db.Find(&keys).Error; err != nil {
+		return nil, err
+	}
+	return keys, nil
+}
+
+// 添加记录
+func AddKey(apikey, name string) error {
+	key := Key{
+		Key:  apikey,
+		Name: name,
+	}
+	if err := db.Create(&key).Error; err != nil {
+		return err
+	}
+	LoadKeysCache()
+	return nil
+}
+
+// 删除记录
+func DeleteKey(id uint) error {
+	if err := db.Delete(&Key{}, id).Error; err != nil {
+		return err
+	}
+	LoadKeysCache()
+	return nil
+}
+
+// 更新记录
+func UpdateKey(id uint, apikey string, userId string) error {
+	key := Key{
+		Key:    apikey,
+		UserId: userId,
+	}
+	if err := db.Model(&Key{}).Where("id = ?", id).Updates(key).Error; err != nil {
+		return err
+	}
+	LoadKeysCache()
+	return nil
+}

store/userdb.go

@@ -0,0 +1,83 @@
+package store
+
+import (
+	"time"
+)
+
+type User struct {
+	IsDelete  bool      `gorm:"default:false" json:"IsDelete"`
+	ID        uint      `gorm:"primarykey" json:"id,omitempty"`
+	Name      string    `gorm:"unique;not null" json:"name,omitempty"`
+	Token     string    `gorm:"unique;not null" json:"token,omitempty"`
+	CreatedAt time.Time `json:"createdAt,omitempty"`
+	UpdatedAt time.Time `json:"updatedAt,omitempty"`
+	// DeletedAt gorm.DeletedAt `gorm:"index" json:"-"`
+}
+
+func CreateUser(u *User) error {
+	result := db.Create(u)
+	if result.Error != nil {
+		return result.Error
+	}
+	LoadAuthCache()
+	return nil
+}
+
+// 添加用户
+func AddUser(name, token string) error {
+	user := &User{Name: name, Token: token}
+	result := db.Create(&user)
+	if result.Error != nil {
+		return result.Error
+	}
+	LoadAuthCache()
+	return nil
+}
+
+// 删除用户
+func DeleteUser(id uint) error {
+	result := db.Delete(&User{}, id)
+	if result.Error != nil {
+		return result.Error
+	}
+	LoadAuthCache()
+	return nil
+}
+
+// 修改用户
+func UpdateUser(id uint, token string) error {
+	user := &User{Token: token}
+	result := db.Model(&User{}).Where("id = ?", id).Updates(user)
+	if result.Error != nil {
+		return result.Error
+	}
+	LoadAuthCache()
+	return nil
+}
+
+func GetUserByID(id uint) (*User, error) {
+	var user User
+	result := db.Where("id = ?", id).First(&user)
+	if result.Error != nil {
+		return nil, result.Error
+	}
+	return &user, nil
+}
+
+func GetUserByName(name string) (*User, error) {
+	var user User
+	result := db.Where(&User{Name: name}).First(&user)
+	if result.Error != nil {
+		return nil, result.Error
+	}
+	return &user, nil
+}
+
+func GetAllUsers() ([]*User, error) {
+	var users []*User
+	result := db.Find(&users)
+	if result.Error != nil {
+		return nil, result.Error
+	}
+	return users, nil
+}