feat: add tool aes

feat: add memory cache
refactor: rename timeout ctx
2025-07-16 18:51:52 +08:00 · 2025-07-15 18:40:13 +08:00 · 2025-07-15 11:43:15 +08:00 · 2025-07-15 11:18:54 +08:00 · 2025-07-15 10:45:05 +08:00 · 2025-07-14 11:16:34 +08:00
14 changed files with 419 additions and 57 deletions
--- a/api/api.go
+++ b/api/api.go
@ -23,7 +23,7 @@ type option struct {

 func WithName(name string) Option {
 	return func(o *option) {
-		if name == "" {
+		if name != "" {
 			o.name = name
 		}
 	}
@ -31,7 +31,7 @@ func WithName(name string) Option {

 func WithVersion(version string) Option {
 	return func(o *option) {
-		if version == "" {
+		if version != "" {
 			o.version = version
 		}
 	}
@ -39,7 +39,7 @@ func WithVersion(version string) Option {

 func WithAddress(address string) Option {
 	return func(o *option) {
-		if address == "" {
+		if address != "" {
 			o.address = address
 		}
 	}
--- a/database/cache/cache.go
+++ b/database/cache/cache.go
@ -62,6 +62,7 @@ var (
 	marshaler        func(data any) ([]byte, error)     = json.Marshal
 	unmarshaler      func(data []byte, model any) error = json.Unmarshal
 	ErrorKeyNotFound                                    = errors.New("key not found")
+	ErrorStoreFailed                                    = errors.New("store failed")
 	Default          Cache
 )

--- a/database/cache/memory.go
+++ b/database/cache/memory.go
@ -0,0 +1,155 @@
+package cache
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"github.com/dgraph-io/ristretto/v2"
+)
+
+var _ Cache = (*_mem)(nil)
+
+type _mem struct {
+	ctx   context.Context
+	cache *ristretto.Cache[string, []byte]
+}
+
+func newMemory(ctx context.Context, ins *ristretto.Cache[string, []byte]) Cache {
+	return &_mem{
+		ctx:   ctx,
+		cache: ins,
+	}
+}
+
+func (m *_mem) Client() any {
+	return m.cache
+}
+
+func (c *_mem) Close() {
+	c.cache.Close()
+}
+
+func (c *_mem) Del(ctx context.Context, keys ...string) error {
+	for _, key := range keys {
+		c.cache.Del(key)
+	}
+
+	return nil
+}
+
+func (c *_mem) Get(ctx context.Context, key string) ([]byte, error) {
+	val, ok := c.cache.Get(key)
+	if !ok {
+		return val, ErrorKeyNotFound
+	}
+
+	return val, nil
+}
+
+func (c *_mem) GetDel(ctx context.Context, key string) ([]byte, error) {
+	val, err := c.Get(ctx, key)
+	if err != nil {
+		return val, err
+	}
+
+	c.cache.Del(key)
+
+	return val, err
+}
+
+func (c *_mem) GetDelScan(ctx context.Context, key string) Scanner {
+	val, err := c.GetDel(ctx, key)
+	return newScan(val, err)
+}
+
+func (c *_mem) GetEx(ctx context.Context, key string, duration time.Duration) ([]byte, error) {
+	val, err := c.Get(ctx, key)
+	if err != nil {
+		return val, err
+	}
+
+	c.cache.SetWithTTL(key, val, 1, duration)
+
+	return val, err
+}
+
+func (m *_mem) GetExScan(ctx context.Context, key string, duration time.Duration) Scanner {
+	val, err := m.GetEx(ctx, key, duration)
+	return newScan(val, err)
+}
+
+func (m *_mem) GetScan(ctx context.Context, key string) Scanner {
+	val, err := m.Get(ctx, key)
+	return newScan(val, err)
+}
+
+func (m *_mem) Gets(ctx context.Context, keys ...string) ([][]byte, error) {
+	vals := make([][]byte, 0, len(keys))
+
+	for _, key := range keys {
+		val, err := m.Get(ctx, key)
+		if err != nil {
+			if errors.Is(err, ErrorKeyNotFound) {
+				continue
+			}
+
+			return vals, err
+		}
+
+		vals = append(vals, val)
+	}
+
+	if len(vals) != len(keys) {
+		return vals, ErrorKeyNotFound
+	}
+
+	return vals, nil
+}
+
+func (m *_mem) Set(ctx context.Context, key string, value any) error {
+	val, err := handleValue(value)
+	if err != nil {
+		return err
+	}
+
+	if ok := m.cache.Set(key, val, 1); !ok {
+		return ErrorStoreFailed
+	}
+
+	m.cache.Wait()
+
+	return nil
+}
+
+func (m *_mem) SetEx(ctx context.Context, key string, value any, duration time.Duration) error {
+	val, err := handleValue(value)
+	if err != nil {
+		return err
+	}
+
+	if ok := m.cache.SetWithTTL(key, val, 1, duration); !ok {
+		return ErrorStoreFailed
+	}
+
+	m.cache.Wait()
+
+	return nil
+}
+
+func (m *_mem) Sets(ctx context.Context, vm map[string]any) error {
+	for key, value := range vm {
+		val, err := handleValue(value)
+		if err != nil {
+			return err
+		}
+
+		if ok := m.cache.Set(key, val, 1); !ok {
+			return ErrorStoreFailed
+		}
+	}
+
+	m.cache.Wait()
+
+	return nil
+}
--- a/database/cache/new.go
+++ b/database/cache/new.go
@ -4,35 +4,31 @@ import (
 	"context"
 	"fmt"
 	"gitea.loveuer.com/yizhisec/packages/tool"
+	"github.com/dgraph-io/ristretto/v2"
 	"github.com/go-redis/redis/v8"
 	_ "github.com/go-redis/redis/v8"
 	"net/url"
 )

-var (
-	defaultRedis = "redis://127.0.0.1:6379"
-)
-
-func New(opts ...OptionFn) (Cache, error) {
+func New(opts ...Option) (Cache, error) {
 	var (
 		err error
-		cfg = &config{
-			ctx:   context.Background(),
-			redis: &defaultRedis,
+		opt = &option{
+			ctx: context.Background(),
 		}
 	)

-	for _, opt := range opts {
-		opt(cfg)
+	for _, fn := range opts {
+		fn(opt)
 	}

-	if cfg.redis != nil {
+	if opt.redis != nil {
 		var (
 			ins    *url.URL
 			client *redis.Client
 		)

-		if ins, err = url.Parse(*cfg.redis); err != nil {
+		if ins, err = url.Parse(*opt.redis); err != nil {
 			return nil, err
 		}

@ -45,17 +41,33 @@ func New(opts ...OptionFn) (Cache, error) {
 			Password: password,
 		})

-		if err = client.Ping(tool.CtxTimeout(cfg.ctx, 5)).Err(); err != nil {
+		if err = client.Ping(tool.TimeoutCtx(opt.ctx, 5)).Err(); err != nil {
 			return nil, err
 		}

-		return newRedis(cfg.ctx, client), nil
+		return newRedis(opt.ctx, client), nil
+	}
+
+	if opt.memory {
+		var (
+			ins *ristretto.Cache[string, []byte]
+		)
+
+		if ins, err = ristretto.NewCache(&ristretto.Config[string, []byte]{
+			NumCounters: 1e7,
+			MaxCost:     1 << 30,
+			BufferItems: 64,
+		}); err != nil {
+			return nil, err
+		}
+
+		return newMemory(opt.ctx, ins), nil
 	}

 	return nil, fmt.Errorf("invalid cache config")
 }

-func Init(opts ...OptionFn) (err error) {
+func Init(opts ...Option) (err error) {
 	Default, err = New(opts...)
 	return err
 }
--- a/database/cache/new_test.go
+++ b/database/cache/new_test.go
@ -1,6 +1,7 @@
 package cache

 import (
+	"errors"
 	"testing"
 )

@ -71,3 +72,31 @@ func TestNoAuth(t *testing.T) {
 	//	t.Fatal(err)
 	//}
 }
+
+func TestMemory(t *testing.T) {
+	c, err := New(WithMemory())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	bs, err := c.Get(t.Context(), "haha")
+	if err != nil {
+		if !errors.Is(err, ErrorKeyNotFound) {
+			t.Fatal(err)
+		}
+
+		t.Logf("key not found")
+	}
+
+	t.Logf("haha = %s", string(bs))
+
+	if err = c.Set(t.Context(), "haha", "haha"); err != nil {
+		t.Fatal(err)
+	}
+
+	if bs, err = c.Get(t.Context(), "haha"); err != nil {
+		t.Fatal(err)
+	}
+
+	t.Logf("haha = %s", string(bs))
+}
--- a/database/cache/option.go
+++ b/database/cache/option.go
@ -5,23 +5,24 @@ import (
 	"fmt"
 )

-type config struct {
-	ctx   context.Context
-	redis *string
+type option struct {
+	ctx    context.Context
+	redis  *string
+	memory bool
 }

-type OptionFn func(*config)
+type Option func(*option)

-func WithCtx(ctx context.Context) OptionFn {
-	return func(c *config) {
+func WithCtx(ctx context.Context) Option {
+	return func(c *option) {
 		if ctx != nil {
 			c.ctx = ctx
 		}
 	}
 }

-func WithRedis(host string, port int, username, password string) OptionFn {
-	return func(c *config) {
+func WithRedis(host string, port int, username, password string) Option {
+	return func(c *option) {
 		uri := fmt.Sprintf("redis://%s:%d", host, port)
 		if username != "" || password != "" {
 			uri = fmt.Sprintf("redis://%s:%s@%s:%d", username, password, host, port)
@ -30,3 +31,9 @@ func WithRedis(host string, port int, username, password string) OptionFn {
 		c.redis = &uri
 	}
 }
+
+func WithMemory() Option {
+	return func(c *option) {
+		c.memory = true
+	}
+}
--- a/go.mod
+++ b/go.mod
@ -1,10 +1,11 @@
 module gitea.loveuer.com/yizhisec/packages

-go 1.23
+go 1.23.0

 toolchain go1.24.3

 require (
+	github.com/dgraph-io/ristretto/v2 v2.2.0
 	github.com/fatih/color v1.18.0
 	github.com/gin-gonic/gin v1.10.1
 	github.com/glebarez/sqlite v1.11.0
@ -23,7 +24,7 @@ require (
 	filippo.io/edwards25519 v1.1.0 // indirect
 	github.com/bytedance/sonic v1.11.6 // indirect
 	github.com/bytedance/sonic/loader v0.1.1 // indirect
-	github.com/cespare/xxhash/v2 v2.1.1 // indirect
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/cloudwego/base64x v0.1.4 // indirect
 	github.com/cloudwego/iasm v0.2.0 // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
@ -60,7 +61,7 @@ require (
 	golang.org/x/arch v0.8.0 // indirect
 	golang.org/x/net v0.27.0 // indirect
 	golang.org/x/sync v0.11.0 // indirect
-	golang.org/x/sys v0.30.0 // indirect
+	golang.org/x/sys v0.31.0 // indirect
 	golang.org/x/text v0.22.0 // indirect
 	google.golang.org/protobuf v1.34.1 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
--- a/go.sum
+++ b/go.sum
@ -6,8 +6,9 @@ github.com/bytedance/sonic v1.11.6 h1:oUp34TzMlL+OY1OUWxHqsdkgC/Zfc85zGqw9siXjrc
 github.com/bytedance/sonic v1.11.6/go.mod h1:LysEHSvpvDySVdC2f87zGWf6CIKJcAvqab1ZaiQtds4=
 github.com/bytedance/sonic/loader v0.1.1 h1:c+e5Pt1k/cy5wMveRDyk2X4B9hF4g7an8N3zCYjJFNM=
 github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU=
-github.com/cespare/xxhash/v2 v2.1.1 h1:6MnRN8NT7+YBpUIWxHtefFZOKTAPgGjpQSxqLNn0+qY=
 github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
+github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cloudwego/base64x v0.1.4 h1:jwCgWpFanWmN8xoIUHa2rtzmkd5J2plF/dnLS6Xd/0Y=
 github.com/cloudwego/base64x v0.1.4/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w=
 github.com/cloudwego/iasm v0.2.0 h1:1KNIy1I1H9hNNFEEH3DVnI4UujN+1zjpuk6gwHLTssg=
@ -19,6 +20,10 @@ github.com/creack/pty v1.1.7/go.mod h1:lj5s0c3V2DBrqTV7llrYr5NG6My20zk30Fl46Y7Do
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dgraph-io/ristretto/v2 v2.2.0 h1:bkY3XzJcXoMuELV8F+vS8kzNgicwQFAaGINAEJdWGOM=
+github.com/dgraph-io/ristretto/v2 v2.2.0/go.mod h1:RZrm63UmcBAaYWC1DotLYBmTvgkrs0+XhBd7Npn7/zI=
+github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da h1:aIftn67I1fkbMa512G+w+Pxci9hJPB8oMnkcP3iZF38=
+github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
 github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
@ -321,8 +326,8 @@ golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc=
-golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik=
+golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
 golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
--- a/resp/msg.go
+++ b/resp/msg.go
@ -1,15 +1,16 @@
 package resp

 const (
-	Msg200        = "操作成功"
-	Msg400        = "参数错误"
-	Msg401        = "该账号登录已失效, 请重新登录"
-	Msg401NoMulti = "用户已在其他地方登录"
-	Msg403        = "权限不足"
-	Msg404        = "资源不存在"
-	Msg500        = "服务器开小差了"
-	Msg501        = "服务不可用"
-	Msg503        = "服务不可用或正在升级, 请联系管理员"
+	Msg200         = "操作成功"
+	Msg400         = "参数错误"
+	Msg401         = "该账号登录已失效, 请重新登录"
+	Msg401NoMulti  = "用户已在其他地方登录"
+	Msg401Inactive = "当前用户尚未生效, 请稍后再试"
+	Msg403         = "权限不足"
+	Msg404         = "资源不存在"
+	Msg500         = "服务器开小差了"
+	Msg501         = "服务不可用"
+	Msg503         = "服务不可用或正在升级, 请联系管理员"
 )

 func Msg(status int) string {
--- a/resp/resp.go
+++ b/resp/resp.go
@ -43,23 +43,45 @@ func RE(c *gin.Context, err error) {

 func _r(c *gin.Context, r *res, args ...any) {
 	length := len(args)
-	switch length {
-	case 0:
-		break
-	case 1:
-		if msg, ok := args[0].(string); ok {
-			r.Msg = msg
-		}
-	case 2:
-		r.Data = args[1]
-	case 3:
-		r.Err = args[2]
-	case 4:
-		if code, err := cast.ToIntE(args[3]); err == nil {
-			r.Code = code
-		}
+
+	if length == 0 {
+		goto END
 	}

+	if length >= 4 {
+		goto H4
+	}
+
+	if length >= 3 {
+		goto H3
+	}
+
+	if length >= 2 {
+		goto H2
+	}
+	if length >= 1 {
+		goto H1
+	}
+
+H4:
+	if code, err := cast.ToIntE(args[3]); err == nil {
+		r.Code = code
+	}
+H3:
+	if es, ok := args[2].(error); ok {
+		r.Err = es.Error()
+	} else {
+		r.Err = args[2]
+	}
+H2:
+	r.Data = args[1]
+H1:
+	if msg, ok := args[0].(string); ok {
+		r.Msg = msg
+	}
+
+END:
+
 	if r.Msg == "" {
 		r.Msg = Msg(r.Status)
 	}
--- a/resp/resp_test.go
+++ b/resp/resp_test.go
@ -0,0 +1,9 @@
+package resp
+
+import (
+	"testing"
+)
+
+func TestResp(t *testing.T) {
+
+}
--- a/tool/aes.go
+++ b/tool/aes.go
@ -0,0 +1,95 @@
+package tool
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"encoding/base64"
+	"errors"
+	"io"
+)
+
+// AES加密（CBC模式，PKCS7填充）
+func AesEncrypt(data []byte, key []byte) (string, error) {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return "", err
+	}
+
+	// 添加PKCS7填充
+	data = pkcs7Pad(data, block.BlockSize())
+
+	// 创建存储密文的buffer，前aes.BlockSize字节存储IV
+	ciphertext := make([]byte, aes.BlockSize+len(data))
+	iv := ciphertext[:aes.BlockSize]
+	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
+		return "", err
+	}
+
+	// CBC加密
+	mode := cipher.NewCBCEncrypter(block, iv)
+	mode.CryptBlocks(ciphertext[aes.BlockSize:], data)
+
+	return base64.StdEncoding.EncodeToString(ciphertext), nil
+}
+
+// AES解密（CBC模式，PKCS7填充）
+func AesDecrypt(encrypted string, key []byte) ([]byte, error) {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	ciphertext, err := base64.StdEncoding.DecodeString(encrypted)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(ciphertext) < aes.BlockSize {
+		return nil, errors.New("ciphertext too short")
+	}
+
+	iv := ciphertext[:aes.BlockSize]
+	ciphertext = ciphertext[aes.BlockSize:]
+
+	if len(ciphertext)%aes.BlockSize != 0 {
+		return nil, errors.New("ciphertext is not a multiple of the block size")
+	}
+
+	// CBC解密
+	mode := cipher.NewCBCDecrypter(block, iv)
+	mode.CryptBlocks(ciphertext, ciphertext)
+
+	// 去除PKCS7填充
+	return pkcs7Unpad(ciphertext)
+}
+
+// PKCS7填充
+func pkcs7Pad(data []byte, blockSize int) []byte {
+	padding := blockSize - (len(data) % blockSize)
+	padText := bytes.Repeat([]byte{byte(padding)}, padding)
+	return append(data, padText...)
+}
+
+// PKCS7去除填充
+func pkcs7Unpad(data []byte) ([]byte, error) {
+	length := len(data)
+	if length == 0 {
+		return nil, errors.New("empty input")
+	}
+
+	padding := int(data[length-1])
+	if padding > length || padding == 0 {
+		return nil, errors.New("invalid padding")
+	}
+
+	// 验证填充字节是否正确
+	for i := length - padding; i < length; i++ {
+		if int(data[i]) != padding {
+			return nil, errors.New("invalid padding")
+		}
+	}
+
+	return data[:length-padding], nil
+}
--- a/tool/aes_test.go
+++ b/tool/aes_test.go
@ -0,0 +1,25 @@
+package tool
+
+import (
+	"os"
+	"testing"
+)
+
+func TestAes(t *testing.T) {
+	key := os.Getenv("AES_KEY")
+
+	name := "admin"
+	res, err := AesEncrypt([]byte(name), []byte(key))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	t.Logf("res = %s", string(res))
+
+	raw, err := AesDecrypt(res, []byte(key))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	t.Logf("raw = %s", string(raw))
+}
--- a/tool/ctx.go
+++ b/tool/ctx.go
@ -23,7 +23,7 @@ func Timeout(seconds ...int) (ctx context.Context) {
 	return
 }

-func CtxTimeout(ctx context.Context, seconds ...int) context.Context {
+func TimeoutCtx(ctx context.Context, seconds ...int) context.Context {
 	var (
 		duration time.Duration
 	)
Author	SHA1	Message	Date
zhaoyupeng	104745e260	feat: add tool aes	2025-07-16 18:51:52 +08:00
zhaoyupeng	868b959c6f	feat: add memory cache	2025-07-15 18:40:13 +08:00
zhaoyupeng	127c57dc3a	refactor: rename timeout ctx	2025-07-15 11:43:15 +08:00
zhaoyupeng	d6b0b8ea36	fix: api opt setting	2025-07-15 11:18:54 +08:00
zhaoyupeng	23d7841ccf	feat: 添加提示语	2025-07-15 10:45:05 +08:00
zhaoyupeng	247490c35d	fix: resp 多参数问题	2025-07-14 11:16:34 +08:00