Refactor usermanager

This commit is contained in:
Qian Wang 2019-08-03 11:17:09 +01:00
parent 67521efbd6
commit 08cb5d8462
7 changed files with 48 additions and 44 deletions

View File

@ -36,7 +36,7 @@ func (u *ActiveUser) GetSession(sessionID uint32, obfuscator *mux.Obfuscator, un
if sesh = u.sessions[sessionID]; sesh != nil { if sesh = u.sessions[sessionID]; sesh != nil {
return sesh, true, nil return sesh, true, nil
} else { } else {
err := u.panel.Manager.authoriseNewSession(u) err := u.panel.Manager.AuthoriseNewSession(u.arrUID[:], len(u.sessions))
if err != nil { if err != nil {
return nil, false, err return nil, false, err
} }

View File

@ -5,6 +5,7 @@ import (
"encoding/base64" "encoding/base64"
"encoding/json" "encoding/json"
"errors" "errors"
"github.com/cbeuw/Cloak/internal/server/usermanager"
"io/ioutil" "io/ioutil"
"sync" "sync"
"time" "time"
@ -74,7 +75,7 @@ func (sta *State) ParseConfig(conf string) (err error) {
//TODO: implement command & control mode //TODO: implement command & control mode
} else { } else {
manager, err := MakeLocalManager(preParse.DatabasePath) manager, err := usermanager.MakeLocalManager(preParse.DatabasePath)
if err != nil { if err != nil {
return err return err
} }

View File

@ -1,4 +1,4 @@
package server package usermanager
import ( import (
"encoding/binary" "encoding/binary"
@ -63,7 +63,7 @@ func (manager *localManager) registerMux() *gmux.Router {
return r return r
} }
func (manager *localManager) authenticateUser(UID []byte) (int64, int64, error) { func (manager *localManager) AuthenticateUser(UID []byte) (int64, int64, error) {
var upRate, downRate, upCredit, downCredit, expiryTime int64 var upRate, downRate, upCredit, downCredit, expiryTime int64
err := manager.db.View(func(tx *bolt.Tx) error { err := manager.db.View(func(tx *bolt.Tx) error {
bucket := tx.Bucket(UID) bucket := tx.Bucket(UID)
@ -93,11 +93,13 @@ func (manager *localManager) authenticateUser(UID []byte) (int64, int64, error)
return upRate, downRate, nil return upRate, downRate, nil
} }
func (manager *localManager) authoriseNewSession(user *ActiveUser) error { func (manager *localManager) AuthoriseNewSession(UID []byte, numExistingSessions int) error {
var arrUID [16]byte
copy(arrUID[:], UID)
var sessionsCap int var sessionsCap int
var upCredit, downCredit, expiryTime int64 var upCredit, downCredit, expiryTime int64
err := manager.db.View(func(tx *bolt.Tx) error { err := manager.db.View(func(tx *bolt.Tx) error {
bucket := tx.Bucket(user.arrUID[:]) bucket := tx.Bucket(arrUID[:])
if bucket == nil { if bucket == nil {
return ErrUserNotFound return ErrUserNotFound
} }
@ -120,7 +122,7 @@ func (manager *localManager) authoriseNewSession(user *ActiveUser) error {
return ErrUserExpired return ErrUserExpired
} }
//user.sessionsM.RLock() //user.sessionsM.RLock()
if len(user.sessions) >= sessionsCap { if numExistingSessions >= sessionsCap {
//user.sessionsM.RUnlock() //user.sessionsM.RUnlock()
return ErrSessionsCapReached return ErrSessionsCapReached
} }
@ -128,14 +130,14 @@ func (manager *localManager) authoriseNewSession(user *ActiveUser) error {
return nil return nil
} }
func (manager *localManager) uploadStatus(uploads []statusUpdate) ([]statusResponse, error) { func (manager *localManager) UploadStatus(uploads []StatusUpdate) ([]StatusResponse, error) {
var responses []statusResponse var responses []StatusResponse
err := manager.db.Update(func(tx *bolt.Tx) error { err := manager.db.Update(func(tx *bolt.Tx) error {
for _, status := range uploads { for _, status := range uploads {
var resp statusResponse var resp StatusResponse
bucket := tx.Bucket(status.UID) bucket := tx.Bucket(status.UID)
if bucket == nil { if bucket == nil {
resp = statusResponse{ resp = StatusResponse{
status.UID, status.UID,
TERMINATE, TERMINATE,
"User no longer exists", "User no longer exists",
@ -145,9 +147,9 @@ func (manager *localManager) uploadStatus(uploads []statusUpdate) ([]statusRespo
} }
oldUp := int64(Uint64(bucket.Get([]byte("UpCredit")))) oldUp := int64(Uint64(bucket.Get([]byte("UpCredit"))))
newUp := oldUp - status.upUsage newUp := oldUp - status.UpUsage
if newUp <= 0 { if newUp <= 0 {
resp = statusResponse{ resp = StatusResponse{
status.UID, status.UID,
TERMINATE, TERMINATE,
"No upload credit left", "No upload credit left",
@ -162,9 +164,9 @@ func (manager *localManager) uploadStatus(uploads []statusUpdate) ([]statusRespo
} }
oldDown := int64(Uint64(bucket.Get([]byte("DownCredit")))) oldDown := int64(Uint64(bucket.Get([]byte("DownCredit"))))
newDown := oldDown - status.downUsage newDown := oldDown - status.DownUsage
if newDown <= 0 { if newDown <= 0 {
resp = statusResponse{ resp = StatusResponse{
status.UID, status.UID,
TERMINATE, TERMINATE,
"No download credit left", "No download credit left",
@ -180,7 +182,7 @@ func (manager *localManager) uploadStatus(uploads []statusUpdate) ([]statusRespo
expiry := int64(Uint64(bucket.Get([]byte("ExpiryTime")))) expiry := int64(Uint64(bucket.Get([]byte("ExpiryTime"))))
if time.Now().Unix() > expiry { if time.Now().Unix() > expiry {
resp = statusResponse{ resp = StatusResponse{
status.UID, status.UID,
TERMINATE, TERMINATE,
"User has expired", "User has expired",

View File

@ -1,4 +1,4 @@
package server package usermanager
import ( import (
"bytes" "bytes"

View File

@ -1,23 +1,23 @@
package server package usermanager
import ( import (
"errors" "errors"
) )
type statusUpdate struct { type StatusUpdate struct {
UID []byte UID []byte
active bool Active bool
numSession int NumSession int
upUsage int64 UpUsage int64
downUsage int64 DownUsage int64
timestamp int64 Timestamp int64
} }
type statusResponse struct { type StatusResponse struct {
UID []byte UID []byte
action int Action int
message string Message string
} }
const ( const (
@ -32,7 +32,7 @@ var ErrNoDownCredit = errors.New("No download credit left")
var ErrUserExpired = errors.New("User has expired") var ErrUserExpired = errors.New("User has expired")
type UserManager interface { type UserManager interface {
authenticateUser([]byte) (int64, int64, error) AuthenticateUser([]byte) (int64, int64, error)
authoriseNewSession(*ActiveUser) error AuthoriseNewSession([]byte, int) error
uploadStatus([]statusUpdate) ([]statusResponse, error) UploadStatus([]StatusUpdate) ([]StatusResponse, error)
} }

View File

@ -1,6 +1,7 @@
package server package server
import ( import (
"github.com/cbeuw/Cloak/internal/server/usermanager"
"sync" "sync"
"sync/atomic" "sync/atomic"
"time" "time"
@ -10,7 +11,7 @@ import (
) )
type userPanel struct { type userPanel struct {
Manager UserManager Manager usermanager.UserManager
activeUsersM sync.RWMutex activeUsersM sync.RWMutex
activeUsers map[[16]byte]*ActiveUser activeUsers map[[16]byte]*ActiveUser
@ -18,7 +19,7 @@ type userPanel struct {
usageUpdateQueue map[[16]byte]*usagePair usageUpdateQueue map[[16]byte]*usagePair
} }
func MakeUserPanel(manager UserManager) *userPanel { func MakeUserPanel(manager usermanager.UserManager) *userPanel {
ret := &userPanel{ ret := &userPanel{
Manager: manager, Manager: manager,
activeUsers: make(map[[16]byte]*ActiveUser), activeUsers: make(map[[16]byte]*ActiveUser),
@ -37,7 +38,7 @@ func (panel *userPanel) GetUser(UID []byte) (*ActiveUser, error) {
return user, nil return user, nil
} }
upRate, downRate, err := panel.Manager.authenticateUser(UID) upRate, downRate, err := panel.Manager.AuthenticateUser(UID)
if err != nil { if err != nil {
panel.activeUsersM.Unlock() panel.activeUsersM.Unlock()
return nil, err return nil, err
@ -103,7 +104,7 @@ func (panel *userPanel) updateUsageQueueForOne(user *ActiveUser) {
func (panel *userPanel) commitUpdate() error { func (panel *userPanel) commitUpdate() error {
panel.usageUpdateQueueM.Lock() panel.usageUpdateQueueM.Lock()
statuses := make([]statusUpdate, 0, len(panel.usageUpdateQueue)) statuses := make([]usermanager.StatusUpdate, 0, len(panel.usageUpdateQueue))
for arrUID, usage := range panel.usageUpdateQueue { for arrUID, usage := range panel.usageUpdateQueue {
panel.activeUsersM.RLock() panel.activeUsersM.RLock()
user := panel.activeUsers[arrUID] user := panel.activeUsers[arrUID]
@ -112,31 +113,31 @@ func (panel *userPanel) commitUpdate() error {
if user != nil { if user != nil {
numSession = user.NumSession() numSession = user.NumSession()
} }
status := statusUpdate{ status := usermanager.StatusUpdate{
UID: arrUID[:], UID: arrUID[:],
active: panel.isActive(arrUID[:]), Active: panel.isActive(arrUID[:]),
numSession: numSession, NumSession: numSession,
upUsage: *usage.up, UpUsage: *usage.up,
downUsage: *usage.down, DownUsage: *usage.down,
timestamp: time.Now().Unix(), Timestamp: time.Now().Unix(),
} }
statuses = append(statuses, status) statuses = append(statuses, status)
} }
responses, err := panel.Manager.uploadStatus(statuses) responses, err := panel.Manager.UploadStatus(statuses)
if err != nil { if err != nil {
return err return err
} }
for _, resp := range responses { for _, resp := range responses {
var arrUID [16]byte var arrUID [16]byte
copy(arrUID[:], resp.UID) copy(arrUID[:], resp.UID)
switch resp.action { switch resp.Action {
case TERMINATE: case usermanager.TERMINATE:
panel.activeUsersM.RLock() panel.activeUsersM.RLock()
user := panel.activeUsers[arrUID] user := panel.activeUsers[arrUID]
panel.activeUsersM.RUnlock() panel.activeUsersM.RUnlock()
if user != nil { if user != nil {
user.Terminate(resp.message) user.Terminate(resp.Message)
} }
} }
} }