Skip to content

Commit 4ecec9d

Browse files
authored
refactor: use global lock for rate limit ops (#1197)
1 parent 83f9a6f commit 4ecec9d

5 files changed

Lines changed: 283 additions & 132 deletions

File tree

‎internal/controller/user_controller.go‎

Lines changed: 101 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"github.com/tinyauthapp/tinyauth/internal/service"
1212
"github.com/tinyauthapp/tinyauth/internal/utils"
1313
"github.com/tinyauthapp/tinyauth/internal/utils/logger"
14+
"github.com/tinyauthapp/tinyauth/pkg/cache"
1415
"go.uber.org/dig"
1516

1617
"github.com/gin-gonic/gin"
@@ -72,49 +73,59 @@ func (controller *UserController) loginHandler(c *gin.Context) {
7273

7374
controller.log.App.Debug().Str("username", req.Username).Msg("Login attempt")
7475

75-
search, err := controller.auth.SearchUser(req.Username)
76+
var search *model.UserSearch
7677

77-
if err != nil {
78-
if errors.Is(err, service.ErrUserNotFound) {
79-
controller.auth.DummyPasswordCheck(req.Password)
80-
controller.log.App.Warn().Str("username", req.Username).Msg("User not found during login attempt")
81-
controller.log.AuditLoginFailure(req.Username, "unknown", c.ClientIP(), "user not found")
78+
err = controller.auth.WithLoginAttempt(req.Username, func(actions cache.CacheStoreActions[service.LoginAttempt]) error {
79+
search, err = controller.auth.SearchUser(req.Username)
80+
81+
if err != nil {
82+
if errors.Is(err, service.ErrUserNotFound) {
83+
controller.auth.DummyPasswordCheck(req.Password)
84+
controller.log.App.Warn().Str("username", req.Username).Msg("User not found during login attempt")
85+
controller.log.AuditLoginFailure(req.Username, "unknown", c.ClientIP(), "user not found")
86+
c.JSON(401, gin.H{
87+
"status": 401,
88+
"message": "Unauthorized",
89+
})
90+
return fmt.Errorf("user not found")
91+
}
92+
controller.log.App.Error().Err(err).Str("username", req.Username).Msg("Error searching for user during login attempt")
93+
c.JSON(500, gin.H{
94+
"status": 500,
95+
"message": "Internal Server Error",
96+
})
97+
return fmt.Errorf("error searching for user")
98+
}
99+
100+
isLocked, remaining := controller.auth.IsAccountLocked(req.Username, actions)
101+
102+
if isLocked {
103+
controller.log.App.Warn().Str("username", req.Username).Msg("Account is locked due to too many failed login attempts")
104+
controller.log.AuditLoginFailure(req.Username, search.Type.String(), c.ClientIP(), "account locked")
105+
c.Writer.Header().Add("x-tinyauth-lock-locked", "true")
106+
c.Writer.Header().Add("x-tinyauth-lock-reset", time.Now().Add(time.Duration(remaining)*time.Second).Format(time.RFC3339))
107+
c.JSON(429, gin.H{
108+
"status": 429,
109+
"message": fmt.Sprintf("Too many failed login attempts. Try again in %d seconds", remaining),
110+
})
111+
return fmt.Errorf("account locked")
112+
}
113+
114+
if err = controller.auth.CheckUserPassword(*search, req.Password); err != nil {
115+
controller.log.App.Warn().Str("username", req.Username).Msg("Invalid password during login attempt")
116+
controller.auth.RecordLoginAttempt(req.Username, false, actions)
117+
controller.log.AuditLoginFailure(req.Username, search.Type.String(), c.ClientIP(), "invalid password")
82118
c.JSON(401, gin.H{
83119
"status": 401,
84120
"message": "Unauthorized",
85121
})
86-
return
122+
return fmt.Errorf("invalid password")
87123
}
88-
controller.log.App.Error().Err(err).Str("username", req.Username).Msg("Error searching for user during login attempt")
89-
c.JSON(500, gin.H{
90-
"status": 500,
91-
"message": "Internal Server Error",
92-
})
93-
return
94-
}
95-
96-
isLocked, remaining := controller.auth.IsAccountLocked(req.Username)
97124

98-
if isLocked {
99-
controller.log.App.Warn().Str("username", req.Username).Msg("Account is locked due to too many failed login attempts")
100-
controller.log.AuditLoginFailure(req.Username, search.Type.String(), c.ClientIP(), "account locked")
101-
c.Writer.Header().Add("x-tinyauth-lock-locked", "true")
102-
c.Writer.Header().Add("x-tinyauth-lock-reset", time.Now().Add(time.Duration(remaining)*time.Second).Format(time.RFC3339))
103-
c.JSON(429, gin.H{
104-
"status": 429,
105-
"message": fmt.Sprintf("Too many failed login attempts. Try again in %d seconds", remaining),
106-
})
107-
return
108-
}
125+
return nil
126+
})
109127

110-
if err := controller.auth.CheckUserPassword(*search, req.Password); err != nil {
111-
controller.log.App.Warn().Str("username", req.Username).Msg("Invalid password during login attempt")
112-
controller.auth.RecordLoginAttempt(req.Username, false)
113-
controller.log.AuditLoginFailure(req.Username, search.Type.String(), c.ClientIP(), "invalid password")
114-
c.JSON(401, gin.H{
115-
"status": 401,
116-
"message": "Unauthorized",
117-
})
128+
if err != nil {
118129
return
119130
}
120131

@@ -216,7 +227,10 @@ func (controller *UserController) loginHandler(c *gin.Context) {
216227

217228
controller.log.AuditLoginSuccess(req.Username, search.Type.String(), c.ClientIP())
218229

219-
controller.auth.RecordLoginAttempt(req.Username, true)
230+
controller.auth.WithLoginAttempt(req.Username, func(actions cache.CacheStoreActions[service.LoginAttempt]) error {
231+
controller.auth.RecordLoginAttempt(req.Username, true, actions)
232+
return nil
233+
})
220234

221235
c.JSON(200, gin.H{
222236
"status": 200,
@@ -317,56 +331,66 @@ func (controller *UserController) totpHandler(c *gin.Context) {
317331

318332
controller.log.App.Debug().Str("username", context.GetUsername()).Msg("TOTP verification attempt")
319333

320-
isLocked, remaining := controller.auth.IsAccountLocked(context.GetUsername())
334+
var user *model.LocalUser
321335

322-
if isLocked {
323-
controller.log.App.Warn().Str("username", context.GetUsername()).Msg("Account is locked due to too many failed TOTP attempts")
324-
controller.log.AuditLoginFailure(context.GetUsername(), "local", c.ClientIP(), "account locked")
325-
c.Writer.Header().Add("x-tinyauth-lock-locked", "true")
326-
c.Writer.Header().Add("x-tinyauth-lock-reset", time.Now().Add(time.Duration(remaining)*time.Second).Format(time.RFC3339))
327-
c.JSON(429, gin.H{
328-
"status": 429,
329-
"message": fmt.Sprintf("Too many failed TOTP attempts. Try again in %d seconds", remaining),
330-
})
331-
return
332-
}
336+
err = controller.auth.WithLoginAttempt(context.GetUsername(), func(actions cache.CacheStoreActions[service.LoginAttempt]) error {
337+
isLocked, remaining := controller.auth.IsAccountLocked(context.GetUsername(), actions)
333338

334-
user := controller.auth.GetLocalUser(context.GetUsername())
339+
if isLocked {
340+
controller.log.App.Warn().Str("username", context.GetUsername()).Msg("Account is locked due to too many failed TOTP attempts")
341+
controller.log.AuditLoginFailure(context.GetUsername(), "local", c.ClientIP(), "account locked")
342+
c.Writer.Header().Add("x-tinyauth-lock-locked", "true")
343+
c.Writer.Header().Add("x-tinyauth-lock-reset", time.Now().Add(time.Duration(remaining)*time.Second).Format(time.RFC3339))
344+
c.JSON(429, gin.H{
345+
"status": 429,
346+
"message": fmt.Sprintf("Too many failed TOTP attempts. Try again in %d seconds", remaining),
347+
})
348+
return fmt.Errorf("account locked")
349+
}
335350

336-
if user == nil {
337-
controller.log.App.Error().Str("username", context.GetUsername()).Msg("Local user not found during TOTP verification")
338-
c.JSON(401, gin.H{
339-
"status": 401,
340-
"message": "Unauthorized",
341-
})
342-
return
343-
}
351+
user = controller.auth.GetLocalUser(context.GetUsername())
344352

345-
ok := totp.Validate(req.Code, user.TOTPSecret)
353+
if user == nil {
354+
controller.log.App.Error().Str("username", context.GetUsername()).Msg("Local user not found during TOTP verification")
355+
c.JSON(401, gin.H{
356+
"status": 401,
357+
"message": "Unauthorized",
358+
})
359+
return fmt.Errorf("local user not found")
360+
}
346361

347-
if !ok {
348-
controller.log.App.Warn().Str("username", context.GetUsername()).Msg("Invalid TOTP code during verification attempt")
349-
controller.auth.RecordLoginAttempt(context.GetUsername(), false)
350-
controller.log.AuditLoginFailure(context.GetUsername(), "local", c.ClientIP(), "invalid TOTP code")
351-
c.JSON(401, gin.H{
352-
"status": 401,
353-
"message": "Unauthorized",
354-
})
355-
return
356-
}
362+
ok := totp.Validate(req.Code, user.TOTPSecret)
357363

358-
uuid, err := c.Cookie(controller.runtime.SessionCookieName)
364+
if !ok {
365+
controller.log.App.Warn().Str("username", context.GetUsername()).Msg("Invalid TOTP code during verification attempt")
366+
controller.auth.RecordLoginAttempt(context.GetUsername(), false, actions)
367+
controller.log.AuditLoginFailure(context.GetUsername(), "local", c.ClientIP(), "invalid TOTP code")
368+
c.JSON(401, gin.H{
369+
"status": 401,
370+
"message": "Unauthorized",
371+
})
372+
return fmt.Errorf("invalid TOTP code")
373+
}
359374

360-
if err == nil {
361-
_, err = controller.auth.DeleteSession(c, uuid)
362-
if err != nil {
363-
controller.log.App.Error().Err(err).Msg("Failed to delete pending TOTP session after successful verification")
375+
uuid, err := c.Cookie(controller.runtime.SessionCookieName)
376+
377+
if err == nil {
378+
_, err = controller.auth.DeleteSession(c, uuid)
379+
if err != nil {
380+
controller.log.App.Error().Err(err).Msg("Failed to delete pending TOTP session after successful verification")
381+
}
382+
} else {
383+
controller.log.App.Warn().Err(err).Msg("Failed to retrieve session cookie for pending TOTP session, cannot delete it")
364384
}
365-
} else {
366-
controller.log.App.Warn().Err(err).Msg("Failed to retrieve session cookie for pending TOTP session, cannot delete it")
367-
}
368385

369-
controller.auth.RecordLoginAttempt(context.GetUsername(), true)
386+
controller.auth.RecordLoginAttempt(context.GetUsername(), true, actions)
387+
388+
return nil
389+
})
390+
391+
if err != nil {
392+
return
393+
}
370394

371395
sessionCookie := repository.Session{
372396
Username: user.Username,

‎internal/middleware/context_middleware.go‎

Lines changed: 43 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"github.com/tinyauthapp/tinyauth/internal/service"
1313
"github.com/tinyauthapp/tinyauth/internal/utils"
1414
"github.com/tinyauthapp/tinyauth/internal/utils/logger"
15+
"github.com/tinyauthapp/tinyauth/pkg/cache"
1516
"go.uber.org/dig"
1617

1718
"github.com/gin-gonic/gin"
@@ -253,34 +254,58 @@ func (m *ContextMiddleware) cookieAuth(ctx context.Context, uuid string, ip stri
253254
// any response headers to set.
254255
func (m *ContextMiddleware) basicAuth(username string, password string) (*model.UserContext, map[string]string, error) {
255256
headers := make(map[string]string)
257+
errLocked := errors.New("account locked")
256258
userContext := new(model.UserContext)
257-
locked, remaining := m.auth.IsAccountLocked(username)
258259

259-
if locked {
260-
m.log.App.Debug().Msgf("Account for user %s is locked for %d seconds, denying auth", username, remaining)
261-
headers["x-tinyauth-lock-locked"] = "true"
262-
headers["x-tinyauth-lock-reset"] = time.Now().Add(time.Duration(remaining) * time.Second).Format(time.RFC3339)
263-
return nil, headers, nil
264-
}
260+
var search *model.UserSearch
265261

266-
search, err := m.auth.SearchUser(username)
262+
err := m.auth.WithLoginAttempt(username, func(actions cache.CacheStoreActions[service.LoginAttempt]) error {
263+
locked, remaining := m.auth.IsAccountLocked(username, actions)
267264

268-
if err != nil {
269-
if errors.Is(err, service.ErrUserNotFound) {
270-
m.auth.DummyPasswordCheck(password)
265+
if locked {
266+
m.log.App.Debug().Msgf("Account for user %s is locked for %d seconds, denying auth", username, remaining)
267+
headers["x-tinyauth-lock-locked"] = "true"
268+
headers["x-tinyauth-lock-reset"] = time.Now().Add(time.Duration(remaining) * time.Second).Format(time.RFC3339)
269+
return errLocked
271270
}
272-
return nil, nil, fmt.Errorf("error searching for user: %w", err)
273-
}
274271

275-
err = m.auth.CheckUserPassword(*search, password)
272+
var err error
273+
search, err = m.auth.SearchUser(username)
274+
275+
if err != nil {
276+
if errors.Is(err, service.ErrUserNotFound) {
277+
m.auth.DummyPasswordCheck(password)
278+
}
279+
return fmt.Errorf("error searching for user: %w", err)
280+
}
281+
282+
err = m.auth.CheckUserPassword(*search, password)
283+
284+
if err != nil {
285+
m.auth.RecordLoginAttempt(username, false, actions)
286+
return fmt.Errorf("invalid password for basic auth user: %w", err)
287+
}
288+
289+
if search.Type == model.UserLocal {
290+
user := m.auth.GetLocalUser(username)
291+
292+
if user != nil && user.TOTPSecret != "" {
293+
return fmt.Errorf("user with totp not allowed to login via basic auth: %s", username)
294+
}
295+
}
296+
297+
m.auth.RecordLoginAttempt(username, true, actions)
298+
299+
return nil
300+
})
276301

277302
if err != nil {
278-
m.auth.RecordLoginAttempt(username, false)
279-
return nil, nil, fmt.Errorf("invalid password for basic auth user: %w", err)
303+
if errors.Is(err, errLocked) {
304+
return nil, headers, nil
305+
}
306+
return nil, nil, err
280307
}
281308

282-
m.auth.RecordLoginAttempt(username, true)
283-
284309
switch search.Type {
285310
case model.UserLocal:
286311
user := m.auth.GetLocalUser(username)
@@ -289,10 +314,6 @@ func (m *ContextMiddleware) basicAuth(username string, password string) (*model.
289314
return nil, nil, fmt.Errorf("user not found locally: %s", username)
290315
}
291316

292-
if user.TOTPSecret != "" {
293-
return nil, nil, fmt.Errorf("user with totp not allowed to login via basic auth: %s", username)
294-
}
295-
296317
userContext.Local = &model.LocalContext{
297318
BaseContext: model.BaseContext{
298319
Username: user.Username,

0 commit comments

Comments
 (0)