feature/use-channel-for-all #2

Merged
torjus merged 2 commits from feature/use-channel-for-all into master 2022-08-26 10:23:01 +00:00
6 changed files with 52 additions and 22 deletions
Showing only changes of commit 07fe16b72a - Show all commits

View File

@ -2,6 +2,8 @@ package store
import "git.t-juice.club/torjus/apiary/models"
var _ LoginAttemptStore = &CachingStore{}
type CachingStore struct {
backend LoginAttemptStore
@ -33,14 +35,16 @@ func NewCachingStore(backend LoginAttemptStore) *CachingStore {
//TODO: Handle better maybe?
panic(err)
}
var loginCount int
cs.totalLoginsCount = len(all)
for _, attempt := range all {
for attempt := range all {
cs.uniqueUsernames[attempt.Username] = struct{}{}
cs.uniquePasswords[attempt.Password] = struct{}{}
cs.uniqueIPs[attempt.RemoteIP.String()] = struct{}{}
cs.uniqueCountries[attempt.Country] = struct{}{}
loginCount++
}
cs.totalLoginsCount = loginCount
return cs
}
@ -58,7 +62,7 @@ func (s *CachingStore) AddAttempt(l *models.LoginAttempt) error {
return s.backend.AddAttempt(l)
}
func (s *CachingStore) All() ([]models.LoginAttempt, error) {
func (s *CachingStore) All() (<-chan models.LoginAttempt, error) {
return s.backend.All()
}

View File

@ -9,6 +9,8 @@ import (
"git.t-juice.club/torjus/apiary/models"
)
var _ LoginAttemptStore = &MemoryStore{}
type MemoryStore struct {
lock sync.RWMutex
attempts []models.LoginAttempt
@ -32,8 +34,17 @@ func (ms *MemoryStore) AddAttempt(l *models.LoginAttempt) error {
return nil
}
func (ms *MemoryStore) All() ([]models.LoginAttempt, error) {
return ms.attempts, nil
func (ms *MemoryStore) All() (<-chan models.LoginAttempt, error) {
ch := make(chan models.LoginAttempt)
go func() {
ms.lock.RLock()
defer ms.lock.RUnlock()
for _, attempt := range ms.attempts {
ch <- attempt
}
close(ch)
}()
return ch, nil
}
func (ms *MemoryStore) Stats(statType LoginStats, limit int) ([]StatsResult, error) {

View File

@ -8,6 +8,8 @@ import (
"github.com/prometheus/client_golang/prometheus"
)
var _ LoginAttemptStore = &MetricsCollectingStore{}
const tickDuration = 5 * time.Second
type MetricsCollectingStore struct {
@ -89,7 +91,7 @@ func (s *MetricsCollectingStore) AddAttempt(l *models.LoginAttempt) error {
return err
}
func (s *MetricsCollectingStore) All() ([]models.LoginAttempt, error) {
func (s *MetricsCollectingStore) All() (<-chan models.LoginAttempt, error) {
return s.store.All()
}

View File

@ -9,6 +9,8 @@ import (
_ "github.com/jackc/pgx/v4/stdlib"
)
var _ LoginAttemptStore = &PostgresStore{}
const DBSchema = `
CREATE TABLE IF NOT EXISTS login_attempts(
id serial PRIMARY KEY,
@ -61,27 +63,30 @@ RETURNING id;`
return tx.Commit()
}
func (s *PostgresStore) All() ([]models.LoginAttempt, error) {
func (s *PostgresStore) All() (<-chan models.LoginAttempt, error) {
stmt := `SELECT date, remote_ip, username, password, client_version, connection_uuid, country FROM login_attempts`
rows, err := s.db.Query(stmt)
if err != nil {
return nil, err
}
defer rows.Close()
var attempts []models.LoginAttempt
for rows.Next() {
var a models.LoginAttempt
var ip string
if err := rows.Scan(&a.Date, &ip, &a.Username, &a.Password, &a.SSHClientVersion, &a.SSHClientVersion, &a.Country); err != nil {
return nil, err
ch := make(chan models.LoginAttempt)
go func() {
defer rows.Close()
for rows.Next() {
var a models.LoginAttempt
var ip string
if err := rows.Scan(&a.Date, &ip, &a.Username, &a.Password, &a.SSHClientVersion, &a.SSHClientVersion, &a.Country); err != nil {
panic(err)
}
a.RemoteIP = net.ParseIP(ip)
ch <- a
}
a.RemoteIP = net.ParseIP(ip)
attempts = append(attempts, a)
}
close(ch)
}()
return attempts, nil
return ch, nil
}
func (s *PostgresStore) Stats(statType LoginStats, limit int) ([]StatsResult, error) {

View File

@ -38,7 +38,7 @@ type AttemptQuery struct {
}
type LoginAttemptStore interface {
AddAttempt(l *models.LoginAttempt) error
All() ([]models.LoginAttempt, error)
All() (<-chan models.LoginAttempt, error)
Stats(statType LoginStats, limit int) ([]StatsResult, error)
Query(query AttemptQuery) ([]models.LoginAttempt, error)
HealthCheker

View File

@ -25,8 +25,12 @@ func testLoginAttemptStore(s store.LoginAttemptStore, t *testing.T) {
if err != nil {
t.Fatalf("Error getting all attempts: %s", err)
}
if len(all) != len(testAttempts) {
t.Errorf("All returned wrong amount. Got %d want %d", len(all), len(testAttempts))
var count int
for range all {
count++
}
if count != len(testAttempts) {
t.Errorf("All returned wrong amount. Got %d want %d", count, len(testAttempts))
}
stats, err := s.Stats(store.LoginStatsTotals, 1)
if err != nil {
@ -229,7 +233,11 @@ func benchmarkLoginAttemptStore(setupFunc func() store.LoginAttemptStore, b *tes
if err != nil {
b.Fatalf("Error fetchin all: %s", err)
}
_ = len(all)
var count int
for range all {
count++
}
_ = count
}
})
}