soccer-bot/antispam/antispam.go

package antispam

import (
	"fmt"
	"strings"
	"sync"
	"time"
)

func normalizeUser(from string) string {
	from = strings.ToLower(from)

	if strings.Contains(from, "!") {
		split := strings.SplitN(from, "!", 2)
		return fmt.Sprintf("*!%s", split[1])
	}

	return from
}

func normalizeCommand(command string) string {
	return strings.ToLower(command)
}

type AntiSpam struct {
	userTrackingMap  map[string][]time.Time
	userTrackingLock sync.RWMutex

	UserCooldownDuration time.Duration

	UserRateLimitDuration time.Duration
	UserRateLimitCount    int

	commandTrackingMap  map[string][]time.Time
	commandTrackingLock sync.RWMutex

	CommandCooldownDuration time.Duration

	CommandRateLimitDuration time.Duration
	CommandRateLimitCount    int
}

func New() *AntiSpam {
	as := new(AntiSpam)
	as.init()
	return as
}

func (m *AntiSpam) init() {
	m.userTrackingMap = map[string][]time.Time{}
	m.commandTrackingMap = map[string][]time.Time{}

	m.CommandCooldownDuration = 15 * time.Second
	m.CommandRateLimitCount = 2
	m.CommandRateLimitDuration = 1 * time.Minute

	m.UserCooldownDuration = 3 * time.Second
	m.UserRateLimitCount = 4
	m.UserRateLimitDuration = 1 * time.Minute
}

func (m *AntiSpam) checkUserCooldown(from string) (duration time.Duration) {
	from = normalizeUser(from)

	m.userTrackingLock.RLock()
	defer m.userTrackingLock.RUnlock()

	// Are there any times registered at all for this entry?
	if timestamps := m.userTrackingMap[from]; timestamps != nil {

		var mostRecentTimestamp time.Time

		// Figure out most recent time
		for _, timestamp := range timestamps {
			if mostRecentTimestamp.Before(timestamp) {
				mostRecentTimestamp = timestamp
			}
		}

		duration = mostRecentTimestamp.Add(m.UserCooldownDuration).Sub(time.Now())
		if duration < 0 {
			duration = 0
		}
	}

	return
}

func (m *AntiSpam) checkCommandCooldown(command string) (duration time.Duration) {
	command = normalizeCommand(command)

	m.commandTrackingLock.RLock()
	defer m.commandTrackingLock.RUnlock()

	// Are there any times registered at all for this entry?
	if timestamps := m.commandTrackingMap[command]; timestamps != nil {

		var mostRecentTimestamp time.Time

		// Figure out most recent time
		for _, timestamp := range timestamps {
			if mostRecentTimestamp.Before(timestamp) {
				mostRecentTimestamp = timestamp
			}
		}

		duration = mostRecentTimestamp.Add(m.CommandCooldownDuration).Sub(time.Now())
		if duration < 0 {
			duration = 0
		}
	}

	return
}

func (m *AntiSpam) checkUserRateLimit(from string) (duration time.Duration) {
	from = normalizeUser(from)

	m.userTrackingLock.RLock()
	defer m.userTrackingLock.RUnlock()

	// Are there any times registered at all for this entry?
	if timestamps := m.userTrackingMap[from]; timestamps != nil && len(timestamps) >= m.UserRateLimitCount {
		leastRecentTimestamp := timestamps[len(timestamps)-m.UserRateLimitCount]

		duration = leastRecentTimestamp.Add(m.UserRateLimitDuration).Sub(time.Now())
		if duration < 0 {
			duration = 0
		}
	}

	return
}

func (m *AntiSpam) checkCommandRateLimit(command string) (duration time.Duration) {
	command = normalizeCommand(command)

	m.commandTrackingLock.RLock()
	defer m.commandTrackingLock.RUnlock()

	// Are there any times registered at all for this entry?
	if timestamps := m.commandTrackingMap[command]; timestamps != nil && len(timestamps) >= m.CommandRateLimitCount {
		leastRecentTimestamp := timestamps[len(timestamps)-m.CommandRateLimitCount]

		duration = leastRecentTimestamp.Add(m.CommandRateLimitDuration).Sub(time.Now())
		if duration < 0 {
			duration = 0
		}
	}

	return
}

func (m *AntiSpam) track(from string, command string) {
	currentTime := time.Now()

	command = normalizeCommand(command)
	from = normalizeUser(from)

	// User tracking
	go func() {
		// Register current time
		m.userTrackingLock.Lock()
		if m.userTrackingMap[from] == nil {
			m.userTrackingMap[from] = []time.Time{}
		}
		m.userTrackingMap[from] = append(m.userTrackingMap[from], currentTime)
		m.userTrackingLock.Unlock()

		// Wait for time to run down
		var waitDuration time.Duration
		if m.UserCooldownDuration > m.UserRateLimitDuration {
			waitDuration = m.UserCooldownDuration
		} else {
			waitDuration = m.UserRateLimitDuration
		}
		time.Sleep(waitDuration)

		// Remove current time
		m.userTrackingLock.Lock()
		indexToDelete := -1
		for i, time := range m.userTrackingMap[from] {
			if time == currentTime {
				indexToDelete = i
				break
			}
		}
		if indexToDelete >= 0 {
			m.userTrackingMap[from] = append(
				m.userTrackingMap[from][:indexToDelete],
				m.userTrackingMap[from][indexToDelete+1:]...)
			if len(m.userTrackingMap) == 0 {
				m.userTrackingMap[from] = nil
			}
		}
		m.userTrackingLock.Unlock()
	}()

	// Command tracking
	go func() {
		// Register current time
		m.commandTrackingLock.Lock()
		if m.commandTrackingMap[command] == nil {
			m.commandTrackingMap[command] = []time.Time{}
		}
		m.commandTrackingMap[command] = append(m.commandTrackingMap[command], currentTime)
		m.commandTrackingLock.Unlock()

		// Wait for time to run down
		var waitDuration time.Duration
		if m.CommandCooldownDuration > m.CommandRateLimitDuration {
			waitDuration = m.CommandCooldownDuration
		} else {
			waitDuration = m.CommandRateLimitDuration
		}
		time.Sleep(waitDuration)

		// Remove current time
		m.commandTrackingLock.Lock()
		indexToDelete := -1
		for i, time := range m.commandTrackingMap[command] {
			if time == currentTime {
				indexToDelete = i
				break
			}
		}
		if indexToDelete >= 0 {
			m.commandTrackingMap[command] = append(
				m.commandTrackingMap[command][:indexToDelete],
				m.commandTrackingMap[command][indexToDelete+1:]...)
			if len(m.commandTrackingMap) == 0 {
				m.commandTrackingMap[command] = nil
			}
		}
		m.commandTrackingLock.Unlock()
	}()
}

func (m *AntiSpam) Check(from string, command string) (allow bool, waitDuration time.Duration) {
	allow = true

	// if this user has run a command in the last N seconds (user cooldown)
	if currentDuration := m.checkUserCooldown(from); currentDuration > 0 {
		allow = false
		if currentDuration > waitDuration {
			waitDuration = currentDuration
		}
	}

	// if this user has run X commands in the last N seconds (user rate limit)
	if currentDuration := m.checkUserRateLimit(from); currentDuration > 0 {
		allow = false
		if currentDuration > waitDuration {
			waitDuration = currentDuration
		}
	}

	// if command has been run just previously (command cooldown)
	if currentDuration := m.checkCommandCooldown(command); currentDuration > 0 {
		allow = false
		if currentDuration > waitDuration {
			waitDuration = currentDuration
		}
	}

	// if command has been run X (1?) times in the last N seconds (command rate limit)
	if currentDuration := m.checkCommandRateLimit(command); currentDuration > 0 {
		allow = false
		if currentDuration > waitDuration {
			waitDuration = currentDuration
		}
	}

	if allow {
		m.track(from, command)
	}

	return
}