mirror of https://github.com/cbeuw/Cloak
QOS and user managing, bug fixes
Qian Wang
6 years ago
parent
6a6b293164
commit
3534d05055
@ -0,0 +1,58 @@
|
||||
package multiplex
|
||||
|
||||
import (
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/juju/ratelimit"
|
||||
)
|
||||
|
||||
// Valve needs to be universal, across all sessions that belong to a user
|
||||
// gabe please don't sue
|
||||
type Valve struct {
|
||||
// traffic directions from the server's perspective are refered
|
||||
// exclusively as rx and tx.
|
||||
// rx is from client to server, tx is from server to client
|
||||
// DO NOT use terms up or down as this is used in usermanager
|
||||
// for bandwidth limiting
|
||||
rxtb atomic.Value // *ratelimit.Bucket
|
||||
txtb atomic.Value // *ratelimit.Bucket
|
||||
|
||||
rxCredit int64
|
||||
txCredit int64
|
||||
}
|
||||
|
||||
func MakeValve(rxRate, txRate, rxCredit, txCredit int64) *Valve {
|
||||
v := &Valve{
|
||||
rxCredit: rxCredit,
|
||||
txCredit: txCredit,
|
||||
}
|
||||
v.SetRxRate(rxRate)
|
||||
v.SetTxRate(txRate)
|
||||
return v
|
||||
}
|
||||
|
||||
func (v *Valve) SetRxRate(rate int64) {
|
||||
v.rxtb.Store(ratelimit.NewBucketWithRate(float64(rate), rate))
|
||||
}
|
||||
|
||||
func (v *Valve) SetTxRate(rate int64) {
|
||||
v.txtb.Store(ratelimit.NewBucketWithRate(float64(rate), rate))
|
||||
}
|
||||
|
||||
func (v *Valve) rxWait(n int) {
|
||||
v.rxtb.Load().(*ratelimit.Bucket).Wait(int64(n))
|
||||
}
|
||||
|
||||
func (v *Valve) txWait(n int) {
|
||||
v.txtb.Load().(*ratelimit.Bucket).Wait(int64(n))
|
||||
}
|
||||
|
||||
// n can be negative
|
||||
func (v *Valve) AddRxCredit(n int64) int64 {
|
||||
return atomic.AddInt64(&v.rxCredit, n)
|
||||
}
|
||||
|
||||
// n can be negative
|
||||
func (v *Valve) AddTxCredit(n int64) int64 {
|
||||
return atomic.AddInt64(&v.txCredit, n)
|
||||
}
|
||||
@ -0,0 +1,86 @@
|
||||
package usermanager
|
||||
|
||||
import (
|
||||
mux "github.com/cbeuw/Cloak/internal/multiplex"
|
||||
"log"
|
||||
"net"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
/*
|
||||
type userParams struct {
|
||||
sessionsCap uint32
|
||||
upRate int64
|
||||
downRate int64
|
||||
upCredit int64
|
||||
downCredit int64
|
||||
}
|
||||
*/
|
||||
|
||||
type user struct {
|
||||
up *Userpanel
|
||||
|
||||
uid [32]byte
|
||||
|
||||
sessionsCap uint32 //userParams
|
||||
|
||||
valve *mux.Valve
|
||||
|
||||
sessionsM sync.RWMutex
|
||||
sessions map[uint32]*mux.Session
|
||||
}
|
||||
|
||||
func MakeUser(up *Userpanel, uid [32]byte, sessionsCap uint32, upRate, downRate, upCredit, downCredit int64) *user {
|
||||
valve := mux.MakeValve(upRate, downRate, upCredit, downCredit)
|
||||
u := &user{
|
||||
up: up,
|
||||
uid: uid,
|
||||
valve: valve,
|
||||
sessionsCap: sessionsCap,
|
||||
sessions: make(map[uint32]*mux.Session),
|
||||
}
|
||||
return u
|
||||
}
|
||||
|
||||
func (u *user) setSessionsCap(cap uint32) {
|
||||
atomic.StoreUint32(&u.sessionsCap, cap)
|
||||
}
|
||||
|
||||
func (u *user) GetSession(sessionID uint32) *mux.Session {
|
||||
u.sessionsM.RLock()
|
||||
defer u.sessionsM.RUnlock()
|
||||
if sesh, ok := u.sessions[sessionID]; ok {
|
||||
return sesh
|
||||
} else {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (u *user) PutSession(sessionID uint32, sesh *mux.Session) {
|
||||
u.sessionsM.Lock()
|
||||
u.sessions[sessionID] = sesh
|
||||
u.sessionsM.Unlock()
|
||||
}
|
||||
|
||||
func (u *user) DelSession(sessionID uint32) {
|
||||
u.sessionsM.Lock()
|
||||
delete(u.sessions, sessionID)
|
||||
if len(u.sessions) == 0 {
|
||||
u.sessionsM.Unlock()
|
||||
u.up.delActiveUser(u.uid)
|
||||
return
|
||||
}
|
||||
u.sessionsM.Unlock()
|
||||
}
|
||||
|
||||
func (u *user) GetOrCreateSession(sessionID uint32, obfs func(*mux.Frame) []byte, deobfs func([]byte) *mux.Frame, obfsedRead func(net.Conn, []byte) (int, error)) (sesh *mux.Session) {
|
||||
log.Printf("getting sessionID %v\n", sessionID)
|
||||
if sesh = u.GetSession(sessionID); sesh != nil {
|
||||
return
|
||||
} else {
|
||||
sesh = mux.MakeSession(sessionID, u.valve, obfs, deobfs, obfsedRead)
|
||||
u.PutSession(sessionID, sesh)
|
||||
return
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,151 @@
|
||||
package usermanager
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"github.com/boltdb/bolt"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type Userpanel struct {
|
||||
db *bolt.DB
|
||||
|
||||
activeUsersM sync.RWMutex
|
||||
activeUsers map[[32]byte]*user
|
||||
}
|
||||
|
||||
func MakeUserpanel(dbPath string) (*Userpanel, error) {
|
||||
db, err := bolt.Open(dbPath, 0600, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
up := &Userpanel{
|
||||
db: db,
|
||||
activeUsers: make(map[[32]byte]*user),
|
||||
}
|
||||
return up, nil
|
||||
}
|
||||
|
||||
var ErrUserNotFound = errors.New("User does not exist in memory or db")
|
||||
|
||||
// GetUser is used to retrieve a user if s/he is active, or to retrieve the user's infor
|
||||
// from the db and mark it as an active user
|
||||
func (up *Userpanel) GetAndActivateUser(UID [32]byte) (*user, error) {
|
||||
up.activeUsersM.RLock()
|
||||
if user, ok := up.activeUsers[UID]; ok {
|
||||
up.activeUsersM.RUnlock()
|
||||
return user, nil
|
||||
}
|
||||
up.activeUsersM.RUnlock()
|
||||
|
||||
var sessionsCap uint32
|
||||
var upRate, downRate, upCredit, downCredit int64
|
||||
err := up.db.View(func(tx *bolt.Tx) error {
|
||||
b := tx.Bucket(UID[:])
|
||||
if b == nil {
|
||||
return ErrUserNotFound
|
||||
}
|
||||
sessionsCap = binary.BigEndian.Uint32(b.Get([]byte("sessionsCap")))
|
||||
upRate = int64(binary.BigEndian.Uint64(b.Get([]byte("upRate"))))
|
||||
downRate = int64(binary.BigEndian.Uint64(b.Get([]byte("downRate"))))
|
||||
upCredit = int64(binary.BigEndian.Uint64(b.Get([]byte("upCredit")))) // reee brackets
|
||||
downCredit = int64(binary.BigEndian.Uint64(b.Get([]byte("downCredit"))))
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// TODO: put all of these parameters in a struct instead
|
||||
u := MakeUser(up, UID, sessionsCap, upRate, downRate, upCredit, downCredit)
|
||||
up.activeUsersM.Lock()
|
||||
up.activeUsers[UID] = u
|
||||
up.activeUsersM.Unlock()
|
||||
return u, nil
|
||||
}
|
||||
|
||||
func (up *Userpanel) AddNewUser(UID [32]byte, sessionsCap uint32, upRate, downRate, upCredit, downCredit int64) error {
|
||||
err := up.db.Update(func(tx *bolt.Tx) error {
|
||||
b, err := tx.CreateBucket(UID[:])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// FIXME: obnoxious code
|
||||
quad := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(quad, sessionsCap)
|
||||
if err = b.Put([]byte("sessionsCap"), quad); err != nil {
|
||||
return err
|
||||
}
|
||||
oct := make([]byte, 8)
|
||||
binary.BigEndian.PutUint64(oct, uint64(upRate))
|
||||
if err = b.Put([]byte("upRate"), oct); err != nil {
|
||||
return err
|
||||
}
|
||||
binary.BigEndian.PutUint64(oct, uint64(downRate))
|
||||
if err = b.Put([]byte("downRate"), oct); err != nil {
|
||||
return err
|
||||
}
|
||||
binary.BigEndian.PutUint64(oct, uint64(upCredit))
|
||||
if err = b.Put([]byte("upCredit"), oct); err != nil {
|
||||
return err
|
||||
}
|
||||
binary.BigEndian.PutUint64(oct, uint64(downCredit))
|
||||
if err = b.Put([]byte("downCredit"), oct); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (up *Userpanel) updateDBEntryUint32(UID [32]byte, key string, value uint32) error {
|
||||
err := up.db.Update(func(tx *bolt.Tx) error {
|
||||
b := tx.Bucket(UID[:])
|
||||
if b == nil {
|
||||
return ErrUserNotFound
|
||||
}
|
||||
quad := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(quad, value)
|
||||
if err := b.Put([]byte(key), quad); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (up *Userpanel) updateDBEntryInt64(UID [32]byte, key string, value int64) error {
|
||||
err := up.db.Update(func(tx *bolt.Tx) error {
|
||||
b := tx.Bucket(UID[:])
|
||||
if b == nil {
|
||||
return ErrUserNotFound
|
||||
}
|
||||
oct := make([]byte, 8)
|
||||
binary.BigEndian.PutUint64(oct, uint64(value))
|
||||
if err := b.Put([]byte(key), oct); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
// This is used when all sessions of a user close
|
||||
func (up *Userpanel) delActiveUser(UID [32]byte) {
|
||||
up.activeUsersM.Lock()
|
||||
delete(up.activeUsers, UID)
|
||||
up.activeUsersM.Unlock()
|
||||
}
|
||||
|
||||
func (up *Userpanel) getActiveUser(UID [32]byte) *user {
|
||||
up.activeUsersM.RLock()
|
||||
defer up.activeUsersM.RUnlock()
|
||||
return up.activeUsers[UID]
|
||||
}
|
||||
|
||||
func (up *Userpanel) SetSessionsCap(UID [32]byte, newSessionsCap uint32) error {
|
||||
if u := up.getActiveUser(UID); u != nil {
|
||||
u.setSessionsCap(newSessionsCap)
|
||||
}
|
||||
err := up.updateDBEntryUint32(UID, "sessionsCap", newSessionsCap)
|
||||
return err
|
||||
}
|
||||
Loading…
Reference in New Issue