Quota per user
This commit is contained in:
@@ -0,0 +1,304 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
const (
|
||||
quotaActionBlock = "block"
|
||||
quotaActionThrottle = "throttle"
|
||||
)
|
||||
|
||||
var errDataQuotaExceeded = errors.New("data quota exceeded")
|
||||
|
||||
func normalizeQuotaAction(v string) string {
|
||||
if strings.EqualFold(strings.TrimSpace(v), quotaActionThrottle) {
|
||||
return quotaActionThrottle
|
||||
}
|
||||
return quotaActionBlock
|
||||
}
|
||||
|
||||
func quotaThrottleMbpsOrDefault(v int) int {
|
||||
if v <= 0 {
|
||||
return 1
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
type sshTrafficDelta struct {
|
||||
Uplink int64
|
||||
Downlink int64
|
||||
}
|
||||
|
||||
var sshTrafficPersistenceMu sync.Mutex
|
||||
|
||||
func (s *Store) AddSSHUserTrafficBatch(ctx context.Context, deltas map[string]sshTrafficDelta) error {
|
||||
if s == nil || len(deltas) == 0 {
|
||||
return nil
|
||||
}
|
||||
tx, err := s.db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
stmt, err := tx.PrepareContext(ctx, `
|
||||
UPDATE ssh_users SET
|
||||
total_uplink_bytes = GREATEST(total_uplink_bytes + GREATEST($2::BIGINT, 0), 0),
|
||||
total_downlink_bytes = GREATEST(total_downlink_bytes + GREATEST($3::BIGINT, 0), 0)
|
||||
WHERE username = $1`)
|
||||
if err != nil {
|
||||
_ = tx.Rollback()
|
||||
return err
|
||||
}
|
||||
defer stmt.Close()
|
||||
for username, d := range deltas {
|
||||
if strings.TrimSpace(username) == "" || (d.Uplink == 0 && d.Downlink == 0) {
|
||||
continue
|
||||
}
|
||||
if _, err := stmt.ExecContext(ctx, username, d.Uplink, d.Downlink); err != nil {
|
||||
_ = tx.Rollback()
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func (s *Store) ResetSSHUserTraffic(ctx context.Context, username string) error {
|
||||
if s == nil || strings.TrimSpace(username) == "" {
|
||||
return nil
|
||||
}
|
||||
_, err := s.db.ExecContext(ctx, `
|
||||
UPDATE ssh_users
|
||||
SET total_uplink_bytes = 0, total_downlink_bytes = 0
|
||||
WHERE username = $1`, username)
|
||||
return err
|
||||
}
|
||||
|
||||
func initSSHRuntimeUsage(u *UserState, uplink, downlink int64) {
|
||||
if u == nil {
|
||||
return
|
||||
}
|
||||
if uplink < 0 {
|
||||
uplink = 0
|
||||
}
|
||||
if downlink < 0 {
|
||||
downlink = 0
|
||||
}
|
||||
atomic.StoreInt64(&u.TotalUplinkBytes, uplink)
|
||||
atomic.StoreInt64(&u.TotalDownlinkBytes, downlink)
|
||||
atomic.StoreInt64(&u.totalBytes, uplink+downlink)
|
||||
atomic.StoreInt64(&u.pendingUplinkBytes, 0)
|
||||
atomic.StoreInt64(&u.pendingDownlinkBytes, 0)
|
||||
}
|
||||
|
||||
func resetSSHRuntimeUsage(username string) {
|
||||
u, ok := userMgr.Get(username)
|
||||
if !ok || u == nil {
|
||||
return
|
||||
}
|
||||
initSSHRuntimeUsage(u, 0, 0)
|
||||
u.mu.Lock()
|
||||
u.quotaLimiter = nil
|
||||
u.quotaLimiterMbps = 0
|
||||
u.mu.Unlock()
|
||||
}
|
||||
|
||||
func resetSSHUserTrafficAccounting(ctx context.Context, store *Store, username string) error {
|
||||
sshTrafficPersistenceMu.Lock()
|
||||
defer sshTrafficPersistenceMu.Unlock()
|
||||
if err := store.ResetSSHUserTraffic(ctx, username); err != nil {
|
||||
return err
|
||||
}
|
||||
resetSSHRuntimeUsage(username)
|
||||
return nil
|
||||
}
|
||||
|
||||
func sshUserQuotaBlocked(u *UserState) bool {
|
||||
if u == nil {
|
||||
return false
|
||||
}
|
||||
u.mu.Lock()
|
||||
quota := u.Cfg.DataQuotaBytes
|
||||
action := normalizeQuotaAction(u.Cfg.QuotaAction)
|
||||
u.mu.Unlock()
|
||||
return quota > 0 && action == quotaActionBlock && atomic.LoadInt64(&u.totalBytes) >= quota
|
||||
}
|
||||
|
||||
func sshQuotaLimiter(u *UserState, mbps int) *rate.Limiter {
|
||||
mbps = quotaThrottleMbpsOrDefault(mbps)
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
if u.quotaLimiter == nil || u.quotaLimiterMbps != mbps {
|
||||
bps := mbpsToBytesPerSec(mbps)
|
||||
burst := int(bps)
|
||||
if burst < copyBufSize {
|
||||
burst = copyBufSize
|
||||
}
|
||||
u.quotaLimiter = rate.NewLimiter(rate.Limit(bps), burst)
|
||||
u.quotaLimiterMbps = mbps
|
||||
}
|
||||
return u.quotaLimiter
|
||||
}
|
||||
|
||||
func reserveSSHUserBytes(u *UserState, requested int) (allowed int, throttle *rate.Limiter, stopAfter bool) {
|
||||
if u == nil || requested <= 0 {
|
||||
return 0, nil, false
|
||||
}
|
||||
u.mu.Lock()
|
||||
quota := u.Cfg.DataQuotaBytes
|
||||
action := normalizeQuotaAction(u.Cfg.QuotaAction)
|
||||
throttleMbps := u.Cfg.QuotaThrottleMbps
|
||||
u.mu.Unlock()
|
||||
|
||||
n := int64(requested)
|
||||
if quota <= 0 {
|
||||
atomic.AddInt64(&u.totalBytes, n)
|
||||
return requested, nil, false
|
||||
}
|
||||
|
||||
if action == quotaActionThrottle {
|
||||
previous := atomic.AddInt64(&u.totalBytes, n) - n
|
||||
if previous+n > quota {
|
||||
return requested, sshQuotaLimiter(u, throttleMbps), false
|
||||
}
|
||||
return requested, nil, false
|
||||
}
|
||||
|
||||
for {
|
||||
used := atomic.LoadInt64(&u.totalBytes)
|
||||
remaining := quota - used
|
||||
if remaining <= 0 {
|
||||
return 0, nil, true
|
||||
}
|
||||
take := n
|
||||
if take > remaining {
|
||||
take = remaining
|
||||
}
|
||||
if atomic.CompareAndSwapInt64(&u.totalBytes, used, used+take) {
|
||||
return int(take), nil, take < n || used+take >= quota
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func finishSSHUserReservation(u *UserState, uplink bool, reserved, written int) {
|
||||
if u == nil || reserved <= 0 {
|
||||
return
|
||||
}
|
||||
if written < 0 {
|
||||
written = 0
|
||||
}
|
||||
if written > reserved {
|
||||
written = reserved
|
||||
}
|
||||
if written < reserved {
|
||||
atomic.AddInt64(&u.totalBytes, -int64(reserved-written))
|
||||
}
|
||||
if written == 0 {
|
||||
return
|
||||
}
|
||||
if uplink {
|
||||
atomic.AddInt64(&u.TotalUplinkBytes, int64(written))
|
||||
atomic.AddInt64(&u.pendingUplinkBytes, int64(written))
|
||||
} else {
|
||||
atomic.AddInt64(&u.TotalDownlinkBytes, int64(written))
|
||||
atomic.AddInt64(&u.pendingDownlinkBytes, int64(written))
|
||||
}
|
||||
}
|
||||
|
||||
type sshQuotaWriter struct {
|
||||
w io.Writer
|
||||
user *UserState
|
||||
uplink bool
|
||||
}
|
||||
|
||||
func (qw sshQuotaWriter) Write(p []byte) (int, error) {
|
||||
allowed, quotaLimiter, stopAfter := reserveSSHUserBytes(qw.user, len(p))
|
||||
if allowed <= 0 {
|
||||
return 0, errDataQuotaExceeded
|
||||
}
|
||||
if quotaLimiter != nil {
|
||||
if err := quotaLimiter.WaitN(context.Background(), allowed); err != nil {
|
||||
finishSSHUserReservation(qw.user, qw.uplink, allowed, 0)
|
||||
return 0, err
|
||||
}
|
||||
}
|
||||
n, err := qw.w.Write(p[:allowed])
|
||||
finishSSHUserReservation(qw.user, qw.uplink, allowed, n)
|
||||
if err != nil {
|
||||
return n, err
|
||||
}
|
||||
if stopAfter || allowed < len(p) {
|
||||
return n, errDataQuotaExceeded
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func startSSHUserTrafficFlusher(store *Store) {
|
||||
if store == nil {
|
||||
return
|
||||
}
|
||||
go func() {
|
||||
ticker := time.NewTicker(5 * time.Second)
|
||||
defer ticker.Stop()
|
||||
for range ticker.C {
|
||||
flushSSHUserTraffic(store)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func flushSSHUserTraffic(store *Store) {
|
||||
if store == nil {
|
||||
return
|
||||
}
|
||||
sshTrafficPersistenceMu.Lock()
|
||||
defer sshTrafficPersistenceMu.Unlock()
|
||||
deltas := make(map[string]sshTrafficDelta)
|
||||
states := make(map[string]*UserState)
|
||||
for _, u := range userMgr.List() {
|
||||
if u == nil || strings.TrimSpace(u.Cfg.Username) == "" {
|
||||
continue
|
||||
}
|
||||
up := atomic.SwapInt64(&u.pendingUplinkBytes, 0)
|
||||
down := atomic.SwapInt64(&u.pendingDownlinkBytes, 0)
|
||||
if up == 0 && down == 0 {
|
||||
continue
|
||||
}
|
||||
username := u.Cfg.Username
|
||||
deltas[username] = sshTrafficDelta{Uplink: up, Downlink: down}
|
||||
states[username] = u
|
||||
}
|
||||
if len(deltas) == 0 {
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
if err := store.AddSSHUserTrafficBatch(ctx, deltas); err != nil {
|
||||
log.Printf("ssh traffic flush failed: %v", err)
|
||||
for username, d := range deltas {
|
||||
if u := states[username]; u != nil {
|
||||
atomic.AddInt64(&u.pendingUplinkBytes, d.Uplink)
|
||||
atomic.AddInt64(&u.pendingDownlinkBytes, d.Downlink)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func validateQuotaConfig(quotaBytes int64, action string, throttleMbps int) error {
|
||||
if quotaBytes < 0 {
|
||||
return fmt.Errorf("data_quota_bytes must be non-negative")
|
||||
}
|
||||
action = normalizeQuotaAction(action)
|
||||
if quotaBytes > 0 && action == quotaActionThrottle && throttleMbps < 0 {
|
||||
return fmt.Errorf("quota_throttle_mbps must be non-negative")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user