This commit is contained in:
parent
10ccc005be
commit
1ef8ecb7b2
|
|
@ -61,9 +61,16 @@ func (uc *AuthenticationUsecase) recordLogin(ctx context.Context, attempt *Login
|
|||
}
|
||||
|
||||
func (uc *AuthenticationUsecase) Login(ctx context.Context, attempt *LoginAttempt) (*AuthenticationResult, error) {
|
||||
config, _ := uc.security.Current(ctx)
|
||||
config, err := uc.security.Current(ctx)
|
||||
if err != nil || config == nil {
|
||||
return nil, fmt.Errorf("%w: security config unavailable", ErrLoginState)
|
||||
}
|
||||
if config != nil && config.LockEnable {
|
||||
if locked, _ := uc.security.LoginLocked(ctx, attempt.Username); locked {
|
||||
locked, lockErr := uc.security.LoginLocked(ctx, attempt.Username)
|
||||
if lockErr != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrLoginState, lockErr)
|
||||
}
|
||||
if locked {
|
||||
uc.recordLogin(ctx, attempt, false, "账号已锁定", 0)
|
||||
return nil, &AccountLockedError{Minutes: config.LockDuration}
|
||||
}
|
||||
|
|
@ -75,7 +82,10 @@ func (uc *AuthenticationUsecase) Login(ctx context.Context, attempt *LoginAttemp
|
|||
if config.CaptchaTimeout > 0 {
|
||||
ipTTL = time.Duration(config.CaptchaTimeout) * time.Second
|
||||
}
|
||||
failures, _ := uc.security.EnsureLoginIPCounter(ctx, attempt.IP, ipTTL)
|
||||
failures, counterErr := uc.security.EnsureLoginIPCounter(ctx, attempt.IP, ipTTL)
|
||||
if counterErr != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrLoginState, counterErr)
|
||||
}
|
||||
requireCaptcha = config.CaptchaOpen == 0 || failures > config.CaptchaOpen
|
||||
}
|
||||
if requireCaptcha && !uc.security.VerifyCaptcha(ctx, attempt.CaptchaID, attempt.Captcha, true) {
|
||||
|
|
@ -89,9 +99,14 @@ func (uc *AuthenticationUsecase) Login(ctx context.Context, attempt *LoginAttemp
|
|||
_, _ = uc.security.IncrementLoginIP(ctx, attempt.IP, ipTTL)
|
||||
if config != nil && config.LockEnable {
|
||||
lockTTL := time.Duration(config.LockDuration) * time.Minute
|
||||
failures, _ := uc.security.IncrementLoginFailure(ctx, attempt.Username, lockTTL)
|
||||
failures, stateErr := uc.security.IncrementLoginFailure(ctx, attempt.Username, lockTTL)
|
||||
if stateErr != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrLoginState, stateErr)
|
||||
}
|
||||
if int(failures) >= config.LockThreshold {
|
||||
_ = uc.security.LockLogin(ctx, attempt.Username, lockTTL)
|
||||
if stateErr = uc.security.LockLogin(ctx, attempt.Username, lockTTL); stateErr != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrLoginState, stateErr)
|
||||
}
|
||||
}
|
||||
}
|
||||
uc.recordLogin(ctx, attempt, false, "用户名不存在或者密码错误", 0)
|
||||
|
|
@ -130,6 +145,52 @@ func (uc *AuthenticationUsecase) Login(ctx context.Context, attempt *LoginAttemp
|
|||
return &AuthenticationResult{User: user, Token: issued.Value, ExpiresAt: issued.ExpiresAt, NeedChangePassword: user.MustChangePassword}, nil
|
||||
}
|
||||
|
||||
func userHasAuthority(user *User, authorityID uint) bool {
|
||||
if user == nil || authorityID == 0 {
|
||||
return false
|
||||
}
|
||||
if user.AuthorityID == authorityID {
|
||||
return true
|
||||
}
|
||||
for _, authority := range user.Authorities {
|
||||
if authority.AuthorityID == authorityID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (uc *AuthenticationUsecase) currentTokenUser(ctx context.Context, claims *AuthClaims) (*User, error) {
|
||||
if claims == nil || claims.UUID == "" || claims.IssuedAt.IsZero() {
|
||||
return nil, ErrTokenDisabled
|
||||
}
|
||||
user, err := uc.users.UserByUUID(ctx, claims.UUID)
|
||||
if err != nil || user == nil || user.ID != claims.ID || user.Username != claims.Username {
|
||||
return nil, ErrTokenDisabled
|
||||
}
|
||||
if user.Enable != 1 {
|
||||
return nil, &UserDisabledError{UserID: user.ID}
|
||||
}
|
||||
if !userHasAuthority(user, claims.AuthorityID) {
|
||||
return nil, ErrTokenDisabled
|
||||
}
|
||||
passwordVersion := int64(0)
|
||||
if user.PasswordUpdatedAt != nil {
|
||||
passwordVersion = user.PasswordUpdatedAt.UnixNano()
|
||||
}
|
||||
if passwordVersion != claims.PasswordVersion {
|
||||
return nil, ErrTokenDisabled
|
||||
}
|
||||
config, err := uc.security.Current(ctx)
|
||||
if err != nil || config == nil {
|
||||
return nil, ErrTokenDisabled
|
||||
}
|
||||
if config.PwdExpireEnable && config.PwdExpireDays > 0 && user.PasswordUpdatedAt != nil && time.Now().After(user.PasswordUpdatedAt.AddDate(0, 0, config.PwdExpireDays)) {
|
||||
user.MustChangePassword = true
|
||||
}
|
||||
return user, nil
|
||||
}
|
||||
|
||||
// SwitchAuthority changes the current user's active role and re-signs the
|
||||
// current claims. Role switching keeps the original JWT expiry; it
|
||||
// does not treat a role switch as a fresh login or rotate the multipoint
|
||||
|
|
@ -166,12 +227,17 @@ func (uc *AuthenticationUsecase) AuthenticateToken(ctx context.Context, token st
|
|||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
user, err := uc.currentTokenUser(ctx, claims)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
claims.NickName = user.NickName
|
||||
claims.MustChangePwd = user.MustChangePassword
|
||||
result := &TokenAuthentication{Claims: claims}
|
||||
if claims.BufferTime <= 0 || time.Until(claims.ExpiresAt) >= claims.BufferTime {
|
||||
return result, nil
|
||||
}
|
||||
user := &User{ID: claims.ID, UUID: claims.UUID, Username: claims.Username, NickName: claims.NickName, AuthorityID: claims.AuthorityID, MustChangePassword: claims.MustChangePwd}
|
||||
issued, err := uc.issuer.IssueToken(user, claims.AuthorityID, claims.MustChangePwd, 0)
|
||||
issued, err := uc.issuer.IssueToken(user, claims.AuthorityID, user.MustChangePassword, 0)
|
||||
if err != nil {
|
||||
return result, nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,29 +11,42 @@ import (
|
|||
|
||||
type authenticationUserRepo struct {
|
||||
UserRepo
|
||||
user *User
|
||||
user *User
|
||||
findByUUIDErr error
|
||||
}
|
||||
|
||||
func (r *authenticationUserRepo) FindUserByUsername(context.Context, string) (*User, error) {
|
||||
return r.user, nil
|
||||
}
|
||||
|
||||
func (r *authenticationUserRepo) FindUserByUUID(context.Context, string) (*User, error) {
|
||||
return r.user, r.findByUUIDErr
|
||||
}
|
||||
|
||||
func (*authenticationUserRepo) HasAuthorityMenu(context.Context, uint, string) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (*authenticationUserRepo) FillDepartmentNamePaths(context.Context, *User) error { return nil }
|
||||
|
||||
type authenticationSecurityRepo struct{ SecurityRepo }
|
||||
|
||||
func (*authenticationSecurityRepo) SecurityConfig(context.Context) (*SecurityConfig, error) {
|
||||
return &SecurityConfig{CaptchaOpen: 2, CaptchaTimeout: 60}, nil
|
||||
}
|
||||
|
||||
type authenticationCache struct{ activeErr error }
|
||||
type authenticationCache struct {
|
||||
activeErr error
|
||||
getErr error
|
||||
}
|
||||
|
||||
func (c *authenticationCache) Get(_ context.Context, key string) (string, bool, error) {
|
||||
if key == "admin" {
|
||||
return "", false, c.activeErr
|
||||
}
|
||||
if c.getErr != nil {
|
||||
return "", false, c.getErr
|
||||
}
|
||||
return "", false, nil
|
||||
}
|
||||
func (*authenticationCache) Set(context.Context, string, string, time.Duration) error { return nil }
|
||||
|
|
@ -46,12 +59,20 @@ type authenticationSettings struct{ RuntimeSettings }
|
|||
|
||||
func (*authenticationSettings) UseMultipoint() bool { return true }
|
||||
|
||||
type authenticationIssuer struct{ TokenIssuer }
|
||||
type authenticationIssuer struct {
|
||||
TokenIssuer
|
||||
claims *AuthClaims
|
||||
issuedUser *User
|
||||
issuedMust bool
|
||||
}
|
||||
|
||||
func (*authenticationIssuer) IssueToken(*User, uint, bool, time.Duration) (*IssuedToken, error) {
|
||||
func (i *authenticationIssuer) IssueToken(user *User, _ uint, must bool, _ time.Duration) (*IssuedToken, error) {
|
||||
i.issuedUser, i.issuedMust = user, must
|
||||
return &IssuedToken{Value: "token", ExpiresAt: time.Now().Add(time.Hour), TTL: time.Hour}, nil
|
||||
}
|
||||
|
||||
func (i *authenticationIssuer) ParseToken(string) (*AuthClaims, error) { return i.claims, nil }
|
||||
|
||||
func (*authenticationIssuer) ReissueToken(claims *AuthClaims, _ uint) (*IssuedToken, error) {
|
||||
return &IssuedToken{Value: "token", ExpiresAt: claims.ExpiresAt, TTL: time.Until(claims.ExpiresAt)}, nil
|
||||
}
|
||||
|
|
@ -93,6 +114,12 @@ func (a *authenticationAudit) RecordLogin(_ context.Context, value *LoginLog) er
|
|||
return nil
|
||||
}
|
||||
|
||||
type authenticationTokenRepo struct{ APITokenRepo }
|
||||
|
||||
func (*authenticationTokenRepo) IsTokenDisabled(context.Context, string) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func TestLoginRecordsSuccessBeforeMultipointCacheFailure(t *testing.T) {
|
||||
hash, err := bcrypt.GenerateFromPassword([]byte("secret"), bcrypt.MinCost)
|
||||
if err != nil {
|
||||
|
|
@ -137,3 +164,87 @@ func TestSwitchAuthorityReissuesCurrentClaimsWithoutReloadingUser(t *testing.T)
|
|||
t.Fatalf("switch result = %+v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoginRejectsSecurityCacheFailure(t *testing.T) {
|
||||
cacheErr := errors.New("cache unavailable")
|
||||
security := NewSecurityUsecase(&authenticationSecurityRepo{}, &authenticationCache{getErr: cacheErr}, &authenticationSettings{}, nil)
|
||||
uc := NewAuthenticationUsecase(NewUserUsecase(&authenticationUserRepo{}), security, &authenticationIssuer{}, &authenticationAudit{})
|
||||
|
||||
_, err := uc.Login(context.Background(), &LoginAttempt{Username: "admin", Password: "secret", IP: "127.0.0.1"})
|
||||
if !errors.Is(err, ErrLoginState) {
|
||||
t.Fatalf("expected login-state failure, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthenticateTokenRejectsStaleUserState(t *testing.T) {
|
||||
issuedAt := time.Now().Add(-time.Hour).Truncate(time.Second)
|
||||
baseClaims := AuthClaims{ID: 7, UUID: "uuid", Username: "admin", AuthorityID: 888, IssuedAt: issuedAt, ExpiresAt: time.Now().Add(time.Hour)}
|
||||
passwordChanged := issuedAt.Add(time.Minute)
|
||||
tests := []struct {
|
||||
name string
|
||||
user *User
|
||||
want error
|
||||
}{
|
||||
{name: "disabled", user: &User{ID: 7, UUID: "uuid", Username: "admin", AuthorityID: 888, Enable: 0}, want: ErrUserDisabled},
|
||||
{name: "role removed", user: &User{ID: 7, UUID: "uuid", Username: "admin", AuthorityID: 999, Enable: 1}, want: ErrTokenDisabled},
|
||||
{name: "password changed", user: &User{ID: 7, UUID: "uuid", Username: "admin", AuthorityID: 888, Enable: 1, PasswordUpdatedAt: &passwordChanged}, want: ErrTokenDisabled},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
claims := baseClaims
|
||||
issuer := &authenticationIssuer{claims: &claims}
|
||||
security := NewSecurityUsecase(&authenticationSecurityRepo{}, &authenticationCache{}, &authenticationSettings{}, NewTokenUsecase(&authenticationTokenRepo{}))
|
||||
uc := NewAuthenticationUsecase(NewUserUsecase(&authenticationUserRepo{user: test.user}), security, issuer, &authenticationAudit{})
|
||||
_, err := uc.AuthenticateToken(context.Background(), "token")
|
||||
if !errors.Is(err, test.want) {
|
||||
t.Fatalf("AuthenticateToken error = %v, want %v", err, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthenticateTokenRefreshesFromCurrentUser(t *testing.T) {
|
||||
issuedAt := time.Now().Add(-time.Hour).Truncate(time.Second)
|
||||
claims := &AuthClaims{ID: 7, UUID: "uuid", Username: "admin", NickName: "old", AuthorityID: 888, IssuedAt: issuedAt, BufferTime: 2 * time.Hour, ExpiresAt: time.Now().Add(time.Hour)}
|
||||
user := &User{ID: 7, UUID: "uuid", Username: "admin", NickName: "current", AuthorityID: 888, Enable: 1, MustChangePassword: true}
|
||||
issuer := &authenticationIssuer{claims: claims}
|
||||
security := NewSecurityUsecase(&authenticationSecurityRepo{}, &authenticationCache{}, &authenticationSettings{}, NewTokenUsecase(&authenticationTokenRepo{}))
|
||||
uc := NewAuthenticationUsecase(NewUserUsecase(&authenticationUserRepo{user: user}), security, issuer, &authenticationAudit{})
|
||||
|
||||
result, err := uc.AuthenticateToken(context.Background(), "token")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Refreshed == nil || result.Claims.NickName != "current" || !result.Claims.MustChangePwd {
|
||||
t.Fatalf("unexpected authentication result: %+v", result)
|
||||
}
|
||||
if issuer.issuedUser != user || !issuer.issuedMust {
|
||||
t.Fatalf("refresh used stale user state: user=%+v must=%v", issuer.issuedUser, issuer.issuedMust)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthenticateTokenRejectsPasswordChangeWithinIssuedAtSecond(t *testing.T) {
|
||||
issuedAt := time.Now().Add(-time.Minute).Truncate(time.Second)
|
||||
passwordAtIssue := issuedAt.Add(100 * time.Millisecond)
|
||||
passwordChanged := issuedAt.Add(900 * time.Millisecond)
|
||||
if passwordAtIssue.Unix() != passwordChanged.Unix() {
|
||||
t.Fatal("test setup did not keep password changes in one second")
|
||||
}
|
||||
claims := &AuthClaims{
|
||||
ID: 7,
|
||||
UUID: "uuid",
|
||||
Username: "admin",
|
||||
AuthorityID: 888,
|
||||
IssuedAt: issuedAt,
|
||||
ExpiresAt: time.Now().Add(time.Hour),
|
||||
PasswordVersion: passwordAtIssue.UnixNano(),
|
||||
}
|
||||
user := &User{ID: 7, UUID: "uuid", Username: "admin", AuthorityID: 888, Enable: 1, PasswordUpdatedAt: &passwordChanged}
|
||||
issuer := &authenticationIssuer{claims: claims}
|
||||
security := NewSecurityUsecase(&authenticationSecurityRepo{}, &authenticationCache{}, &authenticationSettings{}, NewTokenUsecase(&authenticationTokenRepo{}))
|
||||
uc := NewAuthenticationUsecase(NewUserUsecase(&authenticationUserRepo{user: user}), security, issuer, &authenticationAudit{})
|
||||
|
||||
if _, err := uc.AuthenticateToken(context.Background(), "token"); !errors.Is(err, ErrTokenDisabled) {
|
||||
t.Fatalf("AuthenticateToken error = %v, want ErrTokenDisabled", err)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -55,6 +55,15 @@ type MediaSettings struct {
|
|||
ChunkDir string
|
||||
}
|
||||
|
||||
const DefaultMaxMediaFileSize int64 = 100 << 20
|
||||
|
||||
func (s MediaSettings) EffectiveMaxFileSize() int64 {
|
||||
if s.MaxFileSize > 0 {
|
||||
return s.MaxFileSize
|
||||
}
|
||||
return DefaultMaxMediaFileSize
|
||||
}
|
||||
|
||||
// RuntimeSettings exposes only the active values needed by the application.
|
||||
// The data implementation resolves every call from conf.Runtime so hot reloads
|
||||
// take effect without rebuilding services.
|
||||
|
|
@ -73,18 +82,20 @@ type IssuedToken struct {
|
|||
}
|
||||
|
||||
type AuthClaims struct {
|
||||
UUID string
|
||||
ID uint
|
||||
Username string
|
||||
NickName string
|
||||
AuthorityID uint
|
||||
UserType string
|
||||
BufferTime time.Duration
|
||||
MustChangePwd bool
|
||||
Issuer string
|
||||
Audience []string
|
||||
NotBefore time.Time
|
||||
ExpiresAt time.Time
|
||||
UUID string
|
||||
ID uint
|
||||
Username string
|
||||
NickName string
|
||||
AuthorityID uint
|
||||
UserType string
|
||||
BufferTime time.Duration
|
||||
MustChangePwd bool
|
||||
PasswordVersion int64
|
||||
Issuer string
|
||||
Audience []string
|
||||
IssuedAt time.Time
|
||||
NotBefore time.Time
|
||||
ExpiresAt time.Time
|
||||
}
|
||||
|
||||
var (
|
||||
|
|
|
|||
|
|
@ -26,10 +26,19 @@ type MediaUsecase struct {
|
|||
settings RuntimeSettings
|
||||
}
|
||||
|
||||
var ErrMediaTooLarge = errors.New("文件超过大小上限")
|
||||
|
||||
func NewMediaUsecase(repo MediaRepo, files FileStorage, settings RuntimeSettings) *MediaUsecase {
|
||||
return &MediaUsecase{MediaRepo: repo, files: files, settings: settings}
|
||||
}
|
||||
|
||||
func (uc *MediaUsecase) maxMediaFileSize() int64 {
|
||||
if uc.settings == nil {
|
||||
return MediaSettings{}.EffectiveMaxFileSize()
|
||||
}
|
||||
return uc.settings.MediaSettings().EffectiveMaxFileSize()
|
||||
}
|
||||
|
||||
var allowedMediaExtensions = map[string]bool{
|
||||
".jpg": true, ".jpeg": true, ".png": true, ".gif": true, ".webp": true, ".bmp": true, ".ico": true, ".avif": true,
|
||||
".mp3": true, ".wav": true, ".ogg": true, ".m4a": true, ".flac": true, ".aac": true,
|
||||
|
|
@ -70,10 +79,25 @@ func (uc *MediaUsecase) Upload(ctx context.Context, userID uint, name, suppliedM
|
|||
}
|
||||
key := time.Now().Format("20060102") + "/" + uuid.NewString() + ext
|
||||
hash := md5.New()
|
||||
stored, err := uc.files.Put(ctx, key, io.TeeReader(buffered, hash))
|
||||
maxSize := uc.maxMediaFileSize()
|
||||
limited := &io.LimitedReader{R: buffered, N: maxSize + 1}
|
||||
stored, err := uc.files.Put(ctx, key, io.TeeReader(limited, hash))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
readSize := maxSize + 1 - limited.N
|
||||
if stored == nil {
|
||||
_ = uc.files.Delete(ctx, key)
|
||||
return nil, errors.New("文件存储未返回结果")
|
||||
}
|
||||
if readSize > maxSize || stored.Size > maxSize {
|
||||
_ = uc.files.Delete(ctx, key)
|
||||
return nil, ErrMediaTooLarge
|
||||
}
|
||||
if stored.Size != readSize {
|
||||
_ = uc.files.Delete(ctx, key)
|
||||
return nil, errors.New("文件存储大小不一致")
|
||||
}
|
||||
media := &MediaFile{Name: name, CategoryID: categoryID, URL: stored.URL, Tag: strings.TrimPrefix(filepath.Ext(name), "."), Key: key, Size: stored.Size, Mime: suppliedMIME, MD5: hex.EncodeToString(hash.Sum(nil)), UserID: userID}
|
||||
if save {
|
||||
count, countErr := uc.MediaKeyReferences(ctx, key)
|
||||
|
|
|
|||
|
|
@ -2,6 +2,9 @@ package biz
|
|||
|
||||
import (
|
||||
"context"
|
||||
"crypto/md5"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
|
|
@ -79,3 +82,169 @@ func TestListStorageFiltersChunksWithoutSkippingVisibleObjects(t *testing.T) {
|
|||
t.Fatalf("second cursor/more = %q/%v", cursor, more)
|
||||
}
|
||||
}
|
||||
|
||||
type mediaLimitSettings struct {
|
||||
RuntimeSettings
|
||||
max int64
|
||||
}
|
||||
|
||||
func (s mediaLimitSettings) MediaSettings() MediaSettings {
|
||||
return MediaSettings{MaxFileSize: s.max}
|
||||
}
|
||||
|
||||
type mediaLimitStorage struct {
|
||||
readSize int64
|
||||
reportedSize *int64
|
||||
deleted bool
|
||||
composeCalled bool
|
||||
composed *StoredFile
|
||||
composedHash string
|
||||
}
|
||||
|
||||
func (s *mediaLimitStorage) Put(_ context.Context, _ string, reader io.Reader) (*StoredFile, error) {
|
||||
data, err := io.ReadAll(reader)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
s.readSize = int64(len(data))
|
||||
size := s.readSize
|
||||
if s.reportedSize != nil {
|
||||
size = *s.reportedSize
|
||||
}
|
||||
return &StoredFile{Size: size}, nil
|
||||
}
|
||||
func (*mediaLimitStorage) Open(context.Context, string) (io.ReadCloser, error) {
|
||||
return io.NopCloser(strings.NewReader("")), nil
|
||||
}
|
||||
func (s *mediaLimitStorage) Delete(context.Context, string) error {
|
||||
s.deleted = true
|
||||
return nil
|
||||
}
|
||||
func (s *mediaLimitStorage) Compose(context.Context, []string, string) (*StoredFile, string, error) {
|
||||
s.composeCalled = true
|
||||
return s.composed, s.composedHash, nil
|
||||
}
|
||||
func (*mediaLimitStorage) DeletePrefix(context.Context, string) error { return nil }
|
||||
func (*mediaLimitStorage) List(context.Context, string, string, int) ([]*StoredFile, string, bool, error) {
|
||||
return nil, "", false, nil
|
||||
}
|
||||
|
||||
type mediaUploadRepo struct {
|
||||
MediaRepo
|
||||
session *UploadSession
|
||||
chunks []*UploadChunk
|
||||
claimed bool
|
||||
failed bool
|
||||
}
|
||||
|
||||
func (r *mediaUploadRepo) FindUploadSession(context.Context, uint) (*UploadSession, error) {
|
||||
if r.session == nil {
|
||||
return nil, ErrUploadSessionNotFound
|
||||
}
|
||||
return r.session, nil
|
||||
}
|
||||
func (r *mediaUploadRepo) ClaimUploadSession(context.Context, uint) (bool, error) {
|
||||
return r.claimed, nil
|
||||
}
|
||||
func (r *mediaUploadRepo) FailUploadSession(context.Context, uint) error {
|
||||
r.failed = true
|
||||
return nil
|
||||
}
|
||||
func (r *mediaUploadRepo) ListChunks(context.Context, uint) ([]*UploadChunk, error) {
|
||||
return r.chunks, nil
|
||||
}
|
||||
|
||||
func md5Text(value string) string {
|
||||
sum := md5.Sum([]byte(value))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func TestUploadRejectsDataPastConfiguredLimit(t *testing.T) {
|
||||
storage := &mediaLimitStorage{}
|
||||
uc := NewMediaUsecase(&mediaUploadRepo{}, storage, mediaLimitSettings{max: 4})
|
||||
|
||||
_, err := uc.Upload(context.Background(), 1, "sample.txt", "text/plain", 0, strings.NewReader("12345"), false)
|
||||
if !errors.Is(err, ErrMediaTooLarge) {
|
||||
t.Fatalf("Upload() error = %v, want ErrMediaTooLarge", err)
|
||||
}
|
||||
if storage.readSize != 5 || !storage.deleted {
|
||||
t.Fatalf("storage read/deleted = %d/%v, want 5/true", storage.readSize, storage.deleted)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadRejectsStorageSizeMismatch(t *testing.T) {
|
||||
reported := int64(3)
|
||||
storage := &mediaLimitStorage{reportedSize: &reported}
|
||||
uc := NewMediaUsecase(&mediaUploadRepo{}, storage, mediaLimitSettings{max: 4})
|
||||
|
||||
if _, err := uc.Upload(context.Background(), 1, "sample.txt", "text/plain", 0, strings.NewReader("1234"), false); err == nil {
|
||||
t.Fatal("Upload() error = nil, want storage size mismatch")
|
||||
}
|
||||
if !storage.deleted {
|
||||
t.Fatal("mismatched stored object was not deleted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestInitUploadRejectsInvalidLayout(t *testing.T) {
|
||||
uc := NewMediaUsecase(&mediaUploadRepo{}, &mediaLimitStorage{}, mediaLimitSettings{max: 10})
|
||||
validHash := strings.Repeat("0", md5.Size*2)
|
||||
tests := []struct {
|
||||
name string
|
||||
hash string
|
||||
size int64
|
||||
chunkSize int64
|
||||
total int
|
||||
wantLarge bool
|
||||
}{
|
||||
{name: "invalid hash", hash: "bad", size: 5, chunkSize: 3, total: 2},
|
||||
{name: "zero size", hash: validHash, size: 0, chunkSize: 3, total: 1},
|
||||
{name: "over limit", hash: validHash, size: 11, chunkSize: 3, total: 4, wantLarge: true},
|
||||
{name: "wrong total", hash: validHash, size: 5, chunkSize: 3, total: 3},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
_, _, _, err := uc.InitUpload(context.Background(), 1, "sample.txt", test.hash, test.size, test.chunkSize, test.total)
|
||||
if err == nil {
|
||||
t.Fatal("InitUpload() error = nil")
|
||||
}
|
||||
if test.wantLarge && !errors.Is(err, ErrMediaTooLarge) {
|
||||
t.Fatalf("InitUpload() error = %v, want ErrMediaTooLarge", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSaveChunkRejectsIndexAndSizeMismatch(t *testing.T) {
|
||||
repo := &mediaUploadRepo{session: &UploadSession{ID: 7, UserID: 1, FileSize: 5, ChunkSize: 3, ChunkTotal: 2, Status: "uploading"}}
|
||||
uc := NewMediaUsecase(repo, &mediaLimitStorage{}, mediaLimitSettings{max: 10})
|
||||
|
||||
if err := uc.SaveChunk(context.Background(), 1, 7, 2, md5Text("abc"), strings.NewReader("abc")); err == nil {
|
||||
t.Fatal("SaveChunk() accepted an out-of-range index")
|
||||
}
|
||||
if err := uc.SaveChunk(context.Background(), 1, 7, 0, md5Text("ab"), strings.NewReader("ab")); err == nil {
|
||||
t.Fatal("SaveChunk() accepted a short chunk")
|
||||
}
|
||||
if err := uc.SaveChunk(context.Background(), 1, 7, 0, md5Text("abcd"), strings.NewReader("abcd")); err == nil {
|
||||
t.Fatal("SaveChunk() accepted an oversized chunk")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompleteUploadRejectsMismatchedChunkSizesBeforeCompose(t *testing.T) {
|
||||
repo := &mediaUploadRepo{
|
||||
session: &UploadSession{ID: 7, UserID: 1, FileName: "sample.txt", FileHash: md5Text("abcde"), FileSize: 5, ChunkSize: 3, ChunkTotal: 2, Status: "uploading"},
|
||||
chunks: []*UploadChunk{{Index: 0, Size: 3}, {Index: 1, Size: 1}},
|
||||
claimed: true,
|
||||
}
|
||||
storage := &mediaLimitStorage{}
|
||||
uc := NewMediaUsecase(repo, storage, mediaLimitSettings{max: 10})
|
||||
|
||||
if _, err := uc.CompleteUpload(context.Background(), 1, 7, "text/plain"); err == nil {
|
||||
t.Fatal("CompleteUpload() accepted mismatched chunk sizes")
|
||||
}
|
||||
if storage.composeCalled {
|
||||
t.Fatal("Compose() was called before chunk-size validation completed")
|
||||
}
|
||||
if !repo.failed {
|
||||
t.Fatal("failed upload session was not returned to the uploading state")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -19,6 +19,45 @@ import (
|
|||
|
||||
var ErrUploadSessionNotFound = errors.New("upload session not found")
|
||||
|
||||
func validMD5(value string) bool {
|
||||
if len(value) != md5.Size*2 {
|
||||
return false
|
||||
}
|
||||
_, err := hex.DecodeString(value)
|
||||
return err == nil
|
||||
}
|
||||
|
||||
func (uc *MediaUsecase) validateUploadLayout(size, chunkSize int64, total int) error {
|
||||
if size <= 0 || chunkSize <= 0 || total <= 0 {
|
||||
return errors.New("文件大小和分片参数必须大于 0")
|
||||
}
|
||||
if size > uc.maxMediaFileSize() {
|
||||
return ErrMediaTooLarge
|
||||
}
|
||||
expectedTotal := (size-1)/chunkSize + 1
|
||||
if int64(total) != expectedTotal {
|
||||
return fmt.Errorf("分片数量不匹配: %d/%d", total, expectedTotal)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func expectedChunkSize(session *UploadSession, index int) (int64, error) {
|
||||
if session == nil || session.FileSize <= 0 || session.ChunkSize <= 0 || session.ChunkTotal <= 0 {
|
||||
return 0, errors.New("上传会话分片参数非法")
|
||||
}
|
||||
if index < 0 || index >= session.ChunkTotal {
|
||||
return 0, errors.New("分片序号超出范围")
|
||||
}
|
||||
if index < session.ChunkTotal-1 {
|
||||
return session.ChunkSize, nil
|
||||
}
|
||||
last := session.FileSize - session.ChunkSize*int64(session.ChunkTotal-1)
|
||||
if last <= 0 || last > session.ChunkSize {
|
||||
return 0, errors.New("上传会话分片参数非法")
|
||||
}
|
||||
return last, nil
|
||||
}
|
||||
|
||||
func (uc *MediaUsecase) chunkPrefix(uploadID uint) string {
|
||||
directory := "uploads/chunks"
|
||||
if uc.settings != nil {
|
||||
|
|
@ -37,14 +76,24 @@ func (uc *MediaUsecase) InitUpload(ctx context.Context, userID uint, name, hash
|
|||
if err := validateMediaName(name); err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
if completed, err := uc.FindCompletedSession(ctx, userID, hash); err == nil && completed.MediaID != 0 {
|
||||
if !validMD5(hash) {
|
||||
return nil, nil, nil, errors.New("文件 MD5 非法")
|
||||
}
|
||||
if err := uc.validateUploadLayout(size, chunkSize, total); err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
completed, err := uc.FindCompletedSession(ctx, userID, strings.ToLower(hash))
|
||||
if err == nil && completed.MediaID != 0 {
|
||||
if media, findErr := uc.FindMedia(ctx, completed.MediaID); findErr == nil {
|
||||
copy := &MediaFile{Name: name, URL: media.URL, Tag: media.Tag, Key: media.Key}
|
||||
if createErr := uc.CreateMedia(ctx, copy); createErr == nil {
|
||||
return nil, copy, nil, nil
|
||||
}
|
||||
}
|
||||
} else if err != nil && !errors.Is(err, ErrUploadSessionNotFound) {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
hash = strings.ToLower(hash)
|
||||
session, err := uc.FindUploadingSession(ctx, userID, hash)
|
||||
if errors.Is(err, ErrUploadSessionNotFound) {
|
||||
session = &UploadSession{UserID: userID, FileName: name, FileHash: hash, FileSize: size, ChunkSize: chunkSize, ChunkTotal: total, Status: "uploading"}
|
||||
|
|
@ -53,6 +102,8 @@ func (uc *MediaUsecase) InitUpload(ctx context.Context, userID uint, name, hash
|
|||
}
|
||||
} else if err != nil {
|
||||
return nil, nil, nil, err
|
||||
} else if session.FileSize != size || session.ChunkSize != chunkSize || session.ChunkTotal != total {
|
||||
return nil, nil, nil, errors.New("上传参数与已有会话不一致")
|
||||
}
|
||||
chunks, err := uc.ListChunks(ctx, session.ID)
|
||||
if err != nil {
|
||||
|
|
@ -76,6 +127,14 @@ func (uc *MediaUsecase) SaveChunk(ctx context.Context, userID, uploadID uint, in
|
|||
if session.Status != "uploading" {
|
||||
return errors.New("上传会话状态不允许收片")
|
||||
}
|
||||
expected = strings.ToLower(strings.TrimSpace(expected))
|
||||
if !validMD5(expected) {
|
||||
return errors.New("分片 MD5 非法")
|
||||
}
|
||||
wantSize, err := expectedChunkSize(session, index)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
hash := md5.New()
|
||||
key := uc.chunkKey(uploadID, index)
|
||||
temporary, err := os.CreateTemp("", "kra-upload-chunk-*")
|
||||
|
|
@ -85,11 +144,15 @@ func (uc *MediaUsecase) SaveChunk(ctx context.Context, userID, uploadID uint, in
|
|||
temporaryName := temporary.Name()
|
||||
defer os.Remove(temporaryName)
|
||||
defer temporary.Close()
|
||||
if _, err = io.Copy(io.MultiWriter(temporary, hash), reader); err != nil {
|
||||
written, err := io.Copy(io.MultiWriter(temporary, hash), io.LimitReader(reader, wantSize+1))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if written != wantSize {
|
||||
return fmt.Errorf("分片 %d 大小不匹配: %d/%d", index, written, wantSize)
|
||||
}
|
||||
actual := hex.EncodeToString(hash.Sum(nil))
|
||||
if actual != expected {
|
||||
if !strings.EqualFold(actual, expected) {
|
||||
return fmt.Errorf("分片 %d 校验失败", index)
|
||||
}
|
||||
if _, err = temporary.Seek(0, io.SeekStart); err != nil {
|
||||
|
|
@ -99,7 +162,15 @@ func (uc *MediaUsecase) SaveChunk(ctx context.Context, userID, uploadID uint, in
|
|||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return uc.UpsertChunk(ctx, uploadID, &UploadChunk{Index: index, Hash: actual, Size: stored.Size})
|
||||
if stored == nil || stored.Size != written {
|
||||
_ = uc.files.Delete(ctx, key)
|
||||
return errors.New("分片存储大小不一致")
|
||||
}
|
||||
if err = uc.UpsertChunk(ctx, uploadID, &UploadChunk{Index: index, Hash: actual, Size: written}); err != nil {
|
||||
_ = uc.files.Delete(ctx, key)
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (uc *MediaUsecase) CompleteUpload(ctx context.Context, userID, uploadID uint, mime string) (*MediaFile, error) {
|
||||
session, err := uc.FindUploadSession(ctx, uploadID)
|
||||
|
|
@ -129,26 +200,38 @@ func (uc *MediaUsecase) CompleteUpload(ctx context.Context, userID, uploadID uin
|
|||
}
|
||||
sort.Slice(chunks, func(i, j int) bool { return chunks[i].Index < chunks[j].Index })
|
||||
names := make([]string, 0, len(chunks))
|
||||
var totalSize int64
|
||||
for index, chunk := range chunks {
|
||||
if chunk.Index != index {
|
||||
return fail(errors.New("分片序号不连续"))
|
||||
}
|
||||
expectedSize, sizeErr := expectedChunkSize(session, index)
|
||||
if sizeErr != nil || chunk.Size != expectedSize || chunk.Size > session.FileSize-totalSize {
|
||||
return fail(errors.New("分片大小不一致"))
|
||||
}
|
||||
totalSize += chunk.Size
|
||||
names = append(names, uc.chunkKey(uploadID, index))
|
||||
}
|
||||
if totalSize != session.FileSize {
|
||||
return fail(fmt.Errorf("文件大小不匹配: %d/%d", totalSize, session.FileSize))
|
||||
}
|
||||
ext := strings.ToLower(filepath.Ext(session.FileName))
|
||||
key := time.Now().Format("20060102") + "/" + uuid.NewString() + ext
|
||||
stored, hash, err := uc.files.Compose(ctx, names, key)
|
||||
if err != nil {
|
||||
return fail(err)
|
||||
}
|
||||
if hash != session.FileHash {
|
||||
if stored == nil || stored.Size != session.FileSize {
|
||||
if stored != nil {
|
||||
_ = uc.files.Delete(ctx, key)
|
||||
}
|
||||
return fail(errors.New("合并文件大小不匹配"))
|
||||
}
|
||||
if !strings.EqualFold(hash, session.FileHash) {
|
||||
_ = uc.files.Delete(ctx, key)
|
||||
return fail(errors.New("整文件校验失败"))
|
||||
}
|
||||
// Keep the size declared when the upload session was created. The full-file
|
||||
// MD5 above remains the integrity check even if storage reports another
|
||||
// byte count.
|
||||
media := &MediaFile{Name: session.FileName, URL: stored.URL, Tag: strings.TrimPrefix(filepath.Ext(session.FileName), "."), Key: key, Size: session.FileSize, Mime: mime, MD5: hash, UserID: userID}
|
||||
media := &MediaFile{Name: session.FileName, URL: stored.URL, Tag: strings.TrimPrefix(filepath.Ext(session.FileName), "."), Key: key, Size: stored.Size, Mime: mime, MD5: hash, UserID: userID}
|
||||
if err = uc.CreateMedia(ctx, media); err != nil {
|
||||
return fail(err)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -154,39 +154,47 @@ func NewTaskApplicationUsecase(tasks *TaskUsecase, runtime TaskRuntime) *TaskApp
|
|||
return &TaskApplicationUsecase{tasks: tasks, runtime: runtime}
|
||||
}
|
||||
|
||||
func (uc *TaskApplicationUsecase) syncRuntime(ctx context.Context, id uint) error {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
if err := uc.runtime.ScheduleID(ctx, id); err != nil {
|
||||
repairCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
|
||||
defer cancel()
|
||||
if reloadErr := uc.runtime.Reload(repairCtx); reloadErr != nil {
|
||||
return &TaskScheduleError{Err: errors.Join(err, fmt.Errorf("重载任务运行时失败: %w", reloadErr))}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (uc *TaskApplicationUsecase) Create(ctx context.Context, value *TimedTask) error {
|
||||
if err := uc.tasks.Create(ctx, value); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := uc.runtime.ScheduleID(ctx, value.ID); err != nil {
|
||||
return &TaskScheduleError{Err: err}
|
||||
}
|
||||
return nil
|
||||
return uc.syncRuntime(ctx, value.ID)
|
||||
}
|
||||
|
||||
func (uc *TaskApplicationUsecase) Update(ctx context.Context, value *TimedTask) error {
|
||||
if err := uc.tasks.Update(ctx, value); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := uc.runtime.ScheduleID(ctx, value.ID); err != nil {
|
||||
return &TaskScheduleError{Err: err}
|
||||
}
|
||||
return nil
|
||||
return uc.syncRuntime(ctx, value.ID)
|
||||
}
|
||||
|
||||
func (uc *TaskApplicationUsecase) Delete(ctx context.Context, id uint) error {
|
||||
if err := uc.tasks.DeleteTask(ctx, id); err != nil {
|
||||
return err
|
||||
}
|
||||
uc.runtime.Remove(id)
|
||||
return uc.tasks.DeleteTask(ctx, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (uc *TaskApplicationUsecase) Toggle(ctx context.Context, id uint, enabled bool) error {
|
||||
if err := uc.tasks.ToggleTask(ctx, id, enabled); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := uc.runtime.ScheduleID(ctx, id); err != nil {
|
||||
return &TaskScheduleError{Err: err}
|
||||
}
|
||||
return nil
|
||||
return uc.syncRuntime(ctx, id)
|
||||
}
|
||||
|
||||
func (uc *TaskApplicationUsecase) Trigger(ctx context.Context, id uint) error {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,197 @@
|
|||
package biz
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type applicationTaskRepo struct {
|
||||
mu sync.Mutex
|
||||
deleted bool
|
||||
deleteErr error
|
||||
createdID uint
|
||||
toggleCalls int
|
||||
updateCalls int
|
||||
createdCalls int
|
||||
nameExists bool
|
||||
nameExistsErr error
|
||||
deleteObserved chan struct{}
|
||||
}
|
||||
|
||||
func (r *applicationTaskRepo) CreateTask(_ context.Context, value *TimedTask) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.createdCalls++
|
||||
if value.ID == 0 {
|
||||
value.ID = r.createdID
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (r *applicationTaskRepo) UpdateTask(context.Context, *TimedTask) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.updateCalls++
|
||||
return nil
|
||||
}
|
||||
func (r *applicationTaskRepo) DeleteTask(context.Context, uint) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.deleteErr != nil {
|
||||
return r.deleteErr
|
||||
}
|
||||
r.deleted = true
|
||||
if r.deleteObserved != nil {
|
||||
close(r.deleteObserved)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (r *applicationTaskRepo) FindTask(context.Context, uint) (*TimedTask, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
func (r *applicationTaskRepo) ListTasks(context.Context, int, int, *TimedTask) ([]*TimedTask, int64, error) {
|
||||
return nil, 0, nil
|
||||
}
|
||||
func (r *applicationTaskRepo) ToggleTask(context.Context, uint, bool) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.toggleCalls++
|
||||
return nil
|
||||
}
|
||||
func (r *applicationTaskRepo) RecordTaskLog(context.Context, *TimedTaskLog) error { return nil }
|
||||
func (r *applicationTaskRepo) ListTaskLogs(context.Context, int, int, uint, string) ([]*TimedTaskLog, int64, error) {
|
||||
return nil, 0, nil
|
||||
}
|
||||
func (r *applicationTaskRepo) CleanupLogs(context.Context) error { return nil }
|
||||
func (r *applicationTaskRepo) TaskNameExists(context.Context, string, uint) (bool, error) {
|
||||
return r.nameExists, r.nameExistsErr
|
||||
}
|
||||
|
||||
type applicationTaskRuntime struct {
|
||||
mu sync.Mutex
|
||||
scheduleErr error
|
||||
reloadErr error
|
||||
scheduleCalls int
|
||||
reloadCalls int
|
||||
removeCalls int
|
||||
removeBeforeDelete bool
|
||||
repo *applicationTaskRepo
|
||||
reloadContextActive bool
|
||||
}
|
||||
|
||||
func (r *applicationTaskRuntime) ScheduleID(context.Context, uint) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.scheduleCalls++
|
||||
return r.scheduleErr
|
||||
}
|
||||
func (r *applicationTaskRuntime) Remove(uint) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.removeCalls++
|
||||
if r.repo != nil {
|
||||
r.repo.mu.Lock()
|
||||
r.removeBeforeDelete = !r.repo.deleted
|
||||
r.repo.mu.Unlock()
|
||||
}
|
||||
}
|
||||
func (r *applicationTaskRuntime) TriggerID(context.Context, uint) error { return nil }
|
||||
func (r *applicationTaskRuntime) NextRuns() map[uint]time.Time { return nil }
|
||||
func (r *applicationTaskRuntime) Reload(ctx context.Context) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.reloadCalls++
|
||||
r.reloadContextActive = ctx.Err() == nil
|
||||
return r.reloadErr
|
||||
}
|
||||
func (r *applicationTaskRuntime) Subscribe(uint) chan []byte { return make(chan []byte) }
|
||||
func (r *applicationTaskRuntime) Unsubscribe(uint, chan []byte) {}
|
||||
|
||||
func validApplicationTask(id uint) *TimedTask {
|
||||
const methodName = "biz-test-application-task"
|
||||
RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error { return nil })
|
||||
return &TimedTask{ID: id, Name: "test", Spec: "0 0 * * *", ExecutorType: TaskExecutorMethod, MethodName: methodName, Enabled: true}
|
||||
}
|
||||
|
||||
func TestTaskApplicationRepairsRuntimeAfterScheduleFailure(t *testing.T) {
|
||||
repo := &applicationTaskRepo{createdID: 41}
|
||||
runtime := &applicationTaskRuntime{scheduleErr: errors.New("schedule failed")}
|
||||
uc := NewTaskApplicationUsecase(NewTaskUsecase(repo), runtime)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
if err := uc.Create(ctx, validApplicationTask(0)); err != nil {
|
||||
t.Fatalf("Create() error = %v, want successful reload compensation", err)
|
||||
}
|
||||
if runtime.scheduleCalls != 1 || runtime.reloadCalls != 1 {
|
||||
t.Fatalf("runtime calls = schedule:%d reload:%d, want 1 and 1", runtime.scheduleCalls, runtime.reloadCalls)
|
||||
}
|
||||
if !runtime.reloadContextActive {
|
||||
t.Fatal("reload compensation inherited the canceled request context")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskApplicationReturnsScheduleErrorWhenRepairFails(t *testing.T) {
|
||||
repo := &applicationTaskRepo{createdID: 42}
|
||||
runtime := &applicationTaskRuntime{scheduleErr: errors.New("schedule failed"), reloadErr: errors.New("reload failed")}
|
||||
uc := NewTaskApplicationUsecase(NewTaskUsecase(repo), runtime)
|
||||
|
||||
err := uc.Create(context.Background(), validApplicationTask(0))
|
||||
if !errors.Is(err, ErrTaskSchedule) {
|
||||
t.Fatalf("Create() error = %v, want ErrTaskSchedule", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskApplicationUpdateAndToggleUseRuntimeRepair(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
run func(*TaskApplicationUsecase) error
|
||||
}{
|
||||
{name: "update", run: func(uc *TaskApplicationUsecase) error {
|
||||
return uc.Update(context.Background(), validApplicationTask(7))
|
||||
}},
|
||||
{name: "toggle", run: func(uc *TaskApplicationUsecase) error { return uc.Toggle(context.Background(), 7, true) }},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
repo := &applicationTaskRepo{}
|
||||
runtime := &applicationTaskRuntime{scheduleErr: errors.New("schedule failed")}
|
||||
uc := NewTaskApplicationUsecase(NewTaskUsecase(repo), runtime)
|
||||
if err := test.run(uc); err != nil {
|
||||
t.Fatalf("operation error = %v", err)
|
||||
}
|
||||
if runtime.scheduleCalls != 1 || runtime.reloadCalls != 1 {
|
||||
t.Fatalf("runtime calls = schedule:%d reload:%d, want 1 and 1", runtime.scheduleCalls, runtime.reloadCalls)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskApplicationDeleteRemovesRuntimeAfterDatabase(t *testing.T) {
|
||||
repo := &applicationTaskRepo{}
|
||||
runtime := &applicationTaskRuntime{repo: repo}
|
||||
uc := NewTaskApplicationUsecase(NewTaskUsecase(repo), runtime)
|
||||
if err := uc.Delete(context.Background(), 8); err != nil {
|
||||
t.Fatalf("Delete() error = %v", err)
|
||||
}
|
||||
if runtime.removeBeforeDelete {
|
||||
t.Fatal("runtime entry was removed before the database row")
|
||||
}
|
||||
if runtime.removeCalls != 1 {
|
||||
t.Fatalf("Remove() calls = %d, want 1", runtime.removeCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskApplicationDeleteKeepsRuntimeWhenDatabaseFails(t *testing.T) {
|
||||
repo := &applicationTaskRepo{deleteErr: errors.New("delete failed")}
|
||||
runtime := &applicationTaskRuntime{repo: repo}
|
||||
uc := NewTaskApplicationUsecase(NewTaskUsecase(repo), runtime)
|
||||
if err := uc.Delete(context.Background(), 8); err == nil {
|
||||
t.Fatal("Delete() error = nil, want database error")
|
||||
}
|
||||
if runtime.removeCalls != 0 {
|
||||
t.Fatalf("Remove() calls = %d, want 0", runtime.removeCalls)
|
||||
}
|
||||
}
|
||||
|
|
@ -32,10 +32,33 @@ func (r *apiRepo) APIRoleIDs(ctx context.Context, path, method string) ([]uint,
|
|||
return ids, nil
|
||||
}
|
||||
func (r *apiRepo) SetAPIRoles(ctx context.Context, path, method string, ids []uint) error {
|
||||
access := &authorityAccessRepo{data: r.data}
|
||||
_, allowedAuthorities, strict, err := access.strictAuthorityAccess(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if strict {
|
||||
for _, id := range ids {
|
||||
if !allowedAuthorities[id] {
|
||||
return errors.New("您提交的角色ID不合法")
|
||||
}
|
||||
}
|
||||
if err := r.checkPolicyPathsAuth(ctx, []*biz.API{{Path: path, Method: method}}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
// A path/method need not be present in sys_apis when replacing
|
||||
// all matching policy rows. Do not validate or deduplicate authority IDs.
|
||||
if err := deletePoliciesForPath(tx, path, method); err != nil {
|
||||
if strict {
|
||||
authorityIDs := make([]string, 0, len(allowedAuthorities))
|
||||
for authorityID := range allowedAuthorities {
|
||||
authorityIDs = append(authorityIDs, strconv.FormatUint(uint64(authorityID), 10))
|
||||
}
|
||||
if len(authorityIDs) > 0 {
|
||||
if err := policyScope(tx).Where("v1 = ? AND v2 = ? AND v0 IN ?", path, method, authorityIDs).Delete(&casbinRulePO{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
} else if err := deletePoliciesForPath(tx, path, method); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
|
|
@ -100,42 +123,8 @@ func (r *apiRepo) SetPolicyPaths(ctx context.Context, aid uint, paths []*biz.API
|
|||
if err := (&authorityAccessRepo{data: r.data}).checkAuthorityIDAuth(ctx, aid); err != nil {
|
||||
return err
|
||||
}
|
||||
config := r.data.runtime.Admin()
|
||||
if actor, ok := biz.ActorFromContext(ctx); ok && config != nil && config.System != nil && config.System.UseStrictAuth {
|
||||
var authority authorityPO
|
||||
if err := r.data.gormDB.WithContext(ctx).Where("authority_id = ?", actor.AuthorityID).First(&authority).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var registered []apiPO
|
||||
if err := r.data.gormDB.WithContext(ctx).Find(®istered).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
allowedSet := make(map[string]bool, len(registered))
|
||||
if authority.ParentID == nil || *authority.ParentID == 0 {
|
||||
for _, item := range registered {
|
||||
allowedSet[item.Path+"\x00"+item.Method] = true
|
||||
}
|
||||
} else {
|
||||
policies, err := policyRowsForAuthority(r.data.gormDB.WithContext(ctx), actor.AuthorityID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
policySet := make(map[string]bool, len(policies))
|
||||
for _, item := range policies {
|
||||
policySet[item.V1+"\x00"+item.V2] = true
|
||||
}
|
||||
for _, item := range registered {
|
||||
key := item.Path + "\x00" + item.Method
|
||||
if policySet[key] {
|
||||
allowedSet[key] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, item := range paths {
|
||||
if !allowedSet[item.Path+"\x00"+item.Method] {
|
||||
return errors.New("存在api不在权限列表中")
|
||||
}
|
||||
}
|
||||
if err := r.checkPolicyPathsAuth(ctx, paths); err != nil {
|
||||
return err
|
||||
}
|
||||
db := r.data.gormDB.WithContext(ctx)
|
||||
// The reference enforcer removes the old authority policies before it
|
||||
|
|
@ -165,3 +154,45 @@ func (r *apiRepo) SetPolicyPaths(ctx context.Context, aid uint, paths []*biz.API
|
|||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *apiRepo) checkPolicyPathsAuth(ctx context.Context, paths []*biz.API) error {
|
||||
actor, _, strict, err := (&authorityAccessRepo{data: r.data}).strictAuthorityAccess(ctx)
|
||||
if err != nil || !strict {
|
||||
return err
|
||||
}
|
||||
var authority authorityPO
|
||||
if err := r.data.gormDB.WithContext(ctx).Where("authority_id = ?", actor.AuthorityID).First(&authority).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var registered []apiPO
|
||||
if err := r.data.gormDB.WithContext(ctx).Find(®istered).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
allowedSet := make(map[string]bool, len(registered))
|
||||
if authority.ParentID == nil || *authority.ParentID == 0 {
|
||||
for _, item := range registered {
|
||||
allowedSet[item.Path+"\x00"+item.Method] = true
|
||||
}
|
||||
} else {
|
||||
policies, err := policyRowsForAuthority(r.data.gormDB.WithContext(ctx), actor.AuthorityID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
policySet := make(map[string]bool, len(policies))
|
||||
for _, item := range policies {
|
||||
policySet[item.V1+"\x00"+item.V2] = true
|
||||
}
|
||||
for _, item := range registered {
|
||||
key := item.Path + "\x00" + item.Method
|
||||
if policySet[key] {
|
||||
allowedSet[key] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, item := range paths {
|
||||
if !allowedSet[item.Path+"\x00"+item.Method] {
|
||||
return errors.New("存在api不在权限列表中")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -126,3 +126,40 @@ func TestCheckPolicyStore(t *testing.T) {
|
|||
t.Fatal("CheckPolicyStore succeeded after casbin_rule was dropped")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetAPIRolesStrictOnlyChangesManagedAuthorities(t *testing.T) {
|
||||
data := newPolicyTestData(t)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
rootID, actorID, childID, siblingID := uint(888), uint(2001), uint(2002), uint(3001)
|
||||
if err := db.Create(&[]authorityPO{
|
||||
{AuthorityID: rootID, ParentID: apiPolicyUintPointer(0)},
|
||||
{AuthorityID: actorID, ParentID: &rootID},
|
||||
{AuthorityID: childID, ParentID: &actorID},
|
||||
{AuthorityID: siblingID, ParentID: &rootID},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Create(&apiPO{Path: "/known", Method: "POST"}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Create(&[]casbinRulePO{
|
||||
newPolicyRule(actorID, "/known", "POST"),
|
||||
newPolicyRule(siblingID, "/known", "POST"),
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx := biz.NewActorContext(context.Background(), biz.Actor{AuthorityID: actorID})
|
||||
repo := &apiRepo{data: data}
|
||||
if err := repo.SetAPIRoles(ctx, "/known", "POST", []uint{siblingID}); err == nil {
|
||||
t.Fatal("SetAPIRoles() accepted an out-of-scope authority")
|
||||
}
|
||||
if err := repo.SetAPIRoles(ctx, "/known", "POST", []uint{childID}); err != nil {
|
||||
t.Fatalf("SetAPIRoles() rejected a managed authority: %v", err)
|
||||
}
|
||||
if exists, err := policyExists(db, siblingID, "/known", "POST"); err != nil || !exists {
|
||||
t.Fatalf("out-of-scope policy exists=%v err=%v, want preserved", exists, err)
|
||||
}
|
||||
if exists, err := policyExists(db, childID, "/known", "POST"); err != nil || !exists {
|
||||
t.Fatalf("managed policy exists=%v err=%v, want added", exists, err)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -47,33 +47,105 @@ func (r *authorityAccessRepo) strictAuthorityIDs(ctx context.Context, actorID ui
|
|||
allowed[actorID] = true
|
||||
}
|
||||
walk(actorID)
|
||||
if actor.ParentID != nil && *actor.ParentID != 0 {
|
||||
// A non-root authority manages descendants, never itself. Removing the
|
||||
// actor explicitly also keeps malformed cyclic trees from reopening the
|
||||
// self-management path.
|
||||
delete(allowed, actorID)
|
||||
}
|
||||
return allowed, nil
|
||||
}
|
||||
|
||||
func (r *authorityAccessRepo) checkAuthorityIDAuth(ctx context.Context, targetID uint) error {
|
||||
func (r *authorityAccessRepo) strictAuthorityAccess(ctx context.Context) (biz.Actor, map[uint]bool, bool, error) {
|
||||
config := r.data.runtime.Admin()
|
||||
if config == nil || config.System == nil || !config.System.UseStrictAuth {
|
||||
return nil
|
||||
return biz.Actor{}, nil, false, nil
|
||||
}
|
||||
actor, ok := biz.ActorFromContext(ctx)
|
||||
if !ok {
|
||||
return errors.New("您提交的角色ID不合法")
|
||||
return biz.Actor{}, nil, true, errors.New("您提交的角色ID不合法")
|
||||
}
|
||||
allowed, err := r.strictAuthorityIDs(ctx, actor.AuthorityID)
|
||||
if err != nil {
|
||||
return biz.Actor{}, nil, true, err
|
||||
}
|
||||
return actor, allowed, true, nil
|
||||
}
|
||||
|
||||
func (r *authorityAccessRepo) checkAuthorityIDsAuth(ctx context.Context, targetIDs []uint) error {
|
||||
_, allowed, strict, err := r.strictAuthorityAccess(ctx)
|
||||
if err != nil || !strict {
|
||||
return err
|
||||
}
|
||||
if !allowed[targetID] {
|
||||
return errors.New("您提交的角色ID不合法")
|
||||
for _, targetID := range targetIDs {
|
||||
if !allowed[targetID] {
|
||||
return errors.New("您提交的角色ID不合法")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *authorityAccessRepo) checkAuthorityIDAuth(ctx context.Context, targetID uint) error {
|
||||
return r.checkAuthorityIDsAuth(ctx, []uint{targetID})
|
||||
}
|
||||
|
||||
func managedAuthorityParent(actor biz.Actor, allowed map[uint]bool, targetID uint, parentID *uint, creating bool) (*uint, error) {
|
||||
if creating && (parentID == nil || *parentID == 0) {
|
||||
value := actor.AuthorityID
|
||||
return &value, nil
|
||||
}
|
||||
if !creating && targetID == actor.AuthorityID && allowed[actor.AuthorityID] && (parentID == nil || *parentID == 0) {
|
||||
return parentID, nil
|
||||
}
|
||||
if parentID == nil || *parentID == 0 || *parentID == targetID || (*parentID != actor.AuthorityID && !allowed[*parentID]) {
|
||||
return nil, errors.New("您提交的角色ID不合法")
|
||||
}
|
||||
return parentID, nil
|
||||
}
|
||||
|
||||
func (r *authorityAccessRepo) ensureAuthorityParentAcyclic(ctx context.Context, targetID uint, parentID *uint) error {
|
||||
if parentID == nil || *parentID == 0 {
|
||||
return nil
|
||||
}
|
||||
current := *parentID
|
||||
visited := make(map[uint]struct{})
|
||||
for current != 0 {
|
||||
if current == targetID {
|
||||
return errors.New("角色父级不能形成循环")
|
||||
}
|
||||
if _, seen := visited[current]; seen {
|
||||
return errors.New("角色层级存在循环")
|
||||
}
|
||||
visited[current] = struct{}{}
|
||||
var authority authorityPO
|
||||
if err := r.data.gormDB.WithContext(ctx).Select("authority_id", "parent_id").Where("authority_id = ?", current).First(&authority).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if authority.ParentID == nil {
|
||||
return nil
|
||||
}
|
||||
current = *authority.ParentID
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *authorityAccessRepo) CreateAuthority(ctx context.Context, value *biz.Authority) error {
|
||||
config := r.data.runtime.Admin()
|
||||
if config != nil && config.System != nil && config.System.UseStrictAuth && (value.ParentID == nil || *value.ParentID == 0) {
|
||||
if actor, ok := biz.ActorFromContext(ctx); ok {
|
||||
value.ParentID = &actor.AuthorityID
|
||||
if value.DataScope == 0 {
|
||||
value.DataScope = 1
|
||||
} else if value.DataScope < 1 || value.DataScope > 5 {
|
||||
return errInvalidDataScope
|
||||
}
|
||||
if err := r.checkDataScopeGrant(ctx, value.DataScope, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
actor, allowed, strict, err := r.strictAuthorityAccess(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if strict {
|
||||
value.ParentID, err = managedAuthorityParent(actor, allowed, value.AuthorityID, value.ParentID, true)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
|
|
@ -97,25 +169,58 @@ func (r *authorityAccessRepo) CreateAuthority(ctx context.Context, value *biz.Au
|
|||
value.Menus = []*biz.Menu{{ID: 1, Path: "dashboard", Name: "dashboard", Component: "view/dashboard/index.vue", Sort: 1, Title: "仪表盘", Icon: "setting"}}
|
||||
var dashboard menuPO
|
||||
if err := tx.Where("name = ?", "dashboard").First(&dashboard).Error; err == nil {
|
||||
if err = tx.Create(&authorityMenuPO{SysAuthorityAuthorityID: value.AuthorityID, SysBaseMenuID: dashboard.ID}).Error; err != nil {
|
||||
return err
|
||||
grantDashboard := true
|
||||
if strict && !allowed[actor.AuthorityID] {
|
||||
var count int64
|
||||
if err = tx.Model(&authorityMenuPO{}).Where("sys_authority_authority_id = ? AND sys_base_menu_id = ?", actor.AuthorityID, dashboard.ID).Count(&count).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
grantDashboard = count > 0
|
||||
}
|
||||
if grantDashboard {
|
||||
if err = tx.Create(&authorityMenuPO{SysAuthorityAuthorityID: value.AuthorityID, SysBaseMenuID: dashboard.ID}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
value.Menus = nil
|
||||
}
|
||||
}
|
||||
defaults := []struct{ path, method string }{{"/menu/getMenu", "POST"}, {"/jwt/jsonInBlacklist", "POST"}, {"/base/login", "POST"}, {"/user/changePassword", "POST"}, {"/user/setUserAuthority", "POST"}, {"/user/getUserInfo", "GET"}, {"/user/setSelfInfo", "PUT"}, {"/fileUploadAndDownload/upload", "POST"}, {"/sysDictionary/findSysDictionary", "GET"}}
|
||||
rules := make([]casbinRulePO, 0, len(defaults))
|
||||
for _, item := range defaults {
|
||||
if strict {
|
||||
permitted, policyErr := r.copyPolicyAllowed(ctx, tx, actor.AuthorityID, item.path, item.method)
|
||||
if policyErr != nil {
|
||||
return policyErr
|
||||
}
|
||||
if !permitted {
|
||||
continue
|
||||
}
|
||||
}
|
||||
rules = append(rules, newPolicyRule(value.AuthorityID, item.path, item.method))
|
||||
}
|
||||
if err := tx.Create(&rules).Error; err != nil {
|
||||
return err
|
||||
if len(rules) > 0 {
|
||||
if err := tx.Create(&rules).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
func (r *authorityAccessRepo) CopyAuthority(ctx context.Context, sourceID uint, value *biz.Authority) error {
|
||||
if value.DataScope == 0 {
|
||||
value.DataScope = 1
|
||||
} else if value.DataScope < 1 || value.DataScope > 5 {
|
||||
return errInvalidDataScope
|
||||
}
|
||||
if err := r.checkDataScopeGrant(ctx, value.DataScope, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
actor, allowed, strict, err := r.strictAuthorityAccess(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
config := r.data.runtime.Admin()
|
||||
actor, hasActor := biz.ActorFromContext(ctx)
|
||||
// Reject a duplicate target ID before performing the
|
||||
// hierarchy/Casbin checks that happen later in UpdateCasbin.
|
||||
var existing authorityPO
|
||||
|
|
@ -128,16 +233,14 @@ func (r *authorityAccessRepo) CopyAuthority(ctx context.Context, sourceID uint,
|
|||
// that its requested parent is the actor or one of the actor's allowed
|
||||
// descendants. This preserves the compatible error and prevents a transient role
|
||||
// from being visible when the check is guaranteed to fail.
|
||||
strict := hasActor && config != nil && config.System != nil && config.System.UseStrictAuth
|
||||
if strict {
|
||||
allowed, err := r.strictAuthorityIDs(ctx, actor.AuthorityID)
|
||||
if sourceID != actor.AuthorityID && !allowed[sourceID] {
|
||||
return errors.New("您提交的角色ID不合法")
|
||||
}
|
||||
value.ParentID, err = managedAuthorityParent(actor, allowed, value.AuthorityID, value.ParentID, true)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
parentAllowed := value.ParentID != nil && (allowed[*value.ParentID] || *value.ParentID == actor.AuthorityID)
|
||||
if !parentAllowed {
|
||||
return errors.New("您提交的角色ID不合法")
|
||||
}
|
||||
}
|
||||
po := authorityPO{AuthorityID: value.AuthorityID, AuthorityName: value.AuthorityName, ParentID: value.ParentID, DataScope: value.DataScope, DefaultRouter: value.DefaultRouter}
|
||||
if err := tx.Create(&po).Error; err != nil {
|
||||
|
|
@ -152,6 +255,15 @@ func (r *authorityAccessRepo) CopyAuthority(ctx context.Context, sourceID uint,
|
|||
if err := copyLinks("sys_authority_menus", &menus, map[string]any{"sys_authority_authority_id": sourceID}); err != nil {
|
||||
return err
|
||||
}
|
||||
if strict && len(menus) > 0 {
|
||||
menuIDs := make([]uint, 0, len(menus))
|
||||
for _, menu := range menus {
|
||||
menuIDs = append(menuIDs, menu.SysBaseMenuID)
|
||||
}
|
||||
if err := checkMenuAssignment(tx.WithContext(ctx), actor.AuthorityID, allowed[actor.AuthorityID], menuIDs); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for i := range menus {
|
||||
menus[i].SysAuthorityAuthorityID = value.AuthorityID
|
||||
}
|
||||
|
|
@ -206,6 +318,15 @@ func (r *authorityAccessRepo) CopyAuthority(ctx context.Context, sourceID uint,
|
|||
if err := copyLinks("sys_authority_btns", &buttons, map[string]any{"authority_id": sourceID}); err != nil {
|
||||
return err
|
||||
}
|
||||
if strict && len(buttons) > 0 {
|
||||
requested := make(map[uint][]uint)
|
||||
for _, button := range buttons {
|
||||
requested[button.MenuID] = append(requested[button.MenuID], button.ButtonID)
|
||||
}
|
||||
if err := checkButtonAssignment(tx.WithContext(ctx), actor.AuthorityID, allowed[actor.AuthorityID], requested); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for i := range buttons {
|
||||
buttons[i].AuthorityID = value.AuthorityID
|
||||
}
|
||||
|
|
@ -240,6 +361,38 @@ func (r *authorityAccessRepo) copyPolicyAllowed(ctx context.Context, tx *gorm.DB
|
|||
return policyExists(tx.WithContext(ctx), actorID, path, method)
|
||||
}
|
||||
func (r *authorityAccessRepo) UpdateAuthority(ctx context.Context, value *biz.Authority) error {
|
||||
if err := r.checkAuthorityIDAuth(ctx, value.AuthorityID); err != nil {
|
||||
return err
|
||||
}
|
||||
if value.DataScope < 0 || value.DataScope > 5 {
|
||||
return errInvalidDataScope
|
||||
}
|
||||
if value.DataScope != 0 {
|
||||
var departmentIDs []uint
|
||||
if value.DataScope == 5 {
|
||||
var err error
|
||||
departmentIDs, err = r.DataScopeDepartmentIDs(ctx, value.AuthorityID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := r.checkDataScopeGrant(ctx, value.DataScope, departmentIDs); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
actor, allowed, strict, err := r.strictAuthorityAccess(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if strict {
|
||||
value.ParentID, err = managedAuthorityParent(actor, allowed, value.AuthorityID, value.ParentID, false)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err = r.ensureAuthorityParentAcyclic(ctx, value.AuthorityID, value.ParentID); err != nil {
|
||||
return err
|
||||
}
|
||||
db := r.data.gormDB.WithContext(ctx)
|
||||
var current authorityPO
|
||||
if err := db.Where("authority_id = ?", value.AuthorityID).First(¤t).Error; err != nil {
|
||||
|
|
@ -249,6 +402,9 @@ func (r *authorityAccessRepo) UpdateAuthority(ctx context.Context, value *biz.Au
|
|||
return db.Model(¤t).Updates(updates).Error
|
||||
}
|
||||
func (r *authorityAccessRepo) DeleteAuthority(ctx context.Context, id uint) error {
|
||||
if err := r.checkAuthorityIDAuth(ctx, id); err != nil {
|
||||
return err
|
||||
}
|
||||
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var authority authorityPO
|
||||
if err := tx.Where("authority_id = ?", id).First(&authority).Error; err != nil {
|
||||
|
|
@ -341,11 +497,26 @@ func (r *authorityAccessRepo) ListAuthorities(ctx context.Context) ([]*biz.Autho
|
|||
return out, nil
|
||||
}
|
||||
func (r *authorityAccessRepo) SetAuthorityUsers(ctx context.Context, id uint, ids []uint) error {
|
||||
_, allowed, strict, err := r.strictAuthorityAccess(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if strict && !allowed[id] {
|
||||
return errors.New("您提交的角色ID不合法")
|
||||
}
|
||||
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var oldIDs []uint
|
||||
if err := tx.Model(&userAuthorityPO{}).Where("sys_authority_authority_id = ?", id).Pluck("sys_user_id", &oldIDs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if strict {
|
||||
if err := checkManagedUserIDs(tx, ids, allowed, true); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := checkManagedUserIDs(tx, oldIDs, allowed, false); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := tx.Where("sys_authority_authority_id = ?", id).Delete(&userAuthorityPO{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
|
@ -382,12 +553,111 @@ func (r *authorityAccessRepo) SetAuthorityUsers(ctx context.Context, id uint, id
|
|||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func checkManagedUserIDs(tx *gorm.DB, ids []uint, allowedAuthorities map[uint]bool, requireAll bool) error {
|
||||
unique := make(map[uint]struct{}, len(ids))
|
||||
for _, id := range ids {
|
||||
if id == 0 {
|
||||
if requireAll {
|
||||
return errors.New("您提交的用户ID不合法")
|
||||
}
|
||||
continue
|
||||
}
|
||||
unique[id] = struct{}{}
|
||||
}
|
||||
if len(unique) == 0 {
|
||||
return nil
|
||||
}
|
||||
userIDs := make([]uint, 0, len(unique))
|
||||
for id := range unique {
|
||||
userIDs = append(userIDs, id)
|
||||
}
|
||||
var users []userPO
|
||||
if err := tx.Select("id", "authority_id").Where("id IN ?", userIDs).Find(&users).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if requireAll && len(users) != len(unique) {
|
||||
return errors.New("您提交的用户ID不合法")
|
||||
}
|
||||
managedUsers := make(map[uint]bool, len(users))
|
||||
for _, user := range users {
|
||||
if !allowedAuthorities[user.AuthorityID] {
|
||||
return errors.New("您提交的用户ID不合法")
|
||||
}
|
||||
managedUsers[user.ID] = true
|
||||
}
|
||||
if len(managedUsers) == 0 {
|
||||
return nil
|
||||
}
|
||||
var links []userAuthorityPO
|
||||
if err := tx.Where("sys_user_id IN ?", userIDs).Find(&links).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, link := range links {
|
||||
if managedUsers[link.SysUserID] && !allowedAuthorities[link.SysAuthorityAuthorityID] {
|
||||
return errors.New("您提交的用户ID不合法")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *authorityAccessRepo) checkUserIDAuth(ctx context.Context, id uint, allowSelf bool) error {
|
||||
actor, allowed, strict, err := r.strictAuthorityAccess(ctx)
|
||||
if err != nil || !strict {
|
||||
return err
|
||||
}
|
||||
if allowSelf && actor.UserID != 0 && actor.UserID == id {
|
||||
return nil
|
||||
}
|
||||
return checkManagedUserIDs(r.data.gormDB.WithContext(ctx), []uint{id}, allowed, true)
|
||||
}
|
||||
|
||||
func (r *authorityAccessRepo) checkDepartmentIDsAuth(ctx context.Context, ids []uint) error {
|
||||
actor, allowedAuthorities, strict, err := r.strictAuthorityAccess(ctx)
|
||||
if err != nil || !strict || len(ids) == 0 {
|
||||
return err
|
||||
}
|
||||
if allowedAuthorities[actor.AuthorityID] {
|
||||
return nil
|
||||
}
|
||||
var authority authorityPO
|
||||
if err := r.data.gormDB.WithContext(ctx).Select("authority_id", "data_scope").Where("authority_id = ?", actor.AuthorityID).First(&authority).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if authority.DataScope == 1 {
|
||||
return nil
|
||||
}
|
||||
scope, err := r.ResolveDataScope(ctx, actor.AuthorityID, actor.UserID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
visible := make(map[uint]bool, len(scope.DepartmentIDs))
|
||||
for _, id := range scope.DepartmentIDs {
|
||||
visible[id] = true
|
||||
}
|
||||
for _, id := range ids {
|
||||
if id == 0 || !visible[id] {
|
||||
return errors.New("您提交的部门ID不合法")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *authorityAccessRepo) AuthorityUserIDs(ctx context.Context, id uint) ([]uint, error) {
|
||||
var ids []uint
|
||||
err := r.data.gormDB.WithContext(ctx).Model(&userAuthorityPO{}).Where("sys_authority_authority_id = ?", id).Pluck("sys_user_id", &ids).Error
|
||||
return ids, err
|
||||
}
|
||||
func (r *authorityAccessRepo) SetDataScope(ctx context.Context, id uint, scope int, deptIDs []uint) error {
|
||||
if err := r.checkAuthorityIDAuth(ctx, id); err != nil {
|
||||
return err
|
||||
}
|
||||
if scope < 1 || scope > 5 {
|
||||
return errInvalidDataScope
|
||||
}
|
||||
if err := r.checkDataScopeGrant(ctx, scope, deptIDs); err != nil {
|
||||
return err
|
||||
}
|
||||
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Model(&authorityPO{}).Where("authority_id = ?", id).Update("data_scope", scope).Error; err != nil {
|
||||
return err
|
||||
|
|
@ -407,6 +677,74 @@ func (r *authorityAccessRepo) SetDataScope(ctx context.Context, id uint, scope i
|
|||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (r *authorityAccessRepo) checkDataScopeGrant(ctx context.Context, requestedScope int, departmentIDs []uint) error {
|
||||
actor, allowedAuthorities, strict, err := r.strictAuthorityAccess(ctx)
|
||||
if err != nil || !strict {
|
||||
return err
|
||||
}
|
||||
requestedDepartments := make(map[uint]struct{}, len(departmentIDs))
|
||||
if requestedScope == 5 {
|
||||
for _, departmentID := range departmentIDs {
|
||||
if departmentID == 0 {
|
||||
return errors.New("您提交的部门ID不合法")
|
||||
}
|
||||
requestedDepartments[departmentID] = struct{}{}
|
||||
}
|
||||
if len(requestedDepartments) > 0 {
|
||||
var count int64
|
||||
if err := r.data.gormDB.WithContext(ctx).Model(&departmentPO{}).Where("id IN ?", departmentIDs).Count(&count).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if count != int64(len(requestedDepartments)) {
|
||||
return errors.New("您提交的部门ID不合法")
|
||||
}
|
||||
}
|
||||
}
|
||||
if allowedAuthorities[actor.AuthorityID] {
|
||||
return nil
|
||||
}
|
||||
var actorAuthority authorityPO
|
||||
if err := r.data.gormDB.WithContext(ctx).Select("authority_id", "data_scope").Where("authority_id = ?", actor.AuthorityID).First(&actorAuthority).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
grantable := false
|
||||
switch actorAuthority.DataScope {
|
||||
case 1:
|
||||
grantable = true
|
||||
case 2:
|
||||
grantable = requestedScope == 2 || requestedScope == 3 || requestedScope == 4 || requestedScope == 5
|
||||
case 3:
|
||||
grantable = requestedScope == 3 || requestedScope == 4 || requestedScope == 5
|
||||
case 4:
|
||||
grantable = requestedScope == 4
|
||||
case 5:
|
||||
grantable = requestedScope == 4 || requestedScope == 5
|
||||
default:
|
||||
return errInvalidDataScope
|
||||
}
|
||||
if !grantable {
|
||||
return errInvalidDataScope
|
||||
}
|
||||
if requestedScope != 5 || len(requestedDepartments) == 0 || actorAuthority.DataScope == 1 {
|
||||
return nil
|
||||
}
|
||||
actorScope, err := r.ResolveDataScope(ctx, actor.AuthorityID, actor.UserID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
allowedDepartments := make(map[uint]bool, len(actorScope.DepartmentIDs))
|
||||
for _, departmentID := range actorScope.DepartmentIDs {
|
||||
allowedDepartments[departmentID] = true
|
||||
}
|
||||
for departmentID := range requestedDepartments {
|
||||
if !allowedDepartments[departmentID] {
|
||||
return errors.New("您提交的部门ID不合法")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *authorityAccessRepo) DataScopeDepartmentIDs(ctx context.Context, id uint) ([]uint, error) {
|
||||
var ids []uint
|
||||
err := r.data.gormDB.WithContext(ctx).Model(&authorityDepartmentPO{}).Where("sys_authority_authority_id = ?", id).Pluck("sys_department_id", &ids).Error
|
||||
|
|
@ -415,37 +753,48 @@ func (r *authorityAccessRepo) DataScopeDepartmentIDs(ctx context.Context, id uin
|
|||
func (r *authorityAccessRepo) ResolveDataScope(ctx context.Context, authorityID, userID uint) (biz.DataScope, error) {
|
||||
identity := biz.DataScope{UserID: userID, AuthorityID: authorityID}
|
||||
var user userPO
|
||||
_ = r.data.gormDB.WithContext(ctx).Select("id", "dept_id").First(&user, userID).Error
|
||||
if err := r.data.gormDB.WithContext(ctx).Select("id", "dept_id").First(&user, userID).Error; err != nil {
|
||||
return identity, err
|
||||
}
|
||||
identity.PrimaryDeptID = user.DeptID
|
||||
var authority authorityPO
|
||||
_ = r.data.gormDB.WithContext(ctx).Select("authority_id", "data_scope").First(&authority, "authority_id = ?", authorityID).Error
|
||||
if err := r.data.gormDB.WithContext(ctx).Select("authority_id", "data_scope").First(&authority, "authority_id = ?", authorityID).Error; err != nil {
|
||||
return identity, err
|
||||
}
|
||||
identity.Scope = authority.DataScope
|
||||
if identity.Scope == 0 {
|
||||
identity.Scope = 1
|
||||
if identity.Scope < 1 || identity.Scope > 5 {
|
||||
return identity, errInvalidDataScope
|
||||
}
|
||||
identity.All = identity.Scope == 1
|
||||
if identity.Scope == 4 {
|
||||
identity.OwnerUserID = userID
|
||||
}
|
||||
var ids []uint
|
||||
_ = r.data.gormDB.WithContext(ctx).Model(&userDepartmentPO{}).Where("sys_user_id = ?", userID).Pluck("sys_department_id", &ids).Error
|
||||
selected := make(map[uint]bool, len(ids)+1)
|
||||
for _, id := range ids {
|
||||
selected[id] = true
|
||||
}
|
||||
if user.DeptID != 0 {
|
||||
selected[user.DeptID] = true
|
||||
}
|
||||
ids = ids[:0]
|
||||
for id := range selected {
|
||||
ids = append(ids, id)
|
||||
selected := make(map[uint]bool)
|
||||
if identity.Scope == 2 || identity.Scope == 3 {
|
||||
if err := r.data.gormDB.WithContext(ctx).Model(&userDepartmentPO{}).Where("sys_user_id = ?", userID).Pluck("sys_department_id", &ids).Error; err != nil {
|
||||
return identity, err
|
||||
}
|
||||
selected = make(map[uint]bool, len(ids)+1)
|
||||
for _, id := range ids {
|
||||
selected[id] = true
|
||||
}
|
||||
if user.DeptID != 0 {
|
||||
selected[user.DeptID] = true
|
||||
}
|
||||
ids = ids[:0]
|
||||
for id := range selected {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
}
|
||||
if identity.Scope == 2 && len(ids) > 0 {
|
||||
// Match the administration subtree-union semantics: derive the visible set from
|
||||
// parent_id edges at read time, rather than trusting the denormalized
|
||||
// ancestors string (which may be stale after a department move).
|
||||
var departments []departmentPO
|
||||
_ = r.data.gormDB.WithContext(ctx).Select("id", "parent_id").Find(&departments).Error
|
||||
if err := r.data.gormDB.WithContext(ctx).Select("id", "parent_id").Find(&departments).Error; err != nil {
|
||||
return identity, err
|
||||
}
|
||||
children := make(map[uint][]uint, len(departments))
|
||||
for _, department := range departments {
|
||||
children[department.ParentID] = append(children[department.ParentID], department.ID)
|
||||
|
|
@ -473,7 +822,11 @@ func (r *authorityAccessRepo) ResolveDataScope(ctx context.Context, authorityID,
|
|||
ids = append(ids, id)
|
||||
}
|
||||
} else if identity.Scope == 5 {
|
||||
ids, _ = r.DataScopeDepartmentIDs(ctx, authorityID)
|
||||
var err error
|
||||
ids, err = r.DataScopeDepartmentIDs(ctx, authorityID)
|
||||
if err != nil {
|
||||
return identity, err
|
||||
}
|
||||
}
|
||||
identity.DepartmentIDs = ids
|
||||
return identity, nil
|
||||
|
|
|
|||
|
|
@ -177,3 +177,405 @@ func TestAuthorityCustomDataScopeUsesCompatibleColumns(t *testing.T) {
|
|||
t.Fatalf("non-custom scope retained departments: %v", ids)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetAuthorityUsersStrictRejectsUsersOutsideManagedRoles(t *testing.T) {
|
||||
data := newTransactionTestData(t)
|
||||
enableStrictAuthorityTestMode(data)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
rootID, actorID, childID, siblingID := uint(888), uint(1000), uint(1001), uint(2000)
|
||||
if err := db.Create(&[]authorityPO{
|
||||
{AuthorityID: rootID, ParentID: authorityUintPointer(0)},
|
||||
{AuthorityID: actorID, ParentID: &rootID},
|
||||
{AuthorityID: childID, ParentID: &actorID},
|
||||
{AuthorityID: siblingID, ParentID: &rootID},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
users := []userPO{
|
||||
{Username: "managed", Password: "hash", AuthorityID: childID, Enable: 1},
|
||||
{Username: "outside", Password: "hash", AuthorityID: siblingID, Enable: 1},
|
||||
}
|
||||
if err := db.Create(&users).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx := biz.NewActorContext(context.Background(), biz.Actor{AuthorityID: actorID})
|
||||
repo := &authorityAccessRepo{data: data}
|
||||
if err := repo.SetAuthorityUsers(ctx, childID, []uint{users[1].ID}); err == nil {
|
||||
t.Fatal("SetAuthorityUsers() accepted a user outside the managed role tree")
|
||||
}
|
||||
if err := repo.SetAuthorityUsers(ctx, childID, []uint{users[0].ID}); err != nil {
|
||||
t.Fatalf("SetAuthorityUsers() rejected a managed user: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetAuthorityUsersStrictRejectsMixedRoleUserAlreadyLinked(t *testing.T) {
|
||||
data := newTransactionTestData(t)
|
||||
enableStrictAuthorityTestMode(data)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
rootID, actorID, childID, siblingID := uint(888), uint(1100), uint(1101), uint(2100)
|
||||
if err := db.Create(&[]authorityPO{
|
||||
{AuthorityID: rootID, ParentID: authorityUintPointer(0)},
|
||||
{AuthorityID: actorID, ParentID: &rootID},
|
||||
{AuthorityID: childID, ParentID: &actorID},
|
||||
{AuthorityID: siblingID, ParentID: &rootID},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
user := userPO{Username: "mixed", Password: "hash", AuthorityID: childID, Enable: 1}
|
||||
if err := db.Create(&user).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Create(&[]userAuthorityPO{
|
||||
{SysUserID: user.ID, SysAuthorityAuthorityID: childID},
|
||||
{SysUserID: user.ID, SysAuthorityAuthorityID: siblingID},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx := biz.NewActorContext(context.Background(), biz.Actor{AuthorityID: actorID})
|
||||
if err := (&authorityAccessRepo{data: data}).SetAuthorityUsers(ctx, childID, nil); err == nil {
|
||||
t.Fatal("SetAuthorityUsers() modified a linked user that also has an out-of-scope role")
|
||||
}
|
||||
var count int64
|
||||
if err := db.Model(&userAuthorityPO{}).Where("sys_user_id = ? AND sys_authority_authority_id = ?", user.ID, childID).Count(&count).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if count != 1 {
|
||||
t.Fatalf("rejected update changed the existing target-role link: count=%d", count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateAuthorityStrictRejectsHierarchyCycle(t *testing.T) {
|
||||
data := newTransactionTestData(t)
|
||||
enableStrictAuthorityTestMode(data)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
rootID, actorID, targetID, childID := uint(888), uint(1200), uint(1201), uint(1202)
|
||||
if err := db.Create(&[]authorityPO{
|
||||
{AuthorityID: rootID, ParentID: authorityUintPointer(0)},
|
||||
{AuthorityID: actorID, ParentID: &rootID},
|
||||
{AuthorityID: targetID, AuthorityName: "target", ParentID: &actorID},
|
||||
{AuthorityID: childID, ParentID: &targetID},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx := biz.NewActorContext(context.Background(), biz.Actor{AuthorityID: actorID})
|
||||
err := (&authorityAccessRepo{data: data}).UpdateAuthority(ctx, &biz.Authority{AuthorityID: targetID, AuthorityName: "target", ParentID: &childID})
|
||||
if err == nil {
|
||||
t.Fatal("UpdateAuthority() accepted a parent that forms a cycle")
|
||||
}
|
||||
var stored authorityPO
|
||||
if err := db.Where("authority_id = ?", targetID).First(&stored).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if stored.ParentID == nil || *stored.ParentID != actorID {
|
||||
t.Fatalf("rejected update changed parent_id: %+v", stored.ParentID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetDataScopeValidatesScopeAndStrictDepartments(t *testing.T) {
|
||||
data := newTransactionTestData(t)
|
||||
enableStrictAuthorityTestMode(data)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
rootID, actorID, childID := uint(888), uint(1300), uint(1301)
|
||||
if err := db.Create(&[]authorityPO{
|
||||
{AuthorityID: rootID, ParentID: authorityUintPointer(0)},
|
||||
{AuthorityID: actorID, ParentID: &rootID},
|
||||
{AuthorityID: childID, ParentID: &actorID},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx := biz.NewActorContext(context.Background(), biz.Actor{AuthorityID: actorID})
|
||||
repo := &authorityAccessRepo{data: data}
|
||||
if err := repo.SetDataScope(ctx, childID, 0, nil); !errors.Is(err, errInvalidDataScope) {
|
||||
t.Fatalf("SetDataScope() invalid scope error = %v", err)
|
||||
}
|
||||
if err := repo.SetDataScope(ctx, childID, 5, []uint{999999}); err == nil {
|
||||
t.Fatal("SetDataScope() accepted a missing department in strict mode")
|
||||
}
|
||||
department := departmentPO{Name: "managed"}
|
||||
if err := db.Create(&department).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := repo.SetDataScope(ctx, childID, 5, []uint{department.ID}); err != nil {
|
||||
t.Fatalf("SetDataScope() rejected an existing department: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetDataScopeStrictRejectsDepartmentsOutsideActorScope(t *testing.T) {
|
||||
data := newTransactionTestData(t)
|
||||
enableStrictAuthorityTestMode(data)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
rootID, actorID, childID := uint(888), uint(1350), uint(1351)
|
||||
if err := db.Create(&[]authorityPO{
|
||||
{AuthorityID: rootID, ParentID: authorityUintPointer(0)},
|
||||
{AuthorityID: actorID, ParentID: &rootID, DataScope: 5},
|
||||
{AuthorityID: childID, ParentID: &actorID},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
departments := []departmentPO{{Name: "allowed"}, {Name: "outside"}}
|
||||
if err := db.Create(&departments).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
actorUser := userPO{Username: "scope-actor", Password: "hash", AuthorityID: actorID, DeptID: departments[0].ID, Enable: 1}
|
||||
if err := db.Create(&actorUser).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Create(&authorityDepartmentPO{AuthorityID: actorID, DepartmentID: departments[0].ID}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx := biz.NewActorContext(context.Background(), biz.Actor{UserID: actorUser.ID, AuthorityID: actorID})
|
||||
repo := &authorityAccessRepo{data: data}
|
||||
if err := repo.SetDataScope(ctx, childID, 5, []uint{departments[1].ID}); err == nil {
|
||||
t.Fatal("SetDataScope() accepted a department outside the actor's data scope")
|
||||
}
|
||||
if err := repo.SetDataScope(ctx, childID, 5, []uint{departments[0].ID}); err != nil {
|
||||
t.Fatalf("SetDataScope() rejected a department inside the actor's data scope: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStrictDataScopeGrantRejectsBroaderChildScopes(t *testing.T) {
|
||||
data := newTransactionTestData(t)
|
||||
enableStrictAuthorityTestMode(data)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
rootID, actorID, sourceID, targetID := uint(888), uint(1370), uint(1371), uint(1372)
|
||||
if err := db.Create(&[]authorityPO{
|
||||
{AuthorityID: rootID, ParentID: authorityUintPointer(0), DataScope: 1},
|
||||
{AuthorityID: actorID, ParentID: &rootID, DataScope: 3},
|
||||
{AuthorityID: sourceID, ParentID: &actorID, DataScope: 4},
|
||||
{AuthorityID: targetID, AuthorityName: "target", ParentID: &actorID, DataScope: 4},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
actorUser := userPO{Username: "limited-scope-actor", Password: "hash", AuthorityID: actorID, Enable: 1}
|
||||
if err := db.Create(&actorUser).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx := biz.NewActorContext(context.Background(), biz.Actor{UserID: actorUser.ID, AuthorityID: actorID})
|
||||
repo := &authorityAccessRepo{data: data}
|
||||
|
||||
created := &biz.Authority{AuthorityID: 1373, AuthorityName: "created", ParentID: &actorID}
|
||||
if err := repo.CreateAuthority(ctx, created); !errors.Is(err, errInvalidDataScope) {
|
||||
t.Fatalf("CreateAuthority() broader scope error = %v", err)
|
||||
}
|
||||
if created.DataScope != 1 {
|
||||
t.Fatalf("CreateAuthority() did not normalize zero scope before validation: %d", created.DataScope)
|
||||
}
|
||||
copied := &biz.Authority{AuthorityID: 1374, AuthorityName: "copied", ParentID: &actorID}
|
||||
if err := repo.CopyAuthority(ctx, sourceID, copied); !errors.Is(err, errInvalidDataScope) {
|
||||
t.Fatalf("CopyAuthority() broader scope error = %v", err)
|
||||
}
|
||||
if copied.DataScope != 1 {
|
||||
t.Fatalf("CopyAuthority() did not normalize zero scope before validation: %d", copied.DataScope)
|
||||
}
|
||||
if err := repo.UpdateAuthority(ctx, &biz.Authority{AuthorityID: targetID, AuthorityName: "target", ParentID: &actorID, DataScope: 2}); !errors.Is(err, errInvalidDataScope) {
|
||||
t.Fatalf("UpdateAuthority() broader scope error = %v", err)
|
||||
}
|
||||
if err := repo.SetDataScope(ctx, targetID, 1, nil); !errors.Is(err, errInvalidDataScope) {
|
||||
t.Fatalf("SetDataScope() broader scope error = %v", err)
|
||||
}
|
||||
|
||||
allowed := &biz.Authority{AuthorityID: 1375, AuthorityName: "self-only", ParentID: &actorID, DataScope: 4}
|
||||
if err := repo.CreateAuthority(ctx, allowed); err != nil {
|
||||
t.Fatalf("CreateAuthority() rejected a narrower scope: %v", err)
|
||||
}
|
||||
var stored authorityPO
|
||||
if err := db.Where("authority_id = ?", targetID).First(&stored).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if stored.DataScope != 4 {
|
||||
t.Fatalf("rejected grants changed target data scope to %d", stored.DataScope)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateAuthorityStrictDefaultsStayWithinActorPermissions(t *testing.T) {
|
||||
data := newTransactionTestData(t)
|
||||
enableStrictAuthorityTestMode(data)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
rootID, actorID := uint(888), uint(1400)
|
||||
if err := db.Create(&[]authorityPO{
|
||||
{AuthorityID: rootID, ParentID: authorityUintPointer(0)},
|
||||
{AuthorityID: actorID, ParentID: &rootID},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
dashboard := menuPO{Name: "dashboard", Path: "dashboard"}
|
||||
if err := db.Create(&dashboard).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Create(&apiPO{Path: "/menu/getMenu", Method: "POST"}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx := biz.NewActorContext(context.Background(), biz.Actor{AuthorityID: actorID})
|
||||
repo := &authorityAccessRepo{data: data}
|
||||
first := &biz.Authority{AuthorityID: 1401, AuthorityName: "no-defaults", ParentID: &actorID}
|
||||
if err := repo.CreateAuthority(ctx, first); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var menuLinks int64
|
||||
if err := db.Model(&authorityMenuPO{}).Where("sys_authority_authority_id = ?", first.AuthorityID).Count(&menuLinks).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if menuLinks != 0 {
|
||||
t.Fatalf("new child inherited an unowned default menu: count=%d", menuLinks)
|
||||
}
|
||||
if exists, err := policyExists(db, first.AuthorityID, "/menu/getMenu", "POST"); err != nil || exists {
|
||||
t.Fatalf("new child inherited an unowned default API: exists=%v err=%v", exists, err)
|
||||
}
|
||||
|
||||
if err := db.Create(&authorityMenuPO{SysAuthorityAuthorityID: actorID, SysBaseMenuID: dashboard.ID}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
policy := newPolicyRule(actorID, "/menu/getMenu", "POST")
|
||||
if err := db.Create(&policy).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
second := &biz.Authority{AuthorityID: 1402, AuthorityName: "owned-defaults", ParentID: &actorID}
|
||||
if err := repo.CreateAuthority(ctx, second); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Model(&authorityMenuPO{}).Where("sys_authority_authority_id = ? AND sys_base_menu_id = ?", second.AuthorityID, dashboard.ID).Count(&menuLinks).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if menuLinks != 1 {
|
||||
t.Fatalf("new child did not inherit the owned dashboard menu: count=%d", menuLinks)
|
||||
}
|
||||
if exists, err := policyExists(db, second.AuthorityID, "/menu/getMenu", "POST"); err != nil || !exists {
|
||||
t.Fatalf("new child owned default API exists=%v err=%v", exists, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopyAuthorityStrictValidatesCopiedMenusAndButtons(t *testing.T) {
|
||||
data := newTransactionTestData(t)
|
||||
enableStrictAuthorityTestMode(data)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
rootID, actorID, sourceID := uint(888), uint(1500), uint(1501)
|
||||
if err := db.Create(&[]authorityPO{
|
||||
{AuthorityID: rootID, ParentID: authorityUintPointer(0)},
|
||||
{AuthorityID: actorID, ParentID: &rootID},
|
||||
{AuthorityID: sourceID, ParentID: &actorID},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Create(&[]menuPO{{ID: 10, Name: "owned"}, {ID: 11, Name: "unowned"}}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Create(&[]menuButtonPO{{ID: 31, MenuID: 10, Name: "owned"}, {ID: 32, MenuID: 10, Name: "unowned"}}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Create(&[]authorityMenuPO{
|
||||
{SysAuthorityAuthorityID: actorID, SysBaseMenuID: 10},
|
||||
{SysAuthorityAuthorityID: sourceID, SysBaseMenuID: 10},
|
||||
{SysAuthorityAuthorityID: sourceID, SysBaseMenuID: 11},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Create(&[]authorityButtonPO{
|
||||
{AuthorityID: actorID, MenuID: 10, ButtonID: 31},
|
||||
{AuthorityID: sourceID, MenuID: 10, ButtonID: 31},
|
||||
{AuthorityID: sourceID, MenuID: 10, ButtonID: 32},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx := biz.NewActorContext(context.Background(), biz.Actor{AuthorityID: actorID})
|
||||
repo := &authorityAccessRepo{data: data}
|
||||
if err := repo.CopyAuthority(ctx, sourceID, &biz.Authority{AuthorityID: 1510, AuthorityName: "menu-fail", ParentID: &actorID}); err == nil {
|
||||
t.Fatal("CopyAuthority() copied a menu not assigned to the actor")
|
||||
}
|
||||
if err := db.Where("sys_authority_authority_id = ? AND sys_base_menu_id = ?", sourceID, 11).Delete(&authorityMenuPO{}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := repo.CopyAuthority(ctx, sourceID, &biz.Authority{AuthorityID: 1511, AuthorityName: "button-fail", ParentID: &actorID}); err == nil {
|
||||
t.Fatal("CopyAuthority() copied a button not assigned to the actor")
|
||||
}
|
||||
if err := db.Where("authority_id = ? AND sys_base_menu_btn_id = ?", sourceID, 32).Delete(&authorityButtonPO{}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := repo.CopyAuthority(ctx, sourceID, &biz.Authority{AuthorityID: 1512, AuthorityName: "valid", ParentID: &actorID}); err != nil {
|
||||
t.Fatalf("CopyAuthority() rejected owned permissions: %v", err)
|
||||
}
|
||||
for _, failedID := range []uint{1510, 1511} {
|
||||
var count int64
|
||||
if err := db.Unscoped().Model(&authorityPO{}).Where("authority_id = ?", failedID).Count(&count).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if count != 0 {
|
||||
t.Fatalf("failed copy persisted authority %d", failedID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthorityCreateAndCopyNormalizeZeroDataScope(t *testing.T) {
|
||||
data := newTransactionTestData(t)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
repo := &authorityAccessRepo{data: data}
|
||||
|
||||
created := &biz.Authority{AuthorityID: 1600, AuthorityName: "created"}
|
||||
if err := repo.CreateAuthority(context.Background(), created); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if created.DataScope != 1 {
|
||||
t.Fatalf("created data scope = %d, want 1", created.DataScope)
|
||||
}
|
||||
|
||||
if err := db.Create(&authorityPO{AuthorityID: 1601, AuthorityName: "source", DataScope: 3}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
copied := &biz.Authority{AuthorityID: 1602, AuthorityName: "copied"}
|
||||
if err := repo.CopyAuthority(context.Background(), 1601, copied); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if copied.DataScope != 1 {
|
||||
t.Fatalf("copied data scope = %d, want 1", copied.DataScope)
|
||||
}
|
||||
|
||||
var stored []authorityPO
|
||||
if err := db.Where("authority_id IN ?", []uint{created.AuthorityID, copied.AuthorityID}).Order("authority_id").Find(&stored).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(stored) != 2 || stored[0].DataScope != 1 || stored[1].DataScope != 1 {
|
||||
t.Fatalf("stored normalized data scopes = %+v", stored)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthorityMutationsRejectInvalidDataScope(t *testing.T) {
|
||||
data := newTransactionTestData(t)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
repo := &authorityAccessRepo{data: data}
|
||||
if err := db.Create(&[]authorityPO{
|
||||
{AuthorityID: 1700, AuthorityName: "source", DataScope: 3},
|
||||
{AuthorityID: 1701, AuthorityName: "target", DataScope: 3},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
for _, scope := range []int{-1, 6} {
|
||||
if err := repo.CreateAuthority(context.Background(), &biz.Authority{AuthorityID: uint(1800 + scope + 1), AuthorityName: "invalid-create", DataScope: scope}); !errors.Is(err, errInvalidDataScope) {
|
||||
t.Fatalf("CreateAuthority() scope %d error = %v", scope, err)
|
||||
}
|
||||
if err := repo.CopyAuthority(context.Background(), 1700, &biz.Authority{AuthorityID: uint(1900 + scope + 1), AuthorityName: "invalid-copy", DataScope: scope}); !errors.Is(err, errInvalidDataScope) {
|
||||
t.Fatalf("CopyAuthority() scope %d error = %v", scope, err)
|
||||
}
|
||||
if err := repo.UpdateAuthority(context.Background(), &biz.Authority{AuthorityID: 1701, AuthorityName: "invalid-update", DataScope: scope}); !errors.Is(err, errInvalidDataScope) {
|
||||
t.Fatalf("UpdateAuthority() scope %d error = %v", scope, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateAuthorityZeroDataScopeKeepsStoredValue(t *testing.T) {
|
||||
data := newTransactionTestData(t)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
if err := db.Create(&authorityPO{AuthorityID: 2000, AuthorityName: "before", DataScope: 3}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if err := (&authorityAccessRepo{data: data}).UpdateAuthority(context.Background(), &biz.Authority{AuthorityID: 2000, AuthorityName: "after", DataScope: 0}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var stored authorityPO
|
||||
if err := db.Where("authority_id = ?", 2000).First(&stored).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if stored.DataScope != 3 || stored.AuthorityName != "after" {
|
||||
t.Fatalf("updated authority = %+v, want data scope 3 and name after", stored)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ package data
|
|||
|
||||
import (
|
||||
"database/sql"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"reflect"
|
||||
"strings"
|
||||
|
|
@ -18,6 +19,11 @@ import (
|
|||
// Casbin, while ownership columns on business tables are row-level scope.
|
||||
type dataScopeAuditEnqueue func(dataAccessLogPO)
|
||||
|
||||
var (
|
||||
errDataScopeRequired = errors.New("受控表访问缺少数据权限上下文")
|
||||
errInvalidDataScope = errors.New("数据权限范围不合法")
|
||||
)
|
||||
|
||||
func registerDataScopeCallbacks(db *gorm.DB, enqueue dataScopeAuditEnqueue) {
|
||||
if db == nil {
|
||||
return
|
||||
|
|
@ -40,7 +46,7 @@ func registerDataScopeCallbacks(db *gorm.DB, enqueue dataScopeAuditEnqueue) {
|
|||
}
|
||||
c := db.Callback().Create()
|
||||
if c.Get("data_scope:stamp") == nil {
|
||||
_ = c.Before("gorm:create").Register("data_scope:stamp", stampOwnership)
|
||||
_ = c.Before("gorm:create").Register("data_scope:stamp", stampOwnership(enqueue))
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -52,6 +58,16 @@ func isControlledTable(db *gorm.DB) bool {
|
|||
return db.Statement.Schema != nil && !strings.HasPrefix(db.Statement.Table, "sys_") && (hasScopeField(db, "dept_id") || hasScopeField(db, "created_by"))
|
||||
}
|
||||
|
||||
func skipDataScope(db *gorm.DB) bool {
|
||||
skip, ok := db.Get("data_scope:skip")
|
||||
value, _ := skip.(bool)
|
||||
return ok && value
|
||||
}
|
||||
|
||||
func validDataScope(scope biz.DataScope) bool {
|
||||
return scope.UserID != 0 && scope.AuthorityID != 0 && scope.Scope >= 1 && scope.Scope <= 5 && scope.All == (scope.Scope == 1)
|
||||
}
|
||||
|
||||
func applyDataScope(operation string, enqueue dataScopeAuditEnqueue) func(*gorm.DB) {
|
||||
return func(db *gorm.DB) {
|
||||
if !isControlledTable(db) {
|
||||
|
|
@ -60,15 +76,19 @@ func applyDataScope(operation string, enqueue dataScopeAuditEnqueue) func(*gorm.
|
|||
if _, done := db.Statement.Clauses["data_scope:applied"]; done {
|
||||
return
|
||||
}
|
||||
if skip, ok := db.Get("data_scope:skip"); ok {
|
||||
if value, _ := skip.(bool); value {
|
||||
return
|
||||
}
|
||||
if skipDataScope(db) {
|
||||
return
|
||||
}
|
||||
scope, ok := biz.DataScopeFromContext(db.Statement.Context)
|
||||
if !ok {
|
||||
slog.WarnContext(db.Statement.Context, "数据权限: 业务表访问无身份上下文, 已放行(待补 ctx / 或使用系统上下文)", "mod", "data-scope", "table", db.Statement.Table)
|
||||
recordDataScopeEvent(db, enqueue, "no_identity", operation, "无身份上下文访问受控表, 已放行", biz.DataScope{})
|
||||
slog.WarnContext(db.Statement.Context, "数据权限: 业务表访问无身份上下文, 已拒绝", "mod", "data-scope", "table", db.Statement.Table)
|
||||
recordDataScopeEvent(db, enqueue, "no_identity", operation, "无身份上下文访问受控表, 已拒绝", biz.DataScope{})
|
||||
_ = db.AddError(errDataScopeRequired)
|
||||
return
|
||||
}
|
||||
if !validDataScope(scope) {
|
||||
recordDataScopeEvent(db, enqueue, "invalid_scope", operation, "数据权限上下文不合法, 已拒绝", scope)
|
||||
_ = db.AddError(errInvalidDataScope)
|
||||
return
|
||||
}
|
||||
if (operation == "update" || operation == "delete") && !db.AllowGlobalUpdate && !hasWriteConditions(db) {
|
||||
|
|
@ -86,10 +106,13 @@ func applyDataScope(operation string, enqueue dataScopeAuditEnqueue) func(*gorm.
|
|||
if hasScopeField(db, "dept_id") {
|
||||
ids := scope.DepartmentIDs
|
||||
if len(ids) == 0 {
|
||||
ids = []uint{0}
|
||||
db.Where("1 = 0")
|
||||
return
|
||||
}
|
||||
db.Where(table+".dept_id IN ?", ids)
|
||||
return
|
||||
}
|
||||
db.Where("1 = 0")
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -111,19 +134,28 @@ func auditBlockedWrite(operation string, enqueue dataScopeAuditEnqueue) func(*go
|
|||
}
|
||||
}
|
||||
|
||||
func stampOwnership(db *gorm.DB) {
|
||||
if !isControlledTable(db) {
|
||||
return
|
||||
}
|
||||
scope, ok := biz.DataScopeFromContext(db.Statement.Context)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if hasScopeField(db, "created_by") && scope.UserID != 0 {
|
||||
db.Statement.SetColumn("created_by", scope.UserID, true)
|
||||
}
|
||||
if hasScopeField(db, "dept_id") && scope.PrimaryDeptID != 0 {
|
||||
db.Statement.SetColumn("dept_id", scope.PrimaryDeptID, true)
|
||||
func stampOwnership(enqueue dataScopeAuditEnqueue) func(*gorm.DB) {
|
||||
return func(db *gorm.DB) {
|
||||
if !isControlledTable(db) || skipDataScope(db) {
|
||||
return
|
||||
}
|
||||
scope, ok := biz.DataScopeFromContext(db.Statement.Context)
|
||||
if !ok {
|
||||
recordDataScopeEvent(db, enqueue, "no_identity", "create", "无身份上下文访问受控表, 已拒绝", biz.DataScope{})
|
||||
_ = db.AddError(errDataScopeRequired)
|
||||
return
|
||||
}
|
||||
if !validDataScope(scope) {
|
||||
recordDataScopeEvent(db, enqueue, "invalid_scope", "create", "数据权限上下文不合法, 已拒绝", scope)
|
||||
_ = db.AddError(errInvalidDataScope)
|
||||
return
|
||||
}
|
||||
if hasScopeField(db, "created_by") {
|
||||
db.Statement.SetColumn("created_by", scope.UserID, true)
|
||||
}
|
||||
if hasScopeField(db, "dept_id") {
|
||||
db.Statement.SetColumn("dept_id", scope.PrimaryDeptID, true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -19,6 +19,12 @@ func openDataScopeAuditTestDB(t *testing.T, name string) *gorm.DB {
|
|||
if err = db.AutoMigrate(&dataAccessLogPO{}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// The in-memory SQLite URI is intentionally named so the audit writer can
|
||||
// use a second GORM handle. Clear rows left by a prior -count iteration so
|
||||
// each test starts with an isolated logical database.
|
||||
if err = db.Unscoped().Where("1 = 1").Delete(&dataAccessLogPO{}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return db
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,106 @@
|
|||
package data
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"kra/internal/biz"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type dataScopeRecord struct {
|
||||
ID uint `gorm:"primaryKey"`
|
||||
DeptID uint
|
||||
CreatedBy uint
|
||||
Name string
|
||||
}
|
||||
|
||||
func (dataScopeRecord) TableName() string { return "business_scope_records" }
|
||||
|
||||
func newDataScopeTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
db, err := openWithDriver("sqlite", "file:"+t.Name()+"?mode=memory&cache=shared")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AutoMigrate(&dataScopeRecord{}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
registerDataScopeCallbacks(db, nil)
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = sqlDB.Close() })
|
||||
return db
|
||||
}
|
||||
|
||||
func TestDataScopeCallbacksFailClosedAndAllowExplicitSystemBypass(t *testing.T) {
|
||||
db := newDataScopeTestDB(t)
|
||||
seed := []dataScopeRecord{
|
||||
{DeptID: 10, CreatedBy: 7, Name: "visible"},
|
||||
{DeptID: 20, CreatedBy: 8, Name: "hidden"},
|
||||
{DeptID: 0, CreatedBy: 9, Name: "unassigned"},
|
||||
}
|
||||
if err := db.Set("data_scope:skip", true).Create(&seed).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var rows []dataScopeRecord
|
||||
if err := db.WithContext(context.Background()).Find(&rows).Error; !errors.Is(err, errDataScopeRequired) {
|
||||
t.Fatalf("query without data scope error = %v", err)
|
||||
}
|
||||
|
||||
invalidCtx := biz.NewDataScopeContext(context.Background(), biz.DataScope{UserID: 7, AuthorityID: 1, Scope: 1, All: false})
|
||||
if err := db.WithContext(invalidCtx).Find(&rows).Error; !errors.Is(err, errInvalidDataScope) {
|
||||
t.Fatalf("query with invalid data scope error = %v", err)
|
||||
}
|
||||
|
||||
scopedCtx := biz.NewDataScopeContext(context.Background(), biz.DataScope{UserID: 7, AuthorityID: 1, Scope: 3, DepartmentIDs: []uint{10}})
|
||||
rows = nil
|
||||
if err := db.WithContext(scopedCtx).Order("id").Find(&rows).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(rows) != 1 || rows[0].Name != "visible" {
|
||||
t.Fatalf("scoped query rows = %+v", rows)
|
||||
}
|
||||
|
||||
emptyCtx := biz.NewDataScopeContext(context.Background(), biz.DataScope{UserID: 7, AuthorityID: 1, Scope: 5})
|
||||
rows = nil
|
||||
if err := db.WithContext(emptyCtx).Find(&rows).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(rows) != 0 {
|
||||
t.Fatalf("empty custom scope exposed rows: %+v", rows)
|
||||
}
|
||||
|
||||
rows = nil
|
||||
if err := db.Set("data_scope:skip", true).Order("id").Find(&rows).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(rows) != len(seed) {
|
||||
t.Fatalf("explicit system bypass rows = %d, want %d", len(rows), len(seed))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDataScopeCreateRequiresIdentityAndStampsOwnership(t *testing.T) {
|
||||
db := newDataScopeTestDB(t)
|
||||
if err := db.Create(&dataScopeRecord{Name: "missing"}).Error; !errors.Is(err, errDataScopeRequired) {
|
||||
t.Fatalf("create without data scope error = %v", err)
|
||||
}
|
||||
|
||||
ctx := biz.NewDataScopeContext(context.Background(), biz.DataScope{UserID: 7, AuthorityID: 1, Scope: 3, PrimaryDeptID: 10, DepartmentIDs: []uint{10}})
|
||||
created := dataScopeRecord{Name: "owned", DeptID: 999, CreatedBy: 999}
|
||||
if err := db.WithContext(ctx).Create(&created).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var stored dataScopeRecord
|
||||
if err := db.Set("data_scope:skip", true).First(&stored, created.ID).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if stored.DeptID != 10 || stored.CreatedBy != 7 {
|
||||
t.Fatalf("stamped ownership = dept:%d user:%d, want 10/7", stored.DeptID, stored.CreatedBy)
|
||||
}
|
||||
}
|
||||
|
|
@ -97,7 +97,10 @@ func (r *mediaRepo) DeleteUploadSession(ctx context.Context, id uint) error {
|
|||
}
|
||||
func (r *mediaRepo) UpsertChunk(ctx context.Context, uploadID uint, v *biz.UploadChunk) error {
|
||||
po := uploadChunkPO{UploadID: uploadID, ChunkIndex: v.Index, ChunkHash: v.Hash, Size: v.Size}
|
||||
return r.data.gormDB.WithContext(ctx).Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "upload_id"}, {Name: "chunk_index"}}, DoNothing: true}).Create(&po).Error
|
||||
return r.data.gormDB.WithContext(ctx).Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "upload_id"}, {Name: "chunk_index"}},
|
||||
DoUpdates: clause.AssignmentColumns([]string{"chunk_hash", "size", "updated_at", "deleted_at"}),
|
||||
}).Create(&po).Error
|
||||
}
|
||||
func (r *mediaRepo) ListChunks(ctx context.Context, uploadID uint) ([]*biz.UploadChunk, error) {
|
||||
var pos []uploadChunkPO
|
||||
|
|
|
|||
|
|
@ -263,27 +263,8 @@ func (r *menuRepo) SetAuthorityMenus(ctx context.Context, id uint, ids []uint) e
|
|||
if err := r.data.gormDB.WithContext(ctx).Where("authority_id = ?", id).First(&authorityPO{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
config := r.data.runtime.Admin()
|
||||
if actor, ok := biz.ActorFromContext(ctx); ok && config != nil && config.System != nil && config.System.UseStrictAuth {
|
||||
var authority authorityPO
|
||||
if err := r.data.gormDB.WithContext(ctx).Where("authority_id = ?", actor.AuthorityID).First(&authority).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if authority.ParentID != nil && *authority.ParentID != 0 {
|
||||
var allowedIDs []uint
|
||||
if err := r.data.gormDB.WithContext(ctx).Model(&authorityMenuPO{}).Where("sys_authority_authority_id = ?", actor.AuthorityID).Pluck("sys_base_menu_id", &allowedIDs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
allowed := make(map[uint]bool, len(allowedIDs))
|
||||
for _, menuID := range allowedIDs {
|
||||
allowed[menuID] = true
|
||||
}
|
||||
for _, menuID := range ids {
|
||||
if !allowed[menuID] {
|
||||
return errors.New("添加失败,请勿跨级操作")
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := r.checkMenuAssignmentAuth(ctx, ids); err != nil {
|
||||
return err
|
||||
}
|
||||
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Where("sys_authority_authority_id = ?", id).Delete(&authorityMenuPO{}).Error; err != nil {
|
||||
|
|
@ -346,8 +327,29 @@ func (r *menuRepo) DefaultRouterRoleIDs(ctx context.Context, id uint) ([]uint, e
|
|||
}
|
||||
|
||||
func (r *menuRepo) SetMenuRoles(ctx context.Context, id uint, ids []uint) error {
|
||||
access := &authorityAccessRepo{data: r.data}
|
||||
_, allowedAuthorities, strict, err := access.strictAuthorityAccess(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := r.checkMenuAssignmentAuth(ctx, []uint{id}); err != nil {
|
||||
return err
|
||||
}
|
||||
if strict {
|
||||
for _, authorityID := range ids {
|
||||
if !allowedAuthorities[authorityID] {
|
||||
return errors.New("您提交的角色ID不合法")
|
||||
}
|
||||
}
|
||||
}
|
||||
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Where("sys_base_menu_id = ?", id).Delete(&authorityMenuPO{}).Error; err != nil {
|
||||
if strict {
|
||||
if len(allowedAuthorities) > 0 {
|
||||
if err := tx.Where("sys_base_menu_id = ? AND sys_authority_authority_id IN ?", id, authorityIDs(allowedAuthorities)).Delete(&authorityMenuPO{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
} else if err := tx.Where("sys_base_menu_id = ?", id).Delete(&authorityMenuPO{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
links := make([]authorityMenuPO, 0, len(ids))
|
||||
|
|
@ -360,3 +362,54 @@ func (r *menuRepo) SetMenuRoles(ctx context.Context, id uint, ids []uint) error
|
|||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func authorityIDs(allowed map[uint]bool) []uint {
|
||||
ids := make([]uint, 0, len(allowed))
|
||||
for id := range allowed {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
func (r *menuRepo) checkMenuAssignmentAuth(ctx context.Context, menuIDs []uint) error {
|
||||
actor, allowedAuthorities, strict, err := (&authorityAccessRepo{data: r.data}).strictAuthorityAccess(ctx)
|
||||
if err != nil || !strict {
|
||||
return err
|
||||
}
|
||||
return checkMenuAssignment(r.data.gormDB.WithContext(ctx), actor.AuthorityID, allowedAuthorities[actor.AuthorityID], menuIDs)
|
||||
}
|
||||
|
||||
func checkMenuAssignment(db *gorm.DB, actorID uint, root bool, menuIDs []uint) error {
|
||||
if len(menuIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
unique := make(map[uint]struct{}, len(menuIDs))
|
||||
for _, menuID := range menuIDs {
|
||||
unique[menuID] = struct{}{}
|
||||
}
|
||||
if root {
|
||||
var count int64
|
||||
if err := db.Model(&menuPO{}).Where("id IN ?", menuIDs).Count(&count).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if count != int64(len(unique)) {
|
||||
return errors.New("添加失败,菜单不存在")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
var allowedIDs []uint
|
||||
if err := db.Model(&authorityMenuPO{}).
|
||||
Where("sys_authority_authority_id = ?", actorID).Pluck("sys_base_menu_id", &allowedIDs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
allowed := make(map[uint]bool, len(allowedIDs))
|
||||
for _, menuID := range allowedIDs {
|
||||
allowed[menuID] = true
|
||||
}
|
||||
for menuID := range unique {
|
||||
if !allowed[menuID] {
|
||||
return errors.New("添加失败,请勿跨级操作")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -139,3 +139,46 @@ func TestMenuUpdateRebuildsRelationsWithoutChangingRequestedIdentity(t *testing.
|
|||
t.Fatalf("authority-button links = %d, want 1", links)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetMenuRolesStrictOnlyChangesManagedAuthorities(t *testing.T) {
|
||||
data := newMenuTestData(t)
|
||||
enableStrictAuthorityTestMode(data)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
rootID, actorID, childID, siblingID := uint(888), uint(4001), uint(4002), uint(5001)
|
||||
if err := db.Create(&[]authorityPO{
|
||||
{AuthorityID: rootID, ParentID: authorityUintPointer(0)},
|
||||
{AuthorityID: actorID, ParentID: &rootID},
|
||||
{AuthorityID: childID, ParentID: &actorID},
|
||||
{AuthorityID: siblingID, ParentID: &rootID},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
menu := menuPO{ID: 10, Name: "managed", Path: "managed"}
|
||||
if err := db.Create(&menu).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Create(&[]authorityMenuPO{
|
||||
{SysAuthorityAuthorityID: actorID, SysBaseMenuID: menu.ID},
|
||||
{SysAuthorityAuthorityID: siblingID, SysBaseMenuID: menu.ID},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx := biz.NewActorContext(context.Background(), biz.Actor{AuthorityID: actorID})
|
||||
repo := &menuRepo{data: data}
|
||||
if err := repo.SetMenuRoles(ctx, menu.ID, []uint{siblingID}); err == nil {
|
||||
t.Fatal("SetMenuRoles() accepted an out-of-scope authority")
|
||||
}
|
||||
if err := repo.SetMenuRoles(ctx, menu.ID, []uint{childID}); err != nil {
|
||||
t.Fatalf("SetMenuRoles() rejected a managed authority: %v", err)
|
||||
}
|
||||
var siblingCount, childCount int64
|
||||
if err := db.Model(&authorityMenuPO{}).Where("sys_authority_authority_id = ? AND sys_base_menu_id = ?", siblingID, menu.ID).Count(&siblingCount).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Model(&authorityMenuPO{}).Where("sys_authority_authority_id = ? AND sys_base_menu_id = ?", childID, menu.ID).Count(&childCount).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if siblingCount != 1 || childCount != 1 {
|
||||
t.Fatalf("menu-role counts sibling/child = %d/%d, want 1/1", siblingCount, childCount)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -46,6 +46,12 @@ func (r *permissionRepo) Buttons(ctx context.Context, menuID uint) ([]*biz.MenuB
|
|||
return out, nil
|
||||
}
|
||||
func (r *permissionRepo) SetAuthorityButtons(ctx context.Context, aid uint, buttons map[uint][]uint) error {
|
||||
if err := (&authorityAccessRepo{data: r.data}).checkAuthorityIDAuth(ctx, aid); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := r.checkButtonAssignmentAuth(ctx, buttons); err != nil {
|
||||
return err
|
||||
}
|
||||
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Where("authority_id = ?", aid).Delete(&authorityButtonPO{}).Error; err != nil {
|
||||
return err
|
||||
|
|
@ -74,6 +80,12 @@ func (r *permissionRepo) SelectedButtons(ctx context.Context, aid, menuID uint)
|
|||
return ids, nil
|
||||
}
|
||||
func (r *permissionRepo) SetSelectedButtons(ctx context.Context, aid, menuID uint, ids []uint) error {
|
||||
if err := (&authorityAccessRepo{data: r.data}).checkAuthorityIDAuth(ctx, aid); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := r.checkButtonAssignmentAuth(ctx, map[uint][]uint{menuID: ids}); err != nil {
|
||||
return err
|
||||
}
|
||||
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Where("authority_id = ? AND sys_menu_id = ?", aid, menuID).Delete(&authorityButtonPO{}).Error; err != nil {
|
||||
return err
|
||||
|
|
@ -104,3 +116,53 @@ func (r *permissionRepo) AuthorityButtonIDs(ctx context.Context, aid uint) ([]ui
|
|||
err := r.data.gormDB.WithContext(ctx).Model(&authorityButtonPO{}).Where("authority_id = ?", aid).Pluck("sys_base_menu_btn_id", &ids).Error
|
||||
return ids, err
|
||||
}
|
||||
|
||||
func (r *permissionRepo) checkButtonAssignmentAuth(ctx context.Context, buttons map[uint][]uint) error {
|
||||
actor, allowedAuthorities, strict, err := (&authorityAccessRepo{data: r.data}).strictAuthorityAccess(ctx)
|
||||
if err != nil || !strict {
|
||||
return err
|
||||
}
|
||||
return checkButtonAssignment(r.data.gormDB.WithContext(ctx), actor.AuthorityID, allowedAuthorities[actor.AuthorityID], buttons)
|
||||
}
|
||||
|
||||
func checkButtonAssignment(db *gorm.DB, actorID uint, root bool, buttons map[uint][]uint) error {
|
||||
requested := make(map[[2]uint]struct{})
|
||||
buttonIDs := make([]uint, 0)
|
||||
for menuID, ids := range buttons {
|
||||
for _, buttonID := range ids {
|
||||
key := [2]uint{menuID, buttonID}
|
||||
if _, exists := requested[key]; exists {
|
||||
continue
|
||||
}
|
||||
requested[key] = struct{}{}
|
||||
buttonIDs = append(buttonIDs, buttonID)
|
||||
}
|
||||
}
|
||||
if len(requested) == 0 {
|
||||
return nil
|
||||
}
|
||||
allowed := make(map[[2]uint]bool, len(requested))
|
||||
if root {
|
||||
var rows []menuButtonPO
|
||||
if err := db.Where("id IN ?", buttonIDs).Find(&rows).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, row := range rows {
|
||||
allowed[[2]uint{row.MenuID, row.ID}] = true
|
||||
}
|
||||
} else {
|
||||
var rows []authorityButtonPO
|
||||
if err := db.Where("authority_id = ?", actorID).Find(&rows).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, row := range rows {
|
||||
allowed[[2]uint{row.MenuID, row.ButtonID}] = true
|
||||
}
|
||||
}
|
||||
for key := range requested {
|
||||
if !allowed[key] {
|
||||
return errors.New("添加失败,请勿跨级操作")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import (
|
|||
"context"
|
||||
"testing"
|
||||
|
||||
"kra/internal/biz"
|
||||
"kra/internal/conf"
|
||||
)
|
||||
|
||||
|
|
@ -81,3 +82,41 @@ func TestPermissionSetSelectedButtonsPreservesDuplicates(t *testing.T) {
|
|||
t.Fatalf("selected = %#v, want duplicate button IDs preserved", selected)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetSelectedButtonsStrictRequiresManagedRoleAndOwnedButton(t *testing.T) {
|
||||
data := newMenuTestData(t)
|
||||
enableStrictAuthorityTestMode(data)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
rootID, actorID, childID, siblingID := uint(888), uint(6001), uint(6002), uint(7001)
|
||||
if err := db.Create(&[]authorityPO{
|
||||
{AuthorityID: rootID, ParentID: authorityUintPointer(0)},
|
||||
{AuthorityID: actorID, ParentID: &rootID},
|
||||
{AuthorityID: childID, ParentID: &actorID},
|
||||
{AuthorityID: siblingID, ParentID: &rootID},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Create(&menuPO{ID: 10, Name: "managed", Path: "managed"}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Create(&[]menuButtonPO{
|
||||
{ID: 31, MenuID: 10, Name: "owned"},
|
||||
{ID: 32, MenuID: 10, Name: "unowned"},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Create(&authorityButtonPO{AuthorityID: actorID, MenuID: 10, ButtonID: 31}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx := biz.NewActorContext(context.Background(), biz.Actor{AuthorityID: actorID})
|
||||
repo := &permissionRepo{data: data}
|
||||
if err := repo.SetSelectedButtons(ctx, siblingID, 10, []uint{31}); err == nil {
|
||||
t.Fatal("SetSelectedButtons() accepted an out-of-scope authority")
|
||||
}
|
||||
if err := repo.SetSelectedButtons(ctx, childID, 10, []uint{32}); err == nil {
|
||||
t.Fatal("SetSelectedButtons() accepted a button not assigned to the actor")
|
||||
}
|
||||
if err := repo.SetSelectedButtons(ctx, childID, 10, []uint{31}); err != nil {
|
||||
t.Fatalf("SetSelectedButtons() rejected a managed assignment: %v", err)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -89,7 +89,11 @@ func (i *tokenIssuer) IssueToken(user *biz.User, authorityID uint, mustChangePas
|
|||
if expires <= 0 {
|
||||
expires = settings.Expires
|
||||
}
|
||||
token, claims, err := adminauth.Generate(settings.SigningKey, settings.Issuer, expires, settings.Buffer, user.ID, authorityID, user.UUID, user.Username, user.NickName, mustChangePassword)
|
||||
passwordVersion := int64(0)
|
||||
if user.PasswordUpdatedAt != nil {
|
||||
passwordVersion = user.PasswordUpdatedAt.UnixNano()
|
||||
}
|
||||
token, claims, err := adminauth.Generate(settings.SigningKey, settings.Issuer, expires, settings.Buffer, user.ID, authorityID, user.UUID, user.Username, user.NickName, mustChangePassword, passwordVersion)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
|
@ -102,17 +106,19 @@ func (i *tokenIssuer) ReissueToken(source *biz.AuthClaims, authorityID uint) (*b
|
|||
}
|
||||
settings := i.settings.JWTSettings()
|
||||
claims := &adminauth.Claims{
|
||||
UUID: source.UUID,
|
||||
ID: source.ID,
|
||||
Username: source.Username,
|
||||
NickName: source.NickName,
|
||||
AuthorityID: authorityID,
|
||||
BufferTime: int64(source.BufferTime / time.Second),
|
||||
UserType: source.UserType,
|
||||
MustChangePwd: source.MustChangePwd,
|
||||
UUID: source.UUID,
|
||||
ID: source.ID,
|
||||
Username: source.Username,
|
||||
NickName: source.NickName,
|
||||
AuthorityID: authorityID,
|
||||
BufferTime: int64(source.BufferTime / time.Second),
|
||||
UserType: source.UserType,
|
||||
MustChangePwd: source.MustChangePwd,
|
||||
PasswordVersion: source.PasswordVersion,
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
Audience: jwt.ClaimStrings(append([]string(nil), source.Audience...)),
|
||||
Issuer: source.Issuer,
|
||||
IssuedAt: jwt.NewNumericDate(source.IssuedAt),
|
||||
NotBefore: jwt.NewNumericDate(source.NotBefore),
|
||||
ExpiresAt: jwt.NewNumericDate(source.ExpiresAt),
|
||||
},
|
||||
|
|
@ -142,5 +148,9 @@ func (i *tokenIssuer) ParseToken(token string) (*biz.AuthClaims, error) {
|
|||
}
|
||||
audience := make([]string, len(claims.Audience))
|
||||
copy(audience, claims.Audience)
|
||||
return &biz.AuthClaims{UUID: claims.UUID, ID: claims.ID, Username: claims.Username, NickName: claims.NickName, AuthorityID: claims.AuthorityID, UserType: claims.UserType, BufferTime: time.Duration(claims.BufferTime) * time.Second, MustChangePwd: claims.MustChangePwd, Issuer: claims.Issuer, Audience: audience, NotBefore: claims.NotBefore.Time, ExpiresAt: claims.ExpiresAt.Time}, nil
|
||||
issuedAt := time.Time{}
|
||||
if claims.IssuedAt != nil {
|
||||
issuedAt = claims.IssuedAt.Time
|
||||
}
|
||||
return &biz.AuthClaims{UUID: claims.UUID, ID: claims.ID, Username: claims.Username, NickName: claims.NickName, AuthorityID: claims.AuthorityID, UserType: claims.UserType, BufferTime: time.Duration(claims.BufferTime) * time.Second, MustChangePwd: claims.MustChangePwd, PasswordVersion: claims.PasswordVersion, Issuer: claims.Issuer, Audience: audience, IssuedAt: issuedAt, NotBefore: claims.NotBefore.Time, ExpiresAt: claims.ExpiresAt.Time}, nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -318,6 +318,10 @@ func (r *userRepo) CreateUser(ctx context.Context, user *biz.User) (*biz.User, e
|
|||
}
|
||||
|
||||
func (r *userRepo) CreateUserWithAuthorities(ctx context.Context, user *biz.User, authorityIDs []uint) (*biz.User, error) {
|
||||
requestedAuthorities := append([]uint{user.AuthorityID}, authorityIDs...)
|
||||
if err := (&authorityAccessRepo{data: r.data}).checkAuthorityIDsAuth(ctx, requestedAuthorities); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if user.UUID == "" {
|
||||
user.UUID = uuid.NewString()
|
||||
}
|
||||
|
|
@ -377,9 +381,23 @@ func (r *userRepo) CreateUserWithAuthorities(ctx context.Context, user *biz.User
|
|||
}
|
||||
|
||||
func (r *userRepo) UpdateUser(ctx context.Context, user *biz.User) error {
|
||||
if err := r.checkUserUpdateAuth(ctx, user); err != nil {
|
||||
return err
|
||||
}
|
||||
return r.updateUser(r.data.gormDB.WithContext(ctx), user)
|
||||
}
|
||||
|
||||
func (r *userRepo) checkUserUpdateAuth(ctx context.Context, user *biz.User) error {
|
||||
access := &authorityAccessRepo{data: r.data}
|
||||
if err := access.checkUserIDAuth(ctx, user.ID, false); err != nil {
|
||||
return err
|
||||
}
|
||||
if user.AuthorityID != 0 {
|
||||
return access.checkAuthorityIDAuth(ctx, user.AuthorityID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *userRepo) UpdateSelfUser(ctx context.Context, user *biz.User) error {
|
||||
updates := make(map[string]any)
|
||||
if user.NickName != "" {
|
||||
|
|
@ -453,6 +471,13 @@ func (r *userRepo) updateUser(tx *gorm.DB, user *biz.User) error {
|
|||
}
|
||||
|
||||
func (r *userRepo) UpdateUserWithAuthorities(ctx context.Context, user *biz.User, authorityIDs []uint) error {
|
||||
access := &authorityAccessRepo{data: r.data}
|
||||
if err := r.checkUserUpdateAuth(ctx, user); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := access.checkAuthorityIDsAuth(ctx, authorityIDs); err != nil {
|
||||
return err
|
||||
}
|
||||
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := r.updateUser(tx, user); err != nil {
|
||||
return err
|
||||
|
|
@ -464,6 +489,9 @@ func (r *userRepo) UpdateUserWithAuthorities(ctx context.Context, user *biz.User
|
|||
})
|
||||
}
|
||||
func (r *userRepo) DeleteUser(ctx context.Context, id uint) error {
|
||||
if err := (&authorityAccessRepo{data: r.data}).checkUserIDAuth(ctx, id, false); err != nil {
|
||||
return err
|
||||
}
|
||||
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Where("id = ?", id).Delete(&userPO{}).Error; err != nil {
|
||||
return err
|
||||
|
|
@ -481,6 +509,9 @@ func (r *userRepo) DeleteUser(ctx context.Context, id uint) error {
|
|||
})
|
||||
}
|
||||
func (r *userRepo) UpdatePassword(ctx context.Context, id uint, password string, clearMustChange bool) error {
|
||||
if err := (&authorityAccessRepo{data: r.data}).checkUserIDAuth(ctx, id, true); err != nil {
|
||||
return err
|
||||
}
|
||||
now := time.Now()
|
||||
updates := map[string]any{"password": password, "password_updated_at": now}
|
||||
if clearMustChange {
|
||||
|
|
@ -502,17 +533,18 @@ func (r *userRepo) ListAuthorities(ctx context.Context) ([]*biz.Authority, error
|
|||
}
|
||||
|
||||
func (r *userRepo) SetUserAuthorities(ctx context.Context, id uint, authorityIDs []uint) error {
|
||||
access := &authorityAccessRepo{data: r.data}
|
||||
if err := access.checkUserIDAuth(ctx, id, false); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := access.checkAuthorityIDsAuth(ctx, authorityIDs); err != nil {
|
||||
return err
|
||||
}
|
||||
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var user userPO
|
||||
if err := tx.Where("id = ?", id).First(&user).Error; err != nil {
|
||||
return errors.New("查询用户数据失败")
|
||||
}
|
||||
access := &authorityAccessRepo{data: r.data}
|
||||
for _, authorityID := range authorityIDs {
|
||||
if err := access.checkAuthorityIDAuth(ctx, authorityID); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return setUserAuthorities(tx, id, authorityIDs)
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,95 @@
|
|||
package data
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"kra/internal/biz"
|
||||
)
|
||||
|
||||
func seedStrictUserTree(t *testing.T, data *Data) (actorID, childID, siblingID uint, users []userPO) {
|
||||
t.Helper()
|
||||
enableStrictAuthorityTestMode(data)
|
||||
db := data.gormDB.WithContext(context.Background())
|
||||
rootID := uint(888)
|
||||
actorID, childID, siblingID = 8001, 8002, 9001
|
||||
if err := db.Create(&[]authorityPO{
|
||||
{AuthorityID: rootID, ParentID: authorityUintPointer(0)},
|
||||
{AuthorityID: actorID, ParentID: &rootID},
|
||||
{AuthorityID: childID, ParentID: &actorID},
|
||||
{AuthorityID: siblingID, ParentID: &rootID},
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
users = []userPO{
|
||||
{Username: "actor-user", Password: "hash", AuthorityID: actorID, Enable: 1},
|
||||
{Username: "managed-user", Password: "hash", AuthorityID: childID, Enable: 1},
|
||||
{Username: "outside-user", Password: "hash", AuthorityID: siblingID, Enable: 1},
|
||||
}
|
||||
if err := db.Create(&users).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return actorID, childID, siblingID, users
|
||||
}
|
||||
|
||||
func TestStrictUserMutationsRejectOutsideTargetAndAllowSelfPassword(t *testing.T) {
|
||||
data := newTransactionTestData(t)
|
||||
actorID, _, _, users := seedStrictUserTree(t, data)
|
||||
ctx := biz.NewActorContext(context.Background(), biz.Actor{UserID: users[0].ID, AuthorityID: actorID})
|
||||
repo := &userRepo{data: data}
|
||||
|
||||
if err := repo.UpdateUser(ctx, &biz.User{ID: users[2].ID, NickName: "changed"}); err == nil {
|
||||
t.Fatal("UpdateUser() accepted an out-of-scope target")
|
||||
}
|
||||
if err := repo.DeleteUser(ctx, users[2].ID); err == nil {
|
||||
t.Fatal("DeleteUser() accepted an out-of-scope target")
|
||||
}
|
||||
if err := repo.UpdatePassword(ctx, users[2].ID, "new-hash", false); err == nil {
|
||||
t.Fatal("UpdatePassword() accepted an out-of-scope target")
|
||||
}
|
||||
if err := repo.UpdatePassword(ctx, users[0].ID, "self-hash", true); err != nil {
|
||||
t.Fatalf("UpdatePassword() rejected the current user: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStrictUserCreationAndRoleAssignmentRequireManagedAuthorities(t *testing.T) {
|
||||
data := newTransactionTestData(t)
|
||||
actorID, childID, siblingID, users := seedStrictUserTree(t, data)
|
||||
ctx := biz.NewActorContext(context.Background(), biz.Actor{UserID: users[0].ID, AuthorityID: actorID})
|
||||
repo := &userRepo{data: data}
|
||||
|
||||
if _, err := repo.CreateUserWithAuthorities(ctx, &biz.User{Username: "outside-create", Password: "hash", AuthorityID: siblingID, Enable: 1}, nil); err == nil {
|
||||
t.Fatal("CreateUserWithAuthorities() accepted an out-of-scope primary authority")
|
||||
}
|
||||
if _, err := repo.CreateUserWithAuthorities(ctx, &biz.User{Username: "managed-create", Password: "hash", AuthorityID: childID, Enable: 1}, []uint{childID}); err != nil {
|
||||
t.Fatalf("CreateUserWithAuthorities() rejected managed authorities: %v", err)
|
||||
}
|
||||
if err := repo.SetUserAuthorities(ctx, users[1].ID, []uint{siblingID}); err == nil {
|
||||
t.Fatal("SetUserAuthorities() accepted an out-of-scope authority")
|
||||
}
|
||||
if err := repo.SetUserAuthorities(ctx, users[2].ID, []uint{childID}); err == nil {
|
||||
t.Fatal("SetUserAuthorities() accepted an out-of-scope user target")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStrictUserUpdateRejectsPrimaryAuthorityOutsideManagedTree(t *testing.T) {
|
||||
data := newTransactionTestData(t)
|
||||
actorID, childID, siblingID, users := seedStrictUserTree(t, data)
|
||||
ctx := biz.NewActorContext(context.Background(), biz.Actor{UserID: users[0].ID, AuthorityID: actorID})
|
||||
repo := &userRepo{data: data}
|
||||
|
||||
if err := repo.UpdateUser(ctx, &biz.User{ID: users[1].ID, NickName: "changed", AuthorityID: siblingID}); err == nil {
|
||||
t.Fatal("UpdateUser() accepted a primary authority outside the managed tree")
|
||||
}
|
||||
if err := repo.UpdateUserWithAuthorities(ctx, &biz.User{ID: users[1].ID, NickName: "changed", AuthorityID: siblingID}, []uint{childID}); err == nil {
|
||||
t.Fatal("UpdateUserWithAuthorities() accepted a primary authority outside the managed tree")
|
||||
}
|
||||
|
||||
var stored userPO
|
||||
if err := data.gormDB.WithContext(context.Background()).First(&stored, users[1].ID).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if stored.AuthorityID != childID || stored.NickName == "changed" {
|
||||
t.Fatalf("rejected update changed user: %+v", stored)
|
||||
}
|
||||
}
|
||||
|
|
@ -24,6 +24,9 @@ import (
|
|||
func NewGinEngine(runtime *conf.Runtime, access *service.AccessControlService, handlers *handler.Set, auth *service.AuthService, security *service.SecurityService, audit *service.AuditRecorder, logger *slog.Logger, version string) *gin.Engine {
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
engine := gin.New()
|
||||
if err := engine.SetTrustedProxies(nil); err != nil && logger != nil {
|
||||
logger.Error("disable trusted proxies failed", "error", err)
|
||||
}
|
||||
engine.Use(servermiddleware.RequestMeta(), servermiddleware.Recovery(logger), servermiddleware.AccessLog(runtime, logger, version), servermiddleware.CORS(runtime), servermiddleware.ErrorAudit(logger), servermiddleware.SecurityRateLimit(security))
|
||||
|
||||
prefix := ""
|
||||
|
|
|
|||
|
|
@ -13,6 +13,8 @@ import (
|
|||
|
||||
"kra/internal/conf"
|
||||
"kra/internal/server/handler"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func emptyHandlers() *handler.Set {
|
||||
|
|
@ -99,6 +101,29 @@ func TestSwaggerUsesRootBasePathWithoutRouterPrefix(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestGinDoesNotTrustForwardedIPByDefault(t *testing.T) {
|
||||
engine := NewGinEngine(conf.NewRuntime(nil, &conf.AdminBackend{}), nil, emptyHandlers(), nil, nil, nil, nil, "test")
|
||||
engine.GET("/client-ip", func(c *gin.Context) { c.String(http.StatusOK, c.ClientIP()) })
|
||||
request := httptest.NewRequest(http.MethodGet, "/client-ip", nil)
|
||||
request.RemoteAddr = "203.0.113.10:4321"
|
||||
request.Header.Set("X-Forwarded-For", "198.51.100.20")
|
||||
response := httptest.NewRecorder()
|
||||
|
||||
engine.ServeHTTP(response, request)
|
||||
if response.Code != http.StatusOK || response.Body.String() != "203.0.113.10" {
|
||||
t.Fatalf("client IP = %q, status=%d", response.Body.String(), response.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnnouncementDataSourceRequiresAuthentication(t *testing.T) {
|
||||
engine := NewGinEngine(conf.NewRuntime(nil, &conf.AdminBackend{}), nil, emptyHandlers(), nil, nil, nil, nil, "test")
|
||||
response := httptest.NewRecorder()
|
||||
engine.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/info/getInfoDataSource", nil))
|
||||
if response.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalStorageResponseHeaders(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
for name, body := range map[string]string{"script.html": "<script>alert(1)</script>", "image.png": "png"} {
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
package handler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"kra/internal/server/httpx"
|
||||
|
|
@ -15,13 +17,40 @@ type Media struct{ service *service.MediaService }
|
|||
|
||||
func NewMedia(service *service.MediaService) *Media { return &Media{service: service} }
|
||||
|
||||
func (h *Media) limitMultipartBody(c *gin.Context) int64 {
|
||||
limit := h.service.MediaConfig().EffectiveMaxFileSize()
|
||||
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, limit+(1<<20))
|
||||
return limit
|
||||
}
|
||||
|
||||
func rejectMediaTooLarge(c *gin.Context, err error, message string) bool {
|
||||
var tooLarge *http.MaxBytesError
|
||||
if !errors.As(err, &tooLarge) {
|
||||
return false
|
||||
}
|
||||
c.AbortWithStatusJSON(http.StatusRequestEntityTooLarge, httpx.Response{Code: httpx.CodeError, Data: gin.H{}, Msg: message})
|
||||
return true
|
||||
}
|
||||
|
||||
func (h *Media) Upload(c *gin.Context) {
|
||||
claims := servermiddleware.Claims(c)
|
||||
header, err := c.FormFile("file")
|
||||
if claims == nil || err != nil {
|
||||
if claims == nil {
|
||||
httpx.Fail(c, "接收文件失败")
|
||||
return
|
||||
}
|
||||
limit := h.limitMultipartBody(c)
|
||||
header, err := c.FormFile("file")
|
||||
if err != nil {
|
||||
if rejectMediaTooLarge(c, err, "文件超过大小上限") {
|
||||
return
|
||||
}
|
||||
httpx.Fail(c, "接收文件失败")
|
||||
return
|
||||
}
|
||||
if header.Size > limit {
|
||||
c.AbortWithStatusJSON(http.StatusRequestEntityTooLarge, httpx.Response{Code: httpx.CodeError, Data: gin.H{}, Msg: "文件超过大小上限"})
|
||||
return
|
||||
}
|
||||
opened, err := header.Open()
|
||||
if err != nil {
|
||||
// Let the upload service surface this as the generic upload
|
||||
|
|
@ -191,13 +220,33 @@ func (h *Media) InitUpload(c *gin.Context) {
|
|||
}
|
||||
func (h *Media) SaveChunk(c *gin.Context) {
|
||||
claims := servermiddleware.Claims(c)
|
||||
uploadID, _ := strconv.ParseUint(c.PostForm("uploadId"), 10, 64)
|
||||
index, _ := strconv.Atoi(c.PostForm("chunkIndex"))
|
||||
header, err := c.FormFile("chunk")
|
||||
if claims == nil || err != nil {
|
||||
if claims == nil {
|
||||
httpx.Fail(c, "接收分片失败")
|
||||
return
|
||||
}
|
||||
limit := h.limitMultipartBody(c)
|
||||
header, err := c.FormFile("chunk")
|
||||
if err != nil {
|
||||
if rejectMediaTooLarge(c, err, "分片超过大小上限") {
|
||||
return
|
||||
}
|
||||
httpx.Fail(c, "接收分片失败")
|
||||
return
|
||||
}
|
||||
if header.Size > limit {
|
||||
c.AbortWithStatusJSON(http.StatusRequestEntityTooLarge, httpx.Response{Code: httpx.CodeError, Data: gin.H{}, Msg: "分片超过大小上限"})
|
||||
return
|
||||
}
|
||||
uploadID, parseErr := strconv.ParseUint(c.PostForm("uploadId"), 10, 64)
|
||||
if parseErr != nil || uploadID == 0 {
|
||||
httpx.Fail(c, "上传会话 ID 非法")
|
||||
return
|
||||
}
|
||||
index, parseErr := strconv.Atoi(c.PostForm("chunkIndex"))
|
||||
if parseErr != nil || index < 0 {
|
||||
httpx.Fail(c, "分片序号非法")
|
||||
return
|
||||
}
|
||||
opened, err := header.Open()
|
||||
if err != nil {
|
||||
httpx.Fail(c, "分片读取失败")
|
||||
|
|
|
|||
|
|
@ -34,12 +34,20 @@ func (h *Public) captchaConfig() (int, int, int) {
|
|||
|
||||
func (h *Public) Captcha(c *gin.Context) {
|
||||
keyLong, width, height := h.captchaConfig()
|
||||
security, _ := h.settings.CurrentSecurity(c.Request.Context())
|
||||
security, err := h.settings.CurrentSecurity(c.Request.Context())
|
||||
if err != nil || security == nil {
|
||||
httpx.Fail(c, "安全服务暂不可用")
|
||||
return
|
||||
}
|
||||
openCaptcha := true
|
||||
if security != nil {
|
||||
keyLong, width, height = security.KeyLong, security.ImgWidth, security.ImgHeight
|
||||
ttl := time.Duration(security.CaptchaTimeout) * time.Second
|
||||
failures, _ := h.settings.EnsureLoginIPCounter(c.Request.Context(), c.ClientIP(), ttl)
|
||||
failures, counterErr := h.settings.EnsureLoginIPCounter(c.Request.Context(), c.ClientIP(), ttl)
|
||||
if counterErr != nil {
|
||||
httpx.Fail(c, "安全服务暂不可用")
|
||||
return
|
||||
}
|
||||
openCaptcha = security.CaptchaOpen == 0 || failures > security.CaptchaOpen
|
||||
}
|
||||
driver := base64Captcha.NewDriverDigit(height, width, keyLong, 0.7, 80)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"kra/internal/biz"
|
||||
"kra/internal/conf"
|
||||
"kra/internal/server/httpx"
|
||||
|
|
@ -9,7 +11,12 @@ import (
|
|||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func AccessControl(runtime *conf.Runtime, access *service.AccessControlService) gin.HandlerFunc {
|
||||
type accessController interface {
|
||||
Authorize(context.Context, uint, string, string) (bool, error)
|
||||
ContextWithDataScope(context.Context, uint, uint) (context.Context, error)
|
||||
}
|
||||
|
||||
func AccessControl(runtime *conf.Runtime, access accessController) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
claims := Claims(c)
|
||||
if claims == nil {
|
||||
|
|
@ -29,7 +36,9 @@ func AccessControl(runtime *conf.Runtime, access *service.AccessControlService)
|
|||
}
|
||||
requestContext, err := access.ContextWithDataScope(c.Request.Context(), claims.AuthorityID, claims.ID)
|
||||
if err != nil {
|
||||
requestContext = c.Request.Context()
|
||||
httpx.Write(c, httpx.CodeError, gin.H{}, "数据权限解析失败")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if scope, ok := biz.DataScopeFromContext(requestContext); ok {
|
||||
requestID, _ := c.Get("request_id")
|
||||
|
|
|
|||
|
|
@ -2,17 +2,22 @@ package middleware
|
|||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"kra/internal/biz"
|
||||
"kra/internal/conf"
|
||||
"kra/internal/server/httpx"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const defaultRequestBodyLimit int64 = 8 << 20
|
||||
|
||||
// AccessLog is the single global request/response capture point, matching
|
||||
// the reference middleware ordering and making every HTTP request observable.
|
||||
func AccessLog(runtime *conf.Runtime, logger *slog.Logger, version string) gin.HandlerFunc {
|
||||
|
|
@ -20,20 +25,45 @@ func AccessLog(runtime *conf.Runtime, logger *slog.Logger, version string) gin.H
|
|||
started := time.Now()
|
||||
var requestBody []byte
|
||||
multipart := strings.Contains(c.GetHeader("Content-Type"), "multipart/form-data")
|
||||
mediaUpload := multipart && isMediaUploadRoute(c.FullPath())
|
||||
config := runtime.Admin()
|
||||
bodyLimit := defaultRequestBodyLimit
|
||||
if mediaUpload {
|
||||
bodyLimit = biz.DefaultMaxMediaFileSize + (1 << 20)
|
||||
if config != nil && config.Media != nil && config.Media.MaxFileSize > 0 {
|
||||
bodyLimit = config.Media.MaxFileSize + (1 << 20)
|
||||
}
|
||||
}
|
||||
bytesIn := c.Request.ContentLength
|
||||
if c.Request.Body != nil && !multipart {
|
||||
requestBody, _ = io.ReadAll(c.Request.Body)
|
||||
c.Request.Body = io.NopCloser(bytes.NewReader(requestBody))
|
||||
bytesIn = int64(len(requestBody))
|
||||
}
|
||||
if bytesIn < 0 {
|
||||
bytesIn = 0
|
||||
}
|
||||
maxBytes := 1 << 20
|
||||
writer := &captureWriter{ResponseWriter: c.Writer, maxBytes: maxBytes}
|
||||
c.Writer = writer
|
||||
c.Header("X-Kra-Version", version)
|
||||
config := runtime.Admin()
|
||||
requestReadFailed := c.Request.ContentLength > bodyLimit
|
||||
if requestReadFailed {
|
||||
c.AbortWithStatusJSON(http.StatusRequestEntityTooLarge, gin.H{"code": httpx.CodeError, "msg": "请求体超过大小上限"})
|
||||
} else if c.Request.Body != nil && !mediaUpload {
|
||||
limited := http.MaxBytesReader(c.Writer, c.Request.Body, bodyLimit)
|
||||
var err error
|
||||
requestBody, err = io.ReadAll(limited)
|
||||
if err != nil {
|
||||
requestReadFailed = true
|
||||
var tooLarge *http.MaxBytesError
|
||||
if errors.As(err, &tooLarge) {
|
||||
c.AbortWithStatusJSON(http.StatusRequestEntityTooLarge, gin.H{"code": httpx.CodeError, "msg": "请求体超过大小上限"})
|
||||
} else {
|
||||
c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"code": httpx.CodeError, "msg": "请求体读取失败"})
|
||||
}
|
||||
} else {
|
||||
c.Request.Body = io.NopCloser(bytes.NewReader(requestBody))
|
||||
bytesIn = int64(len(requestBody))
|
||||
}
|
||||
} else if c.Request.Body != nil {
|
||||
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, bodyLimit)
|
||||
}
|
||||
if bytesIn < 0 {
|
||||
bytesIn = 0
|
||||
}
|
||||
logLimit := 1024
|
||||
if config != nil && config.Zap != nil && config.Zap.AccessLogMaxBytes > 0 {
|
||||
logLimit = int(config.Zap.AccessLogMaxBytes)
|
||||
|
|
@ -46,7 +76,9 @@ func AccessLog(runtime *conf.Runtime, logger *slog.Logger, version string) gin.H
|
|||
}
|
||||
c.Set(ctxReqBodyKey, requestText)
|
||||
c.Set(ctxRespBufferKey, &writer.body)
|
||||
c.Next()
|
||||
if !requestReadFailed {
|
||||
c.Next()
|
||||
}
|
||||
if logger == nil {
|
||||
return
|
||||
}
|
||||
|
|
@ -86,6 +118,11 @@ func AccessLog(runtime *conf.Runtime, logger *slog.Logger, version string) gin.H
|
|||
}
|
||||
}
|
||||
|
||||
func isMediaUploadRoute(route string) bool {
|
||||
return strings.HasSuffix(route, "/fileUploadAndDownload/upload") ||
|
||||
strings.HasSuffix(route, "/mediaUpload/chunk")
|
||||
}
|
||||
|
||||
func redactHeaders(headers map[string][]string) map[string]string {
|
||||
out := make(map[string]string, len(headers))
|
||||
for key, values := range headers {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,65 @@
|
|||
package middleware
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"kra/internal/conf"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func TestAccessLogRejectsOversizedRequestBody(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
engine := gin.New()
|
||||
called := false
|
||||
engine.Use(AccessLog(conf.NewRuntime(nil, &conf.AdminBackend{}), nil, "test"))
|
||||
engine.POST("/payload", func(c *gin.Context) { called = true })
|
||||
request := httptest.NewRequest(http.MethodPost, "/payload", strings.NewReader(strings.Repeat("a", int(defaultRequestBodyLimit+1))))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
response := httptest.NewRecorder()
|
||||
|
||||
engine.ServeHTTP(response, request)
|
||||
if response.Code != http.StatusRequestEntityTooLarge || called {
|
||||
t.Fatalf("status=%d called=%v body=%s", response.Code, called, response.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAccessLogRejectsOversizedMultipartOnOrdinaryRoute(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
engine := gin.New()
|
||||
called := false
|
||||
engine.Use(AccessLog(conf.NewRuntime(nil, &conf.AdminBackend{}), nil, "test"))
|
||||
engine.POST("/login", func(c *gin.Context) { called = true })
|
||||
request := httptest.NewRequest(http.MethodPost, "/login", strings.NewReader(strings.Repeat("a", int(defaultRequestBodyLimit+1))))
|
||||
request.ContentLength = -1
|
||||
request.Header.Set("Content-Type", "multipart/form-data; boundary=test")
|
||||
response := httptest.NewRecorder()
|
||||
|
||||
engine.ServeHTTP(response, request)
|
||||
if response.Code != http.StatusRequestEntityTooLarge || called {
|
||||
t.Fatalf("status=%d called=%v body=%s", response.Code, called, response.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAccessLogAllowsMediaLimitOnlyOnUploadRoute(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
engine := gin.New()
|
||||
called := false
|
||||
engine.Use(AccessLog(conf.NewRuntime(nil, &conf.AdminBackend{}), nil, "test"))
|
||||
engine.POST("/api/fileUploadAndDownload/upload", func(c *gin.Context) {
|
||||
called = true
|
||||
_, _ = io.Copy(io.Discard, c.Request.Body)
|
||||
})
|
||||
request := httptest.NewRequest(http.MethodPost, "/api/fileUploadAndDownload/upload", strings.NewReader(strings.Repeat("a", int(defaultRequestBodyLimit+1))))
|
||||
request.Header.Set("Content-Type", "multipart/form-data; boundary=test")
|
||||
response := httptest.NewRecorder()
|
||||
|
||||
engine.ServeHTTP(response, request)
|
||||
if response.Code != http.StatusOK || !called {
|
||||
t.Fatalf("status=%d called=%v body=%s", response.Code, called, response.Body.String())
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,51 @@
|
|||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"kra/internal/biz"
|
||||
"kra/internal/conf"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type accessControllerStub struct {
|
||||
scopeErr error
|
||||
}
|
||||
|
||||
func (*accessControllerStub) Authorize(context.Context, uint, string, string) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (s *accessControllerStub) ContextWithDataScope(ctx context.Context, authorityID, userID uint) (context.Context, error) {
|
||||
if s.scopeErr != nil {
|
||||
return ctx, s.scopeErr
|
||||
}
|
||||
return biz.NewDataScopeContext(ctx, biz.DataScope{UserID: userID, AuthorityID: authorityID, Scope: 1, All: true}), nil
|
||||
}
|
||||
|
||||
func TestAccessControlFailsClosedWhenDataScopeResolutionFails(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
engine := gin.New()
|
||||
engine.Use(func(c *gin.Context) {
|
||||
c.Set(claimsKey, &biz.AuthClaims{ID: 7, AuthorityID: 888})
|
||||
c.Next()
|
||||
})
|
||||
engine.Use(AccessControl(conf.NewRuntime(nil, &conf.AdminBackend{}), &accessControllerStub{scopeErr: errors.New("database unavailable")}))
|
||||
called := false
|
||||
engine.GET("/protected", func(c *gin.Context) { called = true })
|
||||
|
||||
response := httptest.NewRecorder()
|
||||
engine.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/protected", nil))
|
||||
if called {
|
||||
t.Fatal("protected handler ran after data-scope resolution failed")
|
||||
}
|
||||
if response.Code != http.StatusOK || !strings.Contains(response.Body.String(), "数据权限解析失败") {
|
||||
t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
|
||||
}
|
||||
}
|
||||
|
|
@ -3,6 +3,7 @@ package middleware
|
|||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
|
|
@ -41,7 +42,18 @@ func OperationAudit(runtime *conf.Runtime, service *service.AuditRecorder) gin.H
|
|||
if strings.Contains(c.GetHeader("Content-Type"), "multipart/form-data") {
|
||||
requestBody = []byte("[文件]")
|
||||
} else {
|
||||
requestBody, _ = io.ReadAll(c.Request.Body)
|
||||
limited := http.MaxBytesReader(c.Writer, c.Request.Body, defaultRequestBodyLimit)
|
||||
var err error
|
||||
requestBody, err = io.ReadAll(limited)
|
||||
if err != nil {
|
||||
var tooLarge *http.MaxBytesError
|
||||
if errors.As(err, &tooLarge) {
|
||||
c.AbortWithStatus(http.StatusRequestEntityTooLarge)
|
||||
} else {
|
||||
c.AbortWithStatus(http.StatusBadRequest)
|
||||
}
|
||||
return
|
||||
}
|
||||
c.Request.Body = io.NopCloser(bytes.NewReader(requestBody))
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
|
|
@ -18,7 +19,11 @@ func SecurityRateLimit(settings *service.SecurityService) gin.HandlerFunc {
|
|||
return
|
||||
}
|
||||
config, err := settings.CurrentSecurity(c.Request.Context())
|
||||
if err != nil || config == nil || !config.LimitEnable {
|
||||
if err != nil || config == nil {
|
||||
c.AbortWithStatusJSON(http.StatusServiceUnavailable, gin.H{"code": httpx.CodeError, "msg": "安全服务暂不可用"})
|
||||
return
|
||||
}
|
||||
if !config.LimitEnable {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
|
@ -28,7 +33,11 @@ func SecurityRateLimit(settings *service.SecurityService) gin.HandlerFunc {
|
|||
}
|
||||
key := "KRA_SecLimit" + c.ClientIP() + c.FullPath()
|
||||
count, cacheErr := settings.IncrementRateLimit(c.Request.Context(), key, time.Duration(window)*time.Second)
|
||||
if cacheErr == nil && int(count) > config.LimitCount {
|
||||
if cacheErr != nil {
|
||||
c.AbortWithStatusJSON(http.StatusServiceUnavailable, gin.H{"code": httpx.CodeError, "msg": "安全服务暂不可用"})
|
||||
return
|
||||
}
|
||||
if int(count) > config.LimitCount {
|
||||
c.JSON(200, gin.H{"code": httpx.CodeError, "msg": "请求太过频繁,请稍后再试"})
|
||||
c.Abort()
|
||||
return
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ package middleware
|
|||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
|
|
@ -26,12 +27,14 @@ func (rateLimitSecurityRepo) BackfillPasswordUpdatedAt(context.Context, time.Tim
|
|||
return nil
|
||||
}
|
||||
|
||||
type rateLimitCache struct{}
|
||||
type rateLimitCache struct{ err error }
|
||||
|
||||
func (rateLimitCache) Get(context.Context, string) (string, bool, error) { return "", false, nil }
|
||||
func (rateLimitCache) Set(context.Context, string, string, time.Duration) error { return nil }
|
||||
func (rateLimitCache) Delete(context.Context, string) error { return nil }
|
||||
func (rateLimitCache) Increment(context.Context, string, time.Duration) (int64, error) { return 2, nil }
|
||||
func (rateLimitCache) Get(context.Context, string) (string, bool, error) { return "", false, nil }
|
||||
func (rateLimitCache) Set(context.Context, string, string, time.Duration) error { return nil }
|
||||
func (rateLimitCache) Delete(context.Context, string) error { return nil }
|
||||
func (c rateLimitCache) Increment(context.Context, string, time.Duration) (int64, error) {
|
||||
return 2, c.err
|
||||
}
|
||||
|
||||
func TestSecurityRateLimitMatchesResponseContract(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
|
@ -50,3 +53,18 @@ func TestSecurityRateLimitMatchesResponseContract(t *testing.T) {
|
|||
t.Fatalf("unexpected rate-limit response: %s", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityRateLimitFailsClosedWhenCacheIsUnavailable(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
settings := service.NewSecurityService(biz.NewSecurityUsecase(rateLimitSecurityRepo{}, rateLimitCache{err: errors.New("cache unavailable")}, nil, nil))
|
||||
engine := gin.New()
|
||||
called := false
|
||||
engine.Use(SecurityRateLimit(settings))
|
||||
engine.POST("/base/login", func(c *gin.Context) { called = true })
|
||||
|
||||
response := httptest.NewRecorder()
|
||||
engine.ServeHTTP(response, httptest.NewRequest(http.MethodPost, "/base/login", nil))
|
||||
if response.Code != http.StatusServiceUnavailable || called {
|
||||
t.Fatalf("status=%d called=%v body=%s", response.Code, called, response.Body.String())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ func RegisterAnnouncement(private, public *gin.RouterGroup, h *handler.Announcem
|
|||
privateInfo.PUT("/updateInfo", h.Update)
|
||||
privateInfo.GET("/findInfo", h.Find)
|
||||
privateInfo.GET("/getInfoList", h.List)
|
||||
privateInfo.GET("/getInfoDataSource", h.DataSource)
|
||||
publicInfo := public.Group("/info")
|
||||
publicInfo.GET("/getInfoDataSource", h.DataSource)
|
||||
publicInfo.GET("/getInfoPublic", h.Public)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -65,7 +65,7 @@ func (s *APIService) Groups(ctx context.Context) ([]string, map[string]string, e
|
|||
return nil, nil, err
|
||||
}
|
||||
seen := make(map[string]bool)
|
||||
var groups []string
|
||||
groups := make([]string, 0)
|
||||
groupAPIMap := make(map[string]string)
|
||||
for _, item := range items {
|
||||
if !seen[item.APIGroup] {
|
||||
|
|
|
|||
|
|
@ -103,12 +103,12 @@ func truncateTaskText(value string) string {
|
|||
return value[:limit] + "...(截断)"
|
||||
}
|
||||
|
||||
func (e *TaskExecutor) runMethod(task *biz.TimedTask) error {
|
||||
func (e *TaskExecutor) runMethod(ctx context.Context, task *biz.TimedTask) error {
|
||||
method, ok := biz.TaskMethodByName(task.MethodName)
|
||||
if !ok {
|
||||
return fmt.Errorf("方法 %s 未注册(需在 internal/worker/task_registry.go 经 biz.RegisterTaskMethod 注册)", task.MethodName)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
|
||||
runCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
|
||||
defer cancel()
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
|
|
@ -117,7 +117,7 @@ func (e *TaskExecutor) runMethod(task *biz.TimedTask) error {
|
|||
done <- fmt.Errorf("panic: %v", recovered)
|
||||
}
|
||||
}()
|
||||
done <- method(ctx, json.RawMessage(task.Params))
|
||||
done <- method(runCtx, json.RawMessage(task.Params))
|
||||
}()
|
||||
select {
|
||||
case err := <-done:
|
||||
|
|
@ -125,12 +125,18 @@ func (e *TaskExecutor) runMethod(task *biz.TimedTask) error {
|
|||
return errTaskTimeout
|
||||
}
|
||||
return err
|
||||
case <-ctx.Done():
|
||||
return errTaskTimeout
|
||||
case <-runCtx.Done():
|
||||
if errors.Is(runCtx.Err(), context.DeadlineExceeded) {
|
||||
return errTaskTimeout
|
||||
}
|
||||
return runCtx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (e *TaskExecutor) Run(ctx context.Context, task *biz.TimedTask, trigger string) (log *biz.TimedTaskLog) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
started := time.Now()
|
||||
log = &biz.TimedTaskLog{TaskID: task.ID, TaskName: task.Name, TriggerType: trigger, StartedAt: started, Status: "success"}
|
||||
defer func() {
|
||||
|
|
@ -148,9 +154,9 @@ func (e *TaskExecutor) Run(ctx context.Context, task *biz.TimedTask, trigger str
|
|||
var err error
|
||||
switch task.ExecutorType {
|
||||
case biz.TaskExecutorMethod:
|
||||
err = e.runMethod(task)
|
||||
err = e.runMethod(ctx, task)
|
||||
case biz.TaskExecutorHTTP:
|
||||
runCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
runCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||
log.Output, err = e.runHTTP(runCtx, task)
|
||||
cancel()
|
||||
default:
|
||||
|
|
|
|||
|
|
@ -1,8 +1,15 @@
|
|||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"kra/internal/biz"
|
||||
)
|
||||
|
||||
func TestPrivateIP(t *testing.T) {
|
||||
|
|
@ -15,3 +22,76 @@ func TestPrivateIP(t *testing.T) {
|
|||
t.Fatal("public IP was classified as private")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunMethodReportsParentDeadlineAsTimeout(t *testing.T) {
|
||||
const methodName = "worker-test-parent-deadline"
|
||||
biz.RegisterTaskMethod(methodName, "test", func(ctx context.Context, _ json.RawMessage) error {
|
||||
<-ctx.Done()
|
||||
return ctx.Err()
|
||||
})
|
||||
executor := &TaskExecutor{}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
if err := executor.runMethod(ctx, &biz.TimedTask{MethodName: methodName}); !errors.Is(err, errTaskTimeout) {
|
||||
t.Fatalf("runMethod() error = %v, want errTaskTimeout", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunMethodConvertsPanicToError(t *testing.T) {
|
||||
const methodName = "worker-test-panic"
|
||||
biz.RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error {
|
||||
panic("boom")
|
||||
})
|
||||
executor := &TaskExecutor{}
|
||||
err := executor.runMethod(context.Background(), &biz.TimedTask{MethodName: methodName})
|
||||
if err == nil || !strings.Contains(err.Error(), "panic: boom") {
|
||||
t.Fatalf("runMethod() error = %v, want recovered panic", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunMethodUsesParentContext(t *testing.T) {
|
||||
const methodName = "worker-test-parent-context"
|
||||
started := make(chan struct{})
|
||||
biz.RegisterTaskMethod(methodName, "test", func(ctx context.Context, _ json.RawMessage) error {
|
||||
close(started)
|
||||
<-ctx.Done()
|
||||
return ctx.Err()
|
||||
})
|
||||
executor := &TaskExecutor{}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
done <- executor.runMethod(ctx, &biz.TimedTask{MethodName: methodName})
|
||||
}()
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("registered method did not start")
|
||||
}
|
||||
cancel()
|
||||
select {
|
||||
case err := <-done:
|
||||
if !errors.Is(err, context.Canceled) {
|
||||
t.Fatalf("runMethod() error = %v, want context.Canceled", err)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("runMethod() did not stop after parent cancellation")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunMethodHonorsDeadlineWhenMethodIgnoresContext(t *testing.T) {
|
||||
const methodName = "worker-test-ignores-context"
|
||||
release := make(chan struct{})
|
||||
biz.RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error {
|
||||
<-release
|
||||
return nil
|
||||
})
|
||||
executor := &TaskExecutor{}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
if err := executor.runMethod(ctx, &biz.TimedTask{MethodName: methodName}); !errors.Is(err, errTaskTimeout) {
|
||||
close(release)
|
||||
t.Fatalf("runMethod() error = %v, want errTaskTimeout", err)
|
||||
}
|
||||
close(release)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,274 @@
|
|||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"kra/internal/biz"
|
||||
)
|
||||
|
||||
type workerTaskRepo struct {
|
||||
mu sync.Mutex
|
||||
items []*biz.TimedTask
|
||||
listErr error
|
||||
logs chan *biz.TimedTaskLog
|
||||
}
|
||||
|
||||
func (r *workerTaskRepo) CreateTask(context.Context, *biz.TimedTask) error { return nil }
|
||||
func (r *workerTaskRepo) UpdateTask(context.Context, *biz.TimedTask) error { return nil }
|
||||
func (r *workerTaskRepo) DeleteTask(context.Context, uint) error { return nil }
|
||||
func (r *workerTaskRepo) ToggleTask(context.Context, uint, bool) error { return nil }
|
||||
func (r *workerTaskRepo) CleanupLogs(context.Context) error { return nil }
|
||||
func (r *workerTaskRepo) TaskNameExists(context.Context, string, uint) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
func (r *workerTaskRepo) FindTask(_ context.Context, id uint) (*biz.TimedTask, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
for _, item := range r.items {
|
||||
if item.ID == id {
|
||||
return cloneTimedTask(item), nil
|
||||
}
|
||||
}
|
||||
return nil, errors.New("task not found")
|
||||
}
|
||||
func (r *workerTaskRepo) ListTasks(context.Context, int, int, *biz.TimedTask) ([]*biz.TimedTask, int64, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.listErr != nil {
|
||||
return nil, 0, r.listErr
|
||||
}
|
||||
items := make([]*biz.TimedTask, 0, len(r.items))
|
||||
for _, item := range r.items {
|
||||
items = append(items, cloneTimedTask(item))
|
||||
}
|
||||
return items, int64(len(items)), nil
|
||||
}
|
||||
func (r *workerTaskRepo) RecordTaskLog(_ context.Context, value *biz.TimedTaskLog) error {
|
||||
if r.logs != nil {
|
||||
copy := *value
|
||||
r.logs <- ©
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (r *workerTaskRepo) ListTaskLogs(context.Context, int, int, uint, string) ([]*biz.TimedTaskLog, int64, error) {
|
||||
return nil, 0, nil
|
||||
}
|
||||
|
||||
func newTestTaskScheduler(repo *workerTaskRepo) *TaskScheduler {
|
||||
tasks := biz.NewTaskUsecase(repo)
|
||||
executor := &TaskExecutor{tasks: tasks}
|
||||
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
return NewTaskScheduler(tasks, nil, executor, logger)
|
||||
}
|
||||
|
||||
func TestReloadPreservesExistingScheduleOnPreparationFailure(t *testing.T) {
|
||||
repo := &workerTaskRepo{}
|
||||
scheduler := newTestTaskScheduler(repo)
|
||||
old := &biz.TimedTask{ID: 1, Name: "old", Spec: "0 0 * * *", Enabled: true}
|
||||
if err := scheduler.Schedule(old); err != nil {
|
||||
t.Fatalf("schedule old task: %v", err)
|
||||
}
|
||||
|
||||
t.Run("query failure", func(t *testing.T) {
|
||||
repo.mu.Lock()
|
||||
repo.listErr = errors.New("database unavailable")
|
||||
repo.mu.Unlock()
|
||||
if err := scheduler.Reload(context.Background()); err == nil {
|
||||
t.Fatal("Reload() error = nil, want query error")
|
||||
}
|
||||
assertOnlyScheduledTask(t, scheduler, old.ID)
|
||||
})
|
||||
|
||||
t.Run("invalid cron", func(t *testing.T) {
|
||||
repo.mu.Lock()
|
||||
repo.listErr = nil
|
||||
repo.items = []*biz.TimedTask{{ID: 2, Name: "invalid", Spec: "bad cron", Enabled: true}}
|
||||
repo.mu.Unlock()
|
||||
if err := scheduler.Reload(context.Background()); err == nil {
|
||||
t.Fatal("Reload() error = nil, want cron error")
|
||||
}
|
||||
assertOnlyScheduledTask(t, scheduler, old.ID)
|
||||
})
|
||||
}
|
||||
|
||||
func assertOnlyScheduledTask(t *testing.T, scheduler *TaskScheduler, id uint) {
|
||||
t.Helper()
|
||||
scheduler.mu.Lock()
|
||||
defer scheduler.mu.Unlock()
|
||||
if len(scheduler.entries) != 1 {
|
||||
t.Fatalf("scheduled task count = %d, want 1", len(scheduler.entries))
|
||||
}
|
||||
if _, ok := scheduler.entries[id]; !ok {
|
||||
t.Fatalf("scheduled task %d was replaced after failed reload", id)
|
||||
}
|
||||
if got := len(scheduler.standard.Entries()) + len(scheduler.seconds.Entries()); got != 1 {
|
||||
t.Fatalf("cron entry count = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrentScheduleKeepsSingleEntry(t *testing.T) {
|
||||
scheduler := newTestTaskScheduler(&workerTaskRepo{})
|
||||
const workers = 32
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < workers; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if err := scheduler.Schedule(&biz.TimedTask{ID: 7, Name: "same", Spec: "0 0 * * *", Enabled: true}); err != nil {
|
||||
t.Errorf("Schedule() error = %v", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
assertOnlyScheduledTask(t, scheduler, 7)
|
||||
}
|
||||
|
||||
func TestManualTriggerSkipsOverlappingExecution(t *testing.T) {
|
||||
const methodName = "worker-test-no-overlap"
|
||||
started := make(chan struct{}, 2)
|
||||
release := make(chan struct{})
|
||||
var calls atomic.Int32
|
||||
biz.RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error {
|
||||
calls.Add(1)
|
||||
started <- struct{}{}
|
||||
<-release
|
||||
return nil
|
||||
})
|
||||
|
||||
task := &biz.TimedTask{ID: 9, Name: "single", ExecutorType: biz.TaskExecutorMethod, MethodName: methodName}
|
||||
repo := &workerTaskRepo{items: []*biz.TimedTask{task}, logs: make(chan *biz.TimedTaskLog, 2)}
|
||||
scheduler := newTestTaskScheduler(repo)
|
||||
if err := scheduler.TriggerID(context.Background(), task.ID); err != nil {
|
||||
t.Fatalf("first manual trigger failed: %v", err)
|
||||
}
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("first task execution did not start")
|
||||
}
|
||||
|
||||
if err := scheduler.TriggerID(context.Background(), task.ID); err == nil {
|
||||
t.Fatal("overlapping manual trigger returned success")
|
||||
}
|
||||
if got := calls.Load(); got != 1 {
|
||||
t.Fatalf("execution count while first run is active = %d, want 1", got)
|
||||
}
|
||||
close(release)
|
||||
select {
|
||||
case <-repo.logs:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("first task execution did not finish")
|
||||
}
|
||||
waitForSchedulerIdle(t, scheduler)
|
||||
|
||||
if err := scheduler.TriggerID(context.Background(), task.ID); err != nil {
|
||||
t.Fatalf("manual trigger failed after the previous run finished: %v", err)
|
||||
}
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("task could not run again after the first execution finished")
|
||||
}
|
||||
select {
|
||||
case <-repo.logs:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("second task execution did not finish")
|
||||
}
|
||||
waitForSchedulerIdle(t, scheduler)
|
||||
if got := calls.Load(); got != 2 {
|
||||
t.Fatalf("total execution count = %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAutomaticAndManualTriggersShareOverlapGate(t *testing.T) {
|
||||
const methodName = "worker-test-auto-manual-overlap"
|
||||
started := make(chan struct{}, 1)
|
||||
release := make(chan struct{})
|
||||
var calls atomic.Int32
|
||||
biz.RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error {
|
||||
calls.Add(1)
|
||||
started <- struct{}{}
|
||||
<-release
|
||||
return nil
|
||||
})
|
||||
|
||||
task := &biz.TimedTask{ID: 10, Name: "shared-gate", ExecutorType: biz.TaskExecutorMethod, MethodName: methodName}
|
||||
repo := &workerTaskRepo{logs: make(chan *biz.TimedTaskLog, 1)}
|
||||
scheduler := newTestTaskScheduler(repo)
|
||||
if !scheduler.dispatch(task, "auto", true) {
|
||||
t.Fatal("automatic dispatch was rejected while idle")
|
||||
}
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("automatic task execution did not start")
|
||||
}
|
||||
if scheduler.Trigger(task) {
|
||||
t.Fatal("manual trigger bypassed an active automatic execution")
|
||||
}
|
||||
if got := calls.Load(); got != 1 {
|
||||
t.Fatalf("execution count while automatic run is active = %d, want 1", got)
|
||||
}
|
||||
close(release)
|
||||
select {
|
||||
case <-repo.logs:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("automatic task execution did not finish")
|
||||
}
|
||||
waitForSchedulerIdle(t, scheduler)
|
||||
}
|
||||
|
||||
func waitForSchedulerIdle(t *testing.T, scheduler *TaskScheduler) {
|
||||
t.Helper()
|
||||
scheduler.runMu.Lock()
|
||||
idle := scheduler.idle
|
||||
scheduler.runMu.Unlock()
|
||||
select {
|
||||
case <-idle:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("scheduler did not become idle")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopCancelsAndWaitsForActiveRuns(t *testing.T) {
|
||||
scheduler := newTestTaskScheduler(&workerTaskRepo{})
|
||||
runCtx, cancel := context.WithCancel(context.Background())
|
||||
scheduler.ctxMu.Lock()
|
||||
scheduler.runContext = runCtx
|
||||
scheduler.cancel = cancel
|
||||
scheduler.ctxMu.Unlock()
|
||||
|
||||
taskCtx, ok := scheduler.beginRun(12)
|
||||
if !ok {
|
||||
t.Fatal("beginRun() rejected an idle scheduler")
|
||||
}
|
||||
finished := make(chan struct{})
|
||||
go func() {
|
||||
<-taskCtx.Done()
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
close(finished)
|
||||
scheduler.finishRun(12)
|
||||
}()
|
||||
|
||||
stopCtx, stopCancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer stopCancel()
|
||||
if err := scheduler.Stop(stopCtx); err != nil {
|
||||
t.Fatalf("Stop() error = %v", err)
|
||||
}
|
||||
select {
|
||||
case <-finished:
|
||||
default:
|
||||
t.Fatal("Stop() returned before the active run finished")
|
||||
}
|
||||
if _, ok := scheduler.beginRun(13); ok {
|
||||
t.Fatal("scheduler accepted a new run after Stop()")
|
||||
}
|
||||
}
|
||||
|
|
@ -3,6 +3,7 @@ package worker
|
|||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"sync"
|
||||
|
|
@ -24,6 +25,10 @@ type TaskScheduler struct {
|
|||
ctxMu sync.RWMutex
|
||||
runContext context.Context
|
||||
cancel context.CancelFunc
|
||||
runMu sync.Mutex
|
||||
running map[uint]struct{}
|
||||
idle chan struct{}
|
||||
stopping bool
|
||||
subMu sync.RWMutex
|
||||
subscribers map[uint]map[chan []byte]struct{}
|
||||
}
|
||||
|
|
@ -34,7 +39,9 @@ type scheduledEntry struct {
|
|||
}
|
||||
|
||||
func NewTaskScheduler(tasks *biz.TaskUsecase, authorities *biz.AuthorityUsecase, executor *TaskExecutor, logger *slog.Logger) *TaskScheduler {
|
||||
return &TaskScheduler{tasks: tasks, authorities: authorities, executor: executor, logger: logger.With("mod", "timedTask"), standard: cron.New(), seconds: cron.New(cron.WithSeconds()), entries: map[uint]scheduledEntry{}, subscribers: map[uint]map[chan []byte]struct{}{}}
|
||||
idle := make(chan struct{})
|
||||
close(idle)
|
||||
return &TaskScheduler{tasks: tasks, authorities: authorities, executor: executor, logger: logger.With("mod", "timedTask"), standard: cron.New(), seconds: cron.New(cron.WithSeconds()), entries: map[uint]scheduledEntry{}, running: map[uint]struct{}{}, idle: idle, subscribers: map[uint]map[chan []byte]struct{}{}}
|
||||
}
|
||||
|
||||
func NewTaskRuntime(scheduler *TaskScheduler) biz.TaskRuntime { return scheduler }
|
||||
|
|
@ -44,21 +51,12 @@ func (s *TaskScheduler) Start(ctx context.Context) error {
|
|||
s.ctxMu.Lock()
|
||||
s.runContext, s.cancel = runContext, cancel
|
||||
s.ctxMu.Unlock()
|
||||
s.runMu.Lock()
|
||||
s.stopping = false
|
||||
s.runMu.Unlock()
|
||||
s.standard.Start()
|
||||
s.seconds.Start()
|
||||
items, _, err := s.tasks.ListTasks(ctx, 0, 0, nil)
|
||||
if err == nil {
|
||||
for _, task := range items {
|
||||
if task.Enabled {
|
||||
if scheduleErr := s.Schedule(task); scheduleErr != nil {
|
||||
s.logger.ErrorContext(ctx, "restore timed task failed", "id", task.ID, "error", scheduleErr)
|
||||
}
|
||||
}
|
||||
}
|
||||
// Report the number of rows read from the database, not only
|
||||
// the enabled rows that were scheduled successfully.
|
||||
s.logger.InfoContext(ctx, "定时任务加载完成", "task_count", len(items))
|
||||
} else {
|
||||
if err := s.Reload(ctx); err != nil {
|
||||
s.logger.WarnContext(ctx, "timed task table is not ready", "error", err)
|
||||
}
|
||||
<-ctx.Done()
|
||||
|
|
@ -67,28 +65,38 @@ func (s *TaskScheduler) Start(ctx context.Context) error {
|
|||
|
||||
func (s *TaskScheduler) Stop(ctx context.Context) error {
|
||||
s.closeSubscribers()
|
||||
s.runMu.Lock()
|
||||
s.stopping = true
|
||||
idle := s.idle
|
||||
if idle == nil {
|
||||
idle = make(chan struct{})
|
||||
close(idle)
|
||||
}
|
||||
s.runMu.Unlock()
|
||||
s.ctxMu.Lock()
|
||||
if s.cancel != nil {
|
||||
s.cancel()
|
||||
}
|
||||
cancel := s.cancel
|
||||
s.ctxMu.Unlock()
|
||||
if cancel != nil {
|
||||
cancel()
|
||||
}
|
||||
standardDone, secondsDone := s.standard.Stop().Done(), s.seconds.Stop().Done()
|
||||
select {
|
||||
case <-standardDone:
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
select {
|
||||
case <-secondsDone:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
for _, done := range []<-chan struct{}{standardDone, secondsDone, idle} {
|
||||
select {
|
||||
case <-done:
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) Remove(id uint) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.removeLocked(id)
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) removeLocked(id uint) {
|
||||
if old, ok := s.entries[id]; ok {
|
||||
if old.seconds {
|
||||
s.seconds.Remove(old.entry)
|
||||
|
|
@ -101,26 +109,33 @@ func (s *TaskScheduler) Remove(id uint) {
|
|||
|
||||
func (s *TaskScheduler) Reload(ctx context.Context) error {
|
||||
s.mu.Lock()
|
||||
for id, old := range s.entries {
|
||||
if old.seconds {
|
||||
s.seconds.Remove(old.entry)
|
||||
} else {
|
||||
s.standard.Remove(old.entry)
|
||||
}
|
||||
delete(s.entries, id)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
items, _, err := s.tasks.ListTasks(ctx, 0, 0, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
type preparedTask struct {
|
||||
task *biz.TimedTask
|
||||
schedule cron.Schedule
|
||||
}
|
||||
prepared := make([]preparedTask, 0, len(items))
|
||||
for _, task := range items {
|
||||
if task.Enabled {
|
||||
if err = s.Schedule(task); err != nil {
|
||||
return err
|
||||
schedule, parseErr := parseTaskSchedule(task)
|
||||
if parseErr != nil {
|
||||
return fmt.Errorf("定时任务 %d 的 cron 表达式非法: %w", task.ID, parseErr)
|
||||
}
|
||||
prepared = append(prepared, preparedTask{task: cloneTimedTask(task), schedule: schedule})
|
||||
}
|
||||
}
|
||||
|
||||
for id := range s.entries {
|
||||
s.removeLocked(id)
|
||||
}
|
||||
for _, item := range prepared {
|
||||
s.scheduleLocked(item.task, item.schedule)
|
||||
}
|
||||
s.logger.InfoContext(ctx, "定时任务重载完成", "task_count", len(items))
|
||||
return nil
|
||||
}
|
||||
|
|
@ -134,8 +149,58 @@ func (s *TaskScheduler) executionContext() context.Context {
|
|||
return context.Background()
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) run(task *biz.TimedTask, trigger string) {
|
||||
log := s.executor.Run(s.executionContext(), task, trigger)
|
||||
func (s *TaskScheduler) beginRun(id uint) (context.Context, bool) {
|
||||
ctx := s.executionContext()
|
||||
s.runMu.Lock()
|
||||
defer s.runMu.Unlock()
|
||||
if s.stopping || ctx.Err() != nil {
|
||||
return nil, false
|
||||
}
|
||||
if s.running == nil {
|
||||
s.running = map[uint]struct{}{}
|
||||
}
|
||||
if _, exists := s.running[id]; exists {
|
||||
return nil, false
|
||||
}
|
||||
if len(s.running) == 0 {
|
||||
s.idle = make(chan struct{})
|
||||
}
|
||||
s.running[id] = struct{}{}
|
||||
return ctx, true
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) finishRun(id uint) {
|
||||
s.runMu.Lock()
|
||||
defer s.runMu.Unlock()
|
||||
if _, exists := s.running[id]; !exists {
|
||||
return
|
||||
}
|
||||
delete(s.running, id)
|
||||
if len(s.running) == 0 && s.idle != nil {
|
||||
close(s.idle)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) dispatch(task *biz.TimedTask, trigger string, async bool) bool {
|
||||
ctx, ok := s.beginRun(task.ID)
|
||||
if !ok {
|
||||
s.logger.Warn("timed task skipped because it is already running or the scheduler is stopping", "task_id", task.ID, "task_name", task.Name, "trigger_type", trigger)
|
||||
return false
|
||||
}
|
||||
run := func() {
|
||||
defer s.finishRun(task.ID)
|
||||
s.run(ctx, task, trigger)
|
||||
}
|
||||
if async {
|
||||
go run()
|
||||
return true
|
||||
}
|
||||
run()
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) run(ctx context.Context, task *biz.TimedTask, trigger string) {
|
||||
log := s.executor.Run(ctx, task, trigger)
|
||||
attributes := []any{"task_id", log.TaskID, "task_name", log.TaskName, "trigger_type", log.TriggerType, "status", log.Status, "duration_ms", log.DurationMS, "started_at", log.StartedAt, "finished_at", log.FinishedAt}
|
||||
if log.ErrorMsg != "" {
|
||||
attributes = append(attributes, "error", log.ErrorMsg)
|
||||
|
|
@ -146,7 +211,9 @@ func (s *TaskScheduler) run(task *biz.TimedTask, trigger string) {
|
|||
s.logger.Error("timed task finished", attributes...)
|
||||
}
|
||||
if log.Status != "success" {
|
||||
ids, err := s.authorities.AuthorityUserIDs(context.Background(), 888)
|
||||
alertCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 3*time.Second)
|
||||
defer cancel()
|
||||
ids, err := s.authorities.AuthorityUserIDs(alertCtx, 888)
|
||||
if err != nil {
|
||||
s.logger.Error("query timed task alert recipients failed", "error", err)
|
||||
return
|
||||
|
|
@ -156,36 +223,73 @@ func (s *TaskScheduler) run(task *biz.TimedTask, trigger string) {
|
|||
}
|
||||
|
||||
func (s *TaskScheduler) Schedule(task *biz.TimedTask) error {
|
||||
s.Remove(task.ID)
|
||||
if task == nil {
|
||||
return errors.New("定时任务不能为空")
|
||||
}
|
||||
var schedule cron.Schedule
|
||||
var err error
|
||||
if task.Enabled {
|
||||
schedule, err = parseTaskSchedule(task)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.removeLocked(task.ID)
|
||||
if !task.Enabled {
|
||||
return nil
|
||||
}
|
||||
copy := *task
|
||||
run := func() {
|
||||
s.run(©, "auto")
|
||||
}
|
||||
var id cron.EntryID
|
||||
var err error
|
||||
if task.WithSeconds {
|
||||
id, err = s.seconds.AddFunc(task.Spec, run)
|
||||
} else {
|
||||
id, err = s.standard.AddFunc(task.Spec, run)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.entries[task.ID] = scheduledEntry{seconds: task.WithSeconds, entry: id}
|
||||
s.mu.Unlock()
|
||||
s.scheduleLocked(cloneTimedTask(task), schedule)
|
||||
return nil
|
||||
}
|
||||
|
||||
func parseTaskSchedule(task *biz.TimedTask) (cron.Schedule, error) {
|
||||
if task.WithSeconds {
|
||||
return cron.NewParser(cron.Second | cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor).Parse(task.Spec)
|
||||
}
|
||||
return cron.ParseStandard(task.Spec)
|
||||
}
|
||||
|
||||
func cloneTimedTask(task *biz.TimedTask) *biz.TimedTask {
|
||||
copy := *task
|
||||
copy.Params = append([]byte(nil), task.Params...)
|
||||
copy.HTTPHeader = append([]byte(nil), task.HTTPHeader...)
|
||||
return ©
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) scheduleLocked(task *biz.TimedTask, schedule cron.Schedule) {
|
||||
run := func() {
|
||||
s.dispatch(task, "auto", false)
|
||||
}
|
||||
if task.WithSeconds {
|
||||
id := s.seconds.Schedule(schedule, cron.FuncJob(run))
|
||||
s.entries[task.ID] = scheduledEntry{seconds: true, entry: id}
|
||||
} else {
|
||||
id := s.standard.Schedule(schedule, cron.FuncJob(run))
|
||||
s.entries[task.ID] = scheduledEntry{entry: id}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) ScheduleID(ctx context.Context, id uint) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
task, err := s.tasks.FindTask(ctx, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.Schedule(task)
|
||||
var schedule cron.Schedule
|
||||
if task.Enabled {
|
||||
schedule, err = parseTaskSchedule(task)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
s.removeLocked(task.ID)
|
||||
if task.Enabled {
|
||||
s.scheduleLocked(cloneTimedTask(task), schedule)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) TriggerID(ctx context.Context, id uint) error {
|
||||
|
|
@ -193,7 +297,9 @@ func (s *TaskScheduler) TriggerID(ctx context.Context, id uint) error {
|
|||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.Trigger(task)
|
||||
if !s.Trigger(task) {
|
||||
return errors.New("任务正在执行或调度器正在停止")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
@ -211,9 +317,11 @@ func (s *TaskScheduler) NextRuns() map[uint]time.Time {
|
|||
return out
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) Trigger(task *biz.TimedTask) {
|
||||
copy := *task
|
||||
go s.run(©, "manual")
|
||||
func (s *TaskScheduler) Trigger(task *biz.TimedTask) bool {
|
||||
if task == nil {
|
||||
return false
|
||||
}
|
||||
return s.dispatch(cloneTimedTask(task), "manual", true)
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) Subscribe(userID uint) chan []byte {
|
||||
|
|
|
|||
|
|
@ -16,23 +16,24 @@ var (
|
|||
)
|
||||
|
||||
type Claims struct {
|
||||
UUID string
|
||||
ID uint
|
||||
Username string
|
||||
NickName string
|
||||
AuthorityID uint `json:"AuthorityId"`
|
||||
BufferTime int64
|
||||
UserType string
|
||||
MustChangePwd bool `json:"mustChangePwd"`
|
||||
UUID string
|
||||
ID uint
|
||||
Username string
|
||||
NickName string
|
||||
AuthorityID uint `json:"AuthorityId"`
|
||||
BufferTime int64
|
||||
UserType string
|
||||
MustChangePwd bool `json:"mustChangePwd"`
|
||||
PasswordVersion int64 `json:"passwordVersion,omitempty"`
|
||||
jwt.RegisteredClaims
|
||||
}
|
||||
|
||||
func Generate(secret, issuer string, expires, buffer time.Duration, userID, authorityID uint, uuid, username, nickname string, mustChange bool) (string, *Claims, error) {
|
||||
func Generate(secret, issuer string, expires, buffer time.Duration, userID, authorityID uint, uuid, username, nickname string, mustChange bool, passwordVersion int64) (string, *Claims, error) {
|
||||
if secret == "" {
|
||||
return "", nil, errors.New("empty JWT signing key")
|
||||
}
|
||||
now := time.Now()
|
||||
claims := &Claims{UUID: uuid, ID: userID, Username: username, NickName: nickname, AuthorityID: authorityID, BufferTime: int64(buffer / time.Second), UserType: "admin", MustChangePwd: mustChange, RegisteredClaims: jwt.RegisteredClaims{Audience: jwt.ClaimStrings{"KRA"}, Issuer: issuer, NotBefore: jwt.NewNumericDate(now.Add(-1000)), ExpiresAt: jwt.NewNumericDate(now.Add(expires))}}
|
||||
claims := &Claims{UUID: uuid, ID: userID, Username: username, NickName: nickname, AuthorityID: authorityID, BufferTime: int64(buffer / time.Second), UserType: "admin", MustChangePwd: mustChange, PasswordVersion: passwordVersion, RegisteredClaims: jwt.RegisteredClaims{Audience: jwt.ClaimStrings{"KRA"}, Issuer: issuer, IssuedAt: jwt.NewNumericDate(now), NotBefore: jwt.NewNumericDate(now.Add(-1000)), ExpiresAt: jwt.NewNumericDate(now.Add(expires))}}
|
||||
token, err := Sign(secret, claims)
|
||||
return token, claims, err
|
||||
}
|
||||
|
|
|
|||
|
|
@ -519,12 +519,15 @@
|
|||
const getGroup = async () => {
|
||||
const res = await getApiGroups()
|
||||
if (res.code === 0) {
|
||||
const groups = res.data.groups
|
||||
const groups = Array.isArray(res.data?.groups) ? res.data.groups : []
|
||||
apiGroupOptions.value = groups.map((item) => ({
|
||||
label: item,
|
||||
value: item
|
||||
}))
|
||||
apiGroupMap.value = res.data.apiGroupMap
|
||||
apiGroupMap.value =
|
||||
res.data?.apiGroupMap && typeof res.data.apiGroupMap === 'object'
|
||||
? res.data.apiGroupMap
|
||||
: {}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue