feat: add tool aes

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
 )
 

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
+}

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.TimeoutCtx(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
 }

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))
+}

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
+	}
+}

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

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=

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
+}

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))
+}