quickshare/src/userstore/user_store.go

package userstore

import (
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/ihexxa/quickshare/src/kvstore"
)

const (
	AdminRole   = "admin"
	UserRole    = "user"
	VisitorRole = "visitor"
	InitNs      = "usersInit"
	IDsNs       = "ids"
	UsersNs     = "users"
	RoleListNs  = "roleList"
	InitTimeKey = "initTime"

	defaultSpaceLimit         = 1024 * 1024 * 1024 // 1GB
	defaultUploadSpeedLimit   = 50 * 1024 * 1024   // 50MB
	defaultDownloadSpeedLimit = 50 * 1024 * 1024   // 50MB
)

var (
	ErrReachedLimit = errors.New("reached space limit")
)

func IsReachedLimitErr(err error) bool {
	return err == ErrReachedLimit
}

type Quota struct {
	SpaceLimit         int64 `json:"spaceLimit,string"`
	UploadSpeedLimit   int   `json:"uploadSpeedLimit"`
	DownloadSpeedLimit int   `json:"downloadSpeedLimit"`
}

type User struct {
	ID        uint64 `json:"id,string"`
	Name      string `json:"name"`
	Pwd       string `json:"pwd"`
	Role      string `json:"role"`
	UsedSpace int64  `json:"usedSpace,string"`
	Quota     *Quota `json:"quota"`
}

type IUserStore interface {
	Init(rootName, rootPwd string) error
	IsInited() bool
	AddUser(user *User) error
	DelUser(id uint64) error
	GetUser(id uint64) (*User, error)
	GetUserByName(name string) (*User, error)
	SetInfo(id uint64, user *User) error
	SetUsed(id uint64, incr bool, capacity int64) error
	SetPwd(id uint64, pwd string) error
	ListUsers() ([]*User, error)
	AddRole(role string) error
	DelRole(role string) error
	ListRoles() (map[string]bool, error)
}

type KVUserStore struct {
	store kvstore.IKVStore
	mtx   *sync.RWMutex
}

func NewKVUserStore(store kvstore.IKVStore) (*KVUserStore, error) {
	_, ok := store.GetStringIn(InitNs, InitTimeKey)
	if !ok {
		var err error
		for _, nsName := range []string{
			IDsNs,
			UsersNs,
			InitNs,
			RoleListNs,
		} {
			if err = store.AddNamespace(nsName); err != nil {
				return nil, err
			}
		}
	}

	return &KVUserStore{
		store: store,
		mtx:   &sync.RWMutex{},
	}, nil
}

func (us *KVUserStore) Init(rootName, rootPwd string) error {
	var err error
	err = us.AddUser(&User{
		ID:   0,
		Name: rootName,
		Pwd:  rootPwd,
		Role: AdminRole,
		Quota: &Quota{
			SpaceLimit:         defaultSpaceLimit,
			UploadSpeedLimit:   defaultUploadSpeedLimit,
			DownloadSpeedLimit: defaultDownloadSpeedLimit,
		},
	})
	if err != nil {
		return err
	}

	for _, role := range []string{AdminRole, UserRole, VisitorRole} {
		err = us.AddRole(role)
		if err != nil {
			return err
		}
	}

	return us.store.SetStringIn(InitNs, InitTimeKey, fmt.Sprintf("%d", time.Now().Unix()))
}

func (us *KVUserStore) IsInited() bool {
	_, ok := us.store.GetStringIn(InitNs, InitTimeKey)
	return ok
}

func (us *KVUserStore) AddUser(user *User) error {
	us.mtx.Lock()
	defer us.mtx.Unlock()

	userID := fmt.Sprint(user.ID)
	_, ok := us.store.GetStringIn(UsersNs, userID)
	if ok {
		return fmt.Errorf("userID (%d) exists", user.ID)
	}
	if user.Name == "" || user.Pwd == "" {
		return errors.New("user name or password can not be empty")
	}
	_, ok = us.store.GetStringIn(IDsNs, user.Name)
	if ok {
		return fmt.Errorf("user name (%s) exists", user.Name)
	}

	err := us.store.SetStringIn(IDsNs, user.Name, userID)
	if err != nil {
		return err
	}
	userBytes, err := json.Marshal(user)
	if err != nil {
		return err
	}
	return us.store.SetStringIn(UsersNs, userID, string(userBytes))
}

func (us *KVUserStore) DelUser(id uint64) error {
	us.mtx.Lock()
	defer us.mtx.Unlock()

	userID := fmt.Sprint(id)
	name, ok := us.store.GetStringIn(UsersNs, userID)
	if !ok {
		return fmt.Errorf("userID (%s) does not exist", userID)
	}

	// TODO: add complement operations if part of the actions fails
	err1 := us.store.DelStringIn(IDsNs, name)
	err2 := us.store.DelStringIn(UsersNs, userID)
	if err1 != nil || err2 != nil {
		return fmt.Errorf("get id(%s) user(%s)", err1, err2)
	}
	return nil
}

func (us *KVUserStore) GetUser(id uint64) (*User, error) {
	us.mtx.RLock()
	defer us.mtx.RUnlock()

	userID := fmt.Sprint(id)

	infoStr, ok := us.store.GetStringIn(UsersNs, userID)
	if !ok {
		return nil, fmt.Errorf("user (%s) not found", userID)
	}
	user := &User{}
	err := json.Unmarshal([]byte(infoStr), user)
	if err != nil {
		return nil, err
	}

	gotID, ok := us.store.GetStringIn(IDsNs, user.Name)
	if !ok {
		return nil, fmt.Errorf("user id (%s) not found", user.Name)
	} else if gotID != userID {
		return nil, fmt.Errorf("user id (%s) not match: got(%s) expected(%s)", user.Name, gotID, userID)
	}

	// TODO: use sync.Pool instead
	return user, nil

}

