Code optimization

This commit is contained in:
刘河
2019-03-29 10:41:57 +08:00
parent 4b0aebd6a5
commit cc6d053b6d
32 changed files with 357 additions and 289 deletions

View File

@@ -2,6 +2,7 @@ package file
import (
"github.com/cnlh/nps/lib/common"
"sort"
"sync"
)
@@ -20,3 +21,15 @@ func GetCsvDb() *Csv {
})
return CsvDb
}
func GetMapKeys(m sync.Map, isSort bool, sortKey, order string) (keys []int) {
if sortKey != "" && isSort {
return sortClientByKey(m, sortKey, order)
}
m.Range(func(key, value interface{}) bool {
keys = append(keys, key.(int))
return true
})
sort.Ints(keys)
return
}

View File

@@ -7,7 +7,6 @@ import (
"github.com/cnlh/nps/lib/common"
"github.com/cnlh/nps/lib/crypt"
"github.com/cnlh/nps/lib/rate"
"github.com/cnlh/nps/vender/github.com/astaxie/beego"
"github.com/cnlh/nps/vender/github.com/astaxie/beego/logs"
"net/http"
"os"
@@ -52,7 +51,7 @@ func (s *Csv) StoreTasksToCsv() {
record := []string{
strconv.Itoa(task.Port),
task.Mode,
task.Target,
task.Target.TargetStr,
common.GetStrByBool(task.Status),
strconv.Itoa(task.Id),
strconv.Itoa(task.Client.Id),
@@ -101,12 +100,13 @@ func (s *Csv) LoadTaskFromCsv() {
post := &Tunnel{
Port: common.GetIntNoErrByStr(item[0]),
Mode: item[1],
Target: item[2],
Status: common.GetBoolByStr(item[3]),
Id: common.GetIntNoErrByStr(item[4]),
Remark: item[6],
Password: item[9],
}
post.Target = new(Target)
post.Target.TargetStr = item[2]
post.Flow = new(Flow)
post.Flow.ExportFlow = int64(common.GetIntNoErrByStr(item[7]))
post.Flow.InletFlow = int64(common.GetIntNoErrByStr(item[8]))
@@ -212,7 +212,7 @@ func (s *Csv) StoreHostToCsv() {
}
record := []string{
host.Host,
host.Target,
host.Target.TargetStr,
strconv.Itoa(host.Client.Id),
host.HeaderChange,
host.HostChange,
@@ -274,6 +274,16 @@ func (s *Csv) LoadClientFromCsv() {
} else {
post.ConfigConnAllow = true
}
if len(item) >= 13 {
post.WebUserName = item[12]
} else {
post.WebUserName = ""
}
if len(item) >= 14 {
post.WebPassword = item[13]
} else {
post.WebPassword = ""
}
s.Clients.Store(post.Id, post)
}
}
@@ -289,7 +299,6 @@ func (s *Csv) LoadHostFromCsv() {
for _, item := range records {
post := &Host{
Host: item[0],
Target: item[1],
HeaderChange: item[3],
HostChange: item[4],
Remark: item[5],
@@ -299,6 +308,8 @@ func (s *Csv) LoadHostFromCsv() {
if post.Client, err = s.GetClient(common.GetIntNoErrByStr(item[2])); err != nil {
continue
}
post.Target = new(Target)
post.Target.TargetStr = item[1]
post.Flow = new(Flow)
post.Flow.ExportFlow = int64(common.GetIntNoErrByStr(item[8]))
post.Flow.InletFlow = int64(common.GetIntNoErrByStr(item[9]))
@@ -344,12 +355,12 @@ func (s *Csv) IsHostExist(h *Host) bool {
}
func (s *Csv) NewHost(t *Host) error {
if s.IsHostExist(t) {
return errors.New("host has exist")
}
if t.Location == "" {
t.Location = "/"
}
if s.IsHostExist(t) {
return errors.New("host has exist")
}
t.Flow = new(Flow)
s.Hosts.Store(t.Id, t)
s.StoreHostToCsv()
@@ -359,7 +370,7 @@ func (s *Csv) NewHost(t *Host) error {
func (s *Csv) GetHost(start, length int, id int, search string) ([]*Host, int) {
list := make([]*Host, 0)
var cnt int
keys := common.GetMapKeys(s.Hosts)
keys := GetMapKeys(s.Hosts, false, "", "")
for _, key := range keys {
if value, ok := s.Hosts.Load(key); ok {
v := value.(*Host)
@@ -387,6 +398,9 @@ func (s *Csv) DelClient(id int) error {
func (s *Csv) NewClient(c *Client) error {
var isNotSet bool
if c.WebUserName != "" && !s.VerifyUserName(c.WebUserName, c.Id) {
return errors.New("web login username duplicate, please reset")
}
reset:
if c.VerifyKey == "" || isNotSet {
isNotSet = true
@@ -396,7 +410,7 @@ reset:
c.Rate = rate.NewRate(int64(2 << 23))
c.Rate.Start()
}
if !s.VerifyVkey(c.VerifyKey, c.id) {
if !s.VerifyVkey(c.VerifyKey, c.Id) {
if isNotSet {
goto reset
}
@@ -425,6 +439,18 @@ func (s *Csv) VerifyVkey(vkey string, id int) (res bool) {
})
return res
}
func (s *Csv) VerifyUserName(username string, id int) (res bool) {
res = true
s.Clients.Range(func(key, value interface{}) bool {
v := value.(*Client)
if v.WebUserName == username && v.Id != id {
res = false
return false
}
return true
})
return res
}
func (s *Csv) GetClientId() int32 {
return atomic.AddInt32(&s.ClientIncreaseId, 1)
@@ -447,10 +473,10 @@ func (s *Csv) UpdateClient(t *Client) error {
return nil
}
func (s *Csv) GetClientList(start, length int, search string, clientId int) ([]*Client, int) {
func (s *Csv) GetClientList(start, length int, search, sort, order string, clientId int) ([]*Client, int) {
list := make([]*Client, 0)
var cnt int
keys := common.GetMapKeys(s.Clients)
keys := GetMapKeys(s.Clients, true, sort, order)
for _, key := range keys {
if value, ok := s.Clients.Load(key); ok {
v := value.(*Client)
@@ -477,11 +503,7 @@ func (s *Csv) GetClientList(start, length int, search string, clientId int) ([]*
func (s *Csv) IsPubClient(id int) bool {
client, err := s.GetClient(id)
if err == nil {
if client.VerifyKey == beego.AppConfig.String("public_vkey") {
return true
} else {
return false
}
return client.NoDisplay
}
return false
}
@@ -589,6 +611,8 @@ func (s *Csv) StoreClientsToCsv() {
strconv.Itoa(int(client.Flow.FlowLimit)),
strconv.Itoa(int(client.MaxConn)),
common.GetStrByBool(client.ConfigConnAllow),
client.WebUserName,
client.WebPassword,
}
err := writer.Write(record)
if err != nil {

View File

@@ -5,13 +5,14 @@ import (
"github.com/pkg/errors"
"strings"
"sync"
"sync/atomic"
"time"
)
type Flow struct {
ExportFlow int64 //出口流量
InletFlow int64 //入口流量
FlowLimit int64 //流量限制,出口+入口 /M
ExportFlow int64
InletFlow int64
FlowLimit int64
sync.RWMutex
}
@@ -25,20 +26,21 @@ func (s *Flow) Add(in, out int64) {
type Client struct {
Cnf *Config
Id int //id
VerifyKey string //验证密钥
Addr string //客户端ip地址
Remark string //备注
Status bool //是否开启
IsConnect bool //是否连接
RateLimit int //速度限制 /kb
Flow *Flow //流量
Rate *rate.Rate //速度控制
NoStore bool
NoDisplay bool
MaxConn int //客户端最大连接数
NowConn int //当前连接数
id int
ConfigConnAllow bool
VerifyKey string //verify key
Addr string //the ip of client
Remark string //remark
Status bool //is allow connect
IsConnect bool //is the client connect
RateLimit int //rate /kb
Flow *Flow //flow setting
Rate *rate.Rate //rate limit
NoStore bool //no store to file
NoDisplay bool //no display on web
MaxConn int //the max connection num of client allow
NowConn int32 //the connection num of now
WebUserName string //the username of web login
WebPassword string //the password of web login
ConfigConnAllow bool //is allow connected by config file
sync.RWMutex
}
@@ -61,30 +63,21 @@ func NewClient(vKey string, noStore bool, noDisplay bool) *Client {
}
func (s *Client) CutConn() {
s.Lock()
defer s.Unlock()
s.NowConn++
atomic.AddInt32(&s.NowConn, 1)
}
func (s *Client) AddConn() {
s.Lock()
defer s.Unlock()
s.NowConn--
atomic.AddInt32(&s.NowConn, -1)
}
func (s *Client) GetConn() bool {
if s.MaxConn == 0 || s.NowConn < s.MaxConn {
if s.MaxConn == 0 || int(s.NowConn) < s.MaxConn {
s.CutConn()
return true
}
return false
}
//modify the hosts and the tunnels by health information
func (s *Client) ModifyTarget() {
}
func (s *Client) HasTunnel(t *Tunnel) (exist bool) {
GetCsvDb().Tasks.Range(func(key, value interface{}) bool {
v := value.(*Tunnel)
@@ -114,13 +107,11 @@ type Tunnel struct {
Id int //Id
Port int //服务端监听端口
ServerIp string
Mode string //启动方式
Target string //目标
TargetArr []string //目标
Status bool //设置是否开启
RunStatus bool //当前运行状态
Client *Client //所属客户端id
Ports string //客户端与服务端传递
Mode string //启动方式
Status bool //设置是否开启
RunStatus bool //当前运行状态
Client *Client //所属客户端id
Ports string //客户端与服务端传递
Flow *Flow
Password string //私密模式密码,唯一
Remark string //备注
@@ -128,7 +119,7 @@ type Tunnel struct {
NoStore bool
LocalPath string
StripPre string
NowIndex int
Target *Target
Health
sync.RWMutex
}
@@ -146,9 +137,16 @@ type Health struct {
sync.RWMutex
}
func (s *Tunnel) GetRandomTarget() (string, error) {
type Target struct {
nowIndex int
TargetStr string
TargetArr []string //目标
sync.RWMutex
}
func (s *Target) GetRandomTarget() (string, error) {
if s.TargetArr == nil {
s.TargetArr = strings.Split(s.Target, "\n")
s.TargetArr = strings.Split(s.TargetStr, "\n")
}
if len(s.TargetArr) == 1 {
return s.TargetArr[0], nil
@@ -158,54 +156,33 @@ func (s *Tunnel) GetRandomTarget() (string, error) {
}
s.Lock()
defer s.Unlock()
if s.NowIndex >= len(s.TargetArr)-1 {
s.NowIndex = -1
if s.nowIndex >= len(s.TargetArr)-1 {
s.nowIndex = -1
}
s.NowIndex++
return s.TargetArr[s.NowIndex], nil
s.nowIndex++
return s.TargetArr[s.nowIndex], nil
}
type Config struct {
U string //socks5验证用户名
P string //socks5验证密码
Compress bool //压缩方式
Crypt bool //是否加密
U string
P string
Compress bool
Crypt bool
}
type Host struct {
Id int
Host string //启动方式
Target string //目标
HeaderChange string //host修改
HostChange string //host修改
Location string //url 路由
Host string //host
HeaderChange string //header change
HostChange string //host change
Location string //url router
Remark string //remark
Scheme string //http https all
NoStore bool
IsClose bool
Flow *Flow
Client *Client
Remark string //备注
NowIndex int
TargetArr []string
NoStore bool
Scheme string //http https all
IsClose bool
Target *Target //目标
Health
sync.RWMutex
}
func (s *Host) GetRandomTarget() (string, error) {
if s.TargetArr == nil {
s.TargetArr = strings.Split(s.Target, "\n")
}
if len(s.TargetArr) == 1 {
return s.TargetArr[0], nil
}
if len(s.TargetArr) == 0 {
return "", errors.New("all inward-bending targets are offline")
}
s.Lock()
defer s.Unlock()
if s.NowIndex >= len(s.TargetArr)-1 {
s.NowIndex = -1
}
s.NowIndex++
return s.TargetArr[s.NowIndex], nil
}

41
lib/file/sort.go Normal file
View File

@@ -0,0 +1,41 @@
package file
import (
"reflect"
"sort"
"sync"
)
// A data structure to hold a key/value pair.
type Pair struct {
key string //sort key
cId int
order string
clientFlow *Flow
}
// A slice of Pairs that implements sort.Interface to sort by Value.
type PairList []*Pair
func (p PairList) Swap(i, j int) { p[i], p[j] = p[j], p[i] }
func (p PairList) Len() int { return len(p) }
func (p PairList) Less(i, j int) bool {
if p[i].order == "desc" {
return reflect.ValueOf(*p[i].clientFlow).FieldByName(p[i].key).Int() < reflect.ValueOf(*p[j].clientFlow).FieldByName(p[j].key).Int()
}
return reflect.ValueOf(*p[i].clientFlow).FieldByName(p[i].key).Int() > reflect.ValueOf(*p[j].clientFlow).FieldByName(p[j].key).Int()
}
// A function to turn a map into a PairList, then sort and return it.
func sortClientByKey(m sync.Map, sortKey, order string) (res []int) {
p := make(PairList, 0)
m.Range(func(key, value interface{}) bool {
p = append(p, &Pair{sortKey, value.(*Client).Id, order, value.(*Client).Flow})
return true
})
sort.Sort(p)
for _, v := range p {
res = append(res, v.cId)
}
return
}