renjue
66 lines
1.6 KiB
Go
66 lines
1.6 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"time"
|
|
)
|
|
|
|
type LoginAttempt struct {
|
|
FailCount int
|
|
WindowStart time.Time
|
|
LockedUntil *time.Time
|
|
}
|
|
|
|
func (s *Store) GetLoginAttempt(ctx context.Context, ip string) (*LoginAttempt, error) {
|
|
var failCount int
|
|
var windowStart string
|
|
var lockedUntil sql.NullString
|
|
err := s.db.QueryRowContext(ctx,
|
|
`SELECT fail_count, window_start, locked_until FROM login_attempts WHERE ip = ?`, ip,
|
|
).Scan(&failCount, &windowStart, &lockedUntil)
|
|
if err == sql.ErrNoRows {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
out := &LoginAttempt{
|
|
FailCount: failCount,
|
|
WindowStart: parseTime(windowStart),
|
|
}
|
|
if lockedUntil.Valid && lockedUntil.String != "" {
|
|
t := parseTime(lockedUntil.String)
|
|
out.LockedUntil = &t
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (s *Store) SetLoginAttempt(ctx context.Context, ip string, attempt LoginAttempt) error {
|
|
var locked string
|
|
if attempt.LockedUntil != nil {
|
|
locked = attempt.LockedUntil.Format("2006-01-02 15:04:05")
|
|
}
|
|
_, err := s.db.ExecContext(ctx, `
|
|
INSERT INTO login_attempts(ip, fail_count, window_start, locked_until)
|
|
VALUES(?,?,?,?)
|
|
ON CONFLICT(ip) DO UPDATE SET
|
|
fail_count = excluded.fail_count,
|
|
window_start = excluded.window_start,
|
|
locked_until = excluded.locked_until
|
|
`, ip, attempt.FailCount, attempt.WindowStart.Format("2006-01-02 15:04:05"), nullIfEmpty(locked))
|
|
return err
|
|
}
|
|
|
|
func (s *Store) ClearLoginAttempt(ctx context.Context, ip string) error {
|
|
_, err := s.db.ExecContext(ctx, `DELETE FROM login_attempts WHERE ip = ?`, ip)
|
|
return err
|
|
}
|
|
|
|
func nullIfEmpty(s string) any {
|
|
if s == "" {
|
|
return nil
|
|
}
|
|
return s
|
|
}
|