func (us *KVUserStore) GetUserByName(name string) (*User, error) {
	us.mtx.RLock()
	defer us.mtx.RUnlock()

	userID, ok := us.store.GetStringIn(IDsNs, name)
	if !ok {
		return nil, fmt.Errorf("user id (%s) not found", name)
	}
	infoStr, ok := us.store.GetStringIn(UsersNs, userID)
	if !ok {
		return nil, fmt.Errorf("user name (%s) not found", userID)
	}

	user := &User{}
	err := json.Unmarshal([]byte(infoStr), user)
	if err != nil {
		return nil, err
	}
	if user.Name != name {
		return nil, fmt.Errorf("user id (%s) not match: got(%s) expected(%s)", userID, user.Name, name)
	}

	// TODO: use sync.Pool instead
	return user, nil

}

func (us *KVUserStore) SetPwd(id uint64, pwd string) error {
	us.mtx.Lock()
	defer us.mtx.Unlock()

	userID := fmt.Sprint(id)
	infoStr, ok := us.store.GetStringIn(UsersNs, userID)
	if !ok {
		return fmt.Errorf("user (%d) does not exist", id)
	}
	gotUser := &User{}
	err := json.Unmarshal([]byte(infoStr), gotUser)
	if err != nil {
		return err
	} else if gotUser.ID != id {
		return fmt.Errorf("user id key(%d) info(%d) does match", id, gotUser.ID)
	}

	gotUser.Pwd = pwd
	infoBytes, err := json.Marshal(gotUser)
	if err != nil {
		return err
	}
	return us.store.SetStringIn(UsersNs, userID, string(infoBytes))
}

func (us *KVUserStore) SetUsed(id uint64, incr bool, capacity int64) error {
	us.mtx.Lock()
	defer us.mtx.Unlock()

	userID := fmt.Sprint(id)
	infoStr, ok := us.store.GetStringIn(UsersNs, userID)
	if !ok {
		return fmt.Errorf("user (%d) does not exist", id)
	}

	gotUser := &User{}
	err := json.Unmarshal([]byte(infoStr), gotUser)
	if err != nil {
		return err
	} else if gotUser.ID != id {
		return fmt.Errorf("user id key(%d) info(%d) does match", id, gotUser.ID)
	}

	if incr && gotUser.UsedSpace+capacity > int64(gotUser.Quota.SpaceLimit) {
		return ErrReachedLimit
	}

	if incr {
		gotUser.UsedSpace = gotUser.UsedSpace + capacity
	} else {
		gotUser.UsedSpace = gotUser.UsedSpace - capacity
	}
	infoBytes, err := json.Marshal(gotUser)
	if err != nil {
		return err
	}
	return us.store.SetStringIn(UsersNs, userID, string(infoBytes))
}

func (us *KVUserStore) SetInfo(id uint64, user *User) error {
	us.mtx.Lock()
	defer us.mtx.Unlock()

	userID := fmt.Sprint(id)
	infoStr, ok := us.store.GetStringIn(UsersNs, userID)
	if !ok {
		return fmt.Errorf("user (%d) does not exist", id)
	}
	gotUser := &User{}
	err := json.Unmarshal([]byte(infoStr), gotUser)
	if err != nil {
		return err
	} else if gotUser.ID != id {
		return fmt.Errorf("user id key(%d) info(%d) does match", id, gotUser.ID)
	}

	// name and password can not be updated here
	if user.Role != "" {
		gotUser.Role = user.Role
	}
	if user.Quota != nil {
		gotUser.Quota = user.Quota
	}
	if user.UsedSpace > 0 {
		gotUser.UsedSpace = user.UsedSpace
	}

	infoBytes, err := json.Marshal(gotUser)
	if err != nil {
		return err
	}

	return us.store.SetStringIn(UsersNs, userID, string(infoBytes))
}

func (us *KVUserStore) ListUsers() ([]*User, error) {
	us.mtx.RLock()
	defer us.mtx.RUnlock()

	idToInfo, err := us.store.ListStringsIn(UsersNs)
	if err != nil {
		return nil, err
	}

	users := []*User{}
	for _, infoStr := range idToInfo {
		user := &User{}
		err = json.Unmarshal([]byte(infoStr), user)
		if err != nil {
			return nil, err
		}
		user.Pwd = ""

		users = append(users, user)
	}

	return users, nil
}

func (us *KVUserStore) AddRole(role string) error {
	us.mtx.Lock()
	defer us.mtx.Unlock()

	_, ok := us.store.GetBoolIn(RoleListNs, role)
	if ok {
		return fmt.Errorf("role (%s) exists", role)
	}

	return us.store.SetBoolIn(RoleListNs, role, true)
}

func (us *KVUserStore) DelRole(role string) error {
	us.mtx.Lock()
	defer us.mtx.Unlock()

	if role == AdminRole || role == UserRole || role == VisitorRole {
		return errors.New("predefined roles can not be deleted")
	}

	return us.store.DelBoolIn(RoleListNs, role)
}

func (us *KVUserStore) ListRoles() (map[string]bool, error) {
	us.mtx.Lock()
	defer us.mtx.Unlock()

	return us.store.ListBoolsIn(RoleListNs)
}