305 lines
7.1 KiB
Go
305 lines
7.1 KiB
Go
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
|
|
}
|