H8-Network/TinySkin
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

抽出限流模块

Browse Source
This commit is contained in:
xmdhs 2023-09-03 23:26:28 +08:00
parent 94bb72eae4
commit 33427656c9
No known key found for this signature in database
GPG Key ID: E809D6D43DEFCC95
1 changed files with 24 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

40
service/yggdrasil/authenticate.go
View File

@ -11,6 +11,7 @@ import (
"github.com/golang-jwt/jwt/v5"
"github.com/google/uuid"
"github.com/xmdhs/authlib-skin/db/cache"
"github.com/xmdhs/authlib-skin/db/ent"
"github.com/xmdhs/authlib-skin/db/ent/user"
"github.com/xmdhs/authlib-skin/model"
@ -24,22 +25,7 @@ var (
)
func (y *Yggdrasil) Authenticate(cxt context.Context, auth yggdrasil.Authenticate) (yggdrasil.Token, error) {
key := []byte("Authenticate" + auth.Username)
v, err := y.cache.Get(key)
if err != nil {
return yggdrasil.Token{}, fmt.Errorf("Authenticate: %w", err)
}
if v != nil {
u := binary.BigEndian.Uint64(v)
t := time.Unix(int64(u), 0)
if time.Now().Before(t) {
return yggdrasil.Token{}, fmt.Errorf("Authenticate: %w", ErrRate)
}
}
b := make([]byte, 8)
binary.BigEndian.PutUint64(b, uint64(time.Now().Add(5*time.Second).Unix()))
err = y.cache.Put(key, b, time.Now().Add(20*time.Second))
err := rate("Authenticate"+auth.Username, y.cache, 5*time.Second)
if err != nil {
return yggdrasil.Token{}, fmt.Errorf("Authenticate: %w", err)
}
@ -116,3 +102,25 @@ func (y *Yggdrasil) Authenticate(cxt context.Context, auth yggdrasil.Authenticat
},
}, nil
}
func rate(k string, c cache.Cache, d time.Duration) error {
key := []byte(k)
v, err := c.Get([]byte(key))
if err != nil {
return fmt.Errorf("rate: %w", err)
}
if v != nil {
u := binary.BigEndian.Uint64(v)
t := time.Unix(int64(u), 0)
if time.Now().Before(t) {
return fmt.Errorf("rate: %w", ErrRate)
}
}
b := make([]byte, 8)
binary.BigEndian.PutUint64(b, uint64(time.Now().Add(d).Unix()))
err = c.Put(key, b, time.Now().Add(d))
if err != nil {
return fmt.Errorf("rate: %w", err)
}
return nil
}