yeyinlaoying/rakshasa
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

优化连接逻辑，更新readme

Browse Source
This commit is contained in:
Mob2003
2023-04-22 19:46:52 +08:00
parent 5a67527907
commit a8dcc3871d
21 changed files with 1310 additions and 1443 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rakshasa_lite/server/node.go
+181 -170
View File
@@ -7,13 +7,16 @@ import (
"encoding/json"
"errors"
"fmt"
"github.com/creack/pty"
"io"
"io/ioutil"
"math/rand"
"net"
"os"
"os/exec"
"rakshasa_lite/common"
"runtime"
"runtime/debug"
"strconv"
"strings"
"sync"
@@ -21,11 +24,29 @@ import (
"time"
"unsafe"
)
var (
shellMapLock sync.Mutex
)
type StartCmdParam struct {
Param string
Size *pty.Winsize
}
type remoteCmd struct {
cmdStatus int32
cmd *exec.Cmd
id uint32
stdin io.WriteCloser
inChan chan []byte
translate func(in []byte) ([]byte, error)
ping, pong int64
}
var (
currentNode = &node{}
clientLock = &lock{}
nodeMap = make(map[string]*node)
clientLock1 = &lock{}
nodeMap sync.Map
upLevelNode []*node //上游节点
upNodeWrite = make(chan []byte, 999)
extNodeIp []string
@@ -64,14 +85,14 @@ func InitCurrentNode() {
addr: currentNode.addr,
}
currentNode.mirrorNode.mirrorNode = currentNode
nodeMap[currentNode.uuid] = currentNode
nodeMap.Store(currentNode.uuid, currentNode)
//fmt.Println("当前节点UUID", currentNode.uuid)
go func() {
for b := range upNodeWrite {
for {
ok := func() bool {
l := clientLock.Lock()
l := clientLock1.Lock()
defer l.Unlock()
if len(upLevelNode) == 0 {
@@ -91,6 +112,40 @@ func InitCurrentNode() {
nodeTickPing()
time.AfterFunc(time.Second*10, checkUpLevelNode)
}
func getNode(arg string) (n *node, err error) {
id, err := strconv.Atoi(arg)
if err == nil {
nodeMap.Range(func(key, value interface{}) bool {
_n := value.(*node)
if _n.id == id {
n = _n
return false
}
return true
})
} else {
nodeMap.Range(func(key, value interface{}) bool {
node := value.(*node)
if fmt.Sprintf("%s:%d", node.mainIp, node.port) == arg {
n = node
return false
} else if fmt.Sprintf("%s:%d", node.addr, node.port) == arg {
n = node
return false
} else if node.uuid == arg {
n = node
return false
}
return true
})
}
if n != nil {
return n, nil
} else {
return connectNew(arg)
}
}
func checkUpLevelNode() {
if len(currentConfig.DstNode) > 0 && len(upLevelNode) == 0 {
@@ -108,29 +163,18 @@ func checkUpLevelNode() {
return
}
}
func() {
l := clientLock.RLock()
defer l.RUnlock()
for _, n := range nodeMap {
if n.uuid != currentNode.uuid {
func() {
l.RUnlock()
defer clientLock.RLock(l)
if len(n.mainIp) == 0 {
getNode(fmt.Sprintf("%s:%d", n.addr, n.port))
}
}()
if len(upLevelNode) > 0 {
return
}
nodeMap.Range(func(key, value interface{}) bool {
n := value.(*node)
if n.uuid != currentNode.uuid {
if len(n.mainIp) == 0 {
getNode(fmt.Sprintf("%s:%d", n.addr, n.port))
}
if len(upLevelNode) > 0 {
return false
}
}
}()
return true
})
}
}
@@ -138,12 +182,9 @@ func checkUpLevelNode() {
time.AfterFunc(time.Second*5, checkUpLevelNode)
}
func nodeTickPing() {
l := clientLock.RLock()
defer l.RUnlock()
now := time.Now().Unix()
for _, n := range nodeMap {
nodeMap.Range(func(key, value interface{}) bool {
n := value.(*node)
if n.uuid != currentNode.uuid {
if n.mainIp != "" {
addr1 := fmt.Sprintf("%s:%d", n.mainIp, n.port)
@@ -167,8 +208,8 @@ func nodeTickPing() {
}
}
}
return true
})
time.AfterFunc(time.Second*1, nodeTickPing)
}
@@ -278,29 +319,28 @@ func connectNew(addr string) (n *node, e error) {
}
n.mainIp = cert.RSADecrypterStr(regmsg.MainIp)
if n.port, err = strconv.Atoi(cert.RSADecrypterStr(regmsg.Port)); err != nil {
if n.port, err = strconv.Atoi(cert.RSADecrypterStr(regmsg.Port)); n.port==0 {
n.port = -1
}
if v, ok := nodeMap[regmsg.UUID]; ok {
if v.conn.node != nil && v.conn.node.uuid == regmsg.UUID && v.conn.closeTag == 0 {
if v, ok := nodeMap.Load(regmsg.UUID); ok {
if v.(*node).conn.node != nil && v.(*node).conn.node.uuid == regmsg.UUID && v.(*node).conn.closeTag == 0 {
n.uuid = "" //清空uuid避免正常的node被删
n.conn.Close("重复注册") //当前的连接关掉
n.conn = v.conn
v.mainIp = cert.RSADecrypterStr(regmsg.MainIp)
if v.port, err = strconv.Atoi(cert.RSADecrypterStr(regmsg.Port)); err != nil {
v.port = -1
v.(*node).mainIp = cert.RSADecrypterStr(regmsg.MainIp)
if v.(*node).port, err = strconv.Atoi(cert.RSADecrypterStr(regmsg.Port)); v.(*node).port==0 {
v.(*node).port = -1
}
n = v
n = v.(*node)
} else {
n.conn.node = n
}
} else {
n.conn.node = n
}
nodeMap[n.uuid] = n
nodeMap.Store(n.uuid, n)
n.reConnectAddrs = []string{addr}
n.Write(common.CMD_GET_NODE, 0, nil)
return n, nil
case <-time.After(time.Second * 10):
return nil, errors.New("time out")
@@ -340,7 +380,6 @@ func (n *node) WriteMsg(msg *common.Msg) {
func (n *node) do(msg *common.Msg) {
var err error
//fmt.Println(common.CmdToName[msg.CmdOpteion])
switch msg.CmdOpteion {
case common.CMD_CONNECT_BYIDADDR:
msg.CmdData = cert.RSADecrypterByPubByte(msg.CmdData)
@@ -436,57 +475,50 @@ func (n *node) do(msg *common.Msg) {
}
case common.CMD_REG:
func() {
l := clientLock.Lock()
defer l.Unlock()
var regmsg RegMsg
err = json.Unmarshal(msg.CmdData, &regmsg)
if err != nil {
regmsg.Err = err.Error()
b, _ := json.Marshal(regmsg)
n.Write(common.CMD_REG_RESULT, 0, b)
return
}
uuid := regmsg.UUID
if uuid == currentNode.uuid {
regmsg.Err = "不能连接自己"
b, _ := json.Marshal(regmsg)
n.Write(common.CMD_REG_RESULT, 0, b)
return
}
n.hostName = cert.RSADecrypterStr(regmsg.Hostname)
n.mainIp = cert.RSADecrypterStr(regmsg.MainIp)
if n.port, err = strconv.Atoi(cert.RSADecrypterStr(regmsg.Port)); err != nil {
n.port = -1
}
n.goos = cert.RSADecrypterStr(regmsg.Goos)
n.addr = n.conn.nodeConn.RemoteAddr().String()
if i := strings.Index(n.addr, ":"); i > -1 {
n.addr = n.addr[:i]
}
resultMsg := regmsg
resultMsg.UUID = currentNode.uuid
resultMsg.Hostname = cert.RSAEncrypterStr(currentNode.hostName)
resultMsg.MainIp = cert.RSAEncrypterStr(currentNode.mainIp)
resultMsg.Port = cert.RSAEncrypterStr(strconv.Itoa(currentNode.port))
resultMsg.Goos = cert.RSAEncrypterStr(currentNode.goos)
b, _ := json.Marshal(resultMsg)
//返回成功结果
var regmsg RegMsg
err = json.Unmarshal(msg.CmdData, &regmsg)
if err != nil {
regmsg.Err = err.Error()
b, _ := json.Marshal(regmsg)
n.Write(common.CMD_REG_RESULT, 0, b)
//储存节点
n.uuid = uuid
if v, ok := nodeMap[uuid]; !ok || v.conn.closeTag > 0 {
n.conn.node = n
nodeMap[regmsg.UUID] = n
return
}
uuid := regmsg.UUID
if uuid == currentNode.uuid {
regmsg.Err = "请求的UUID相同，无法连接自己，请将节点设置为不同的UUID"
b, _ := json.Marshal(regmsg)
n.Write(common.CMD_REG_RESULT, 0, b)
return
}
}
currentNode.broadcastNode()
//把本机所有节点同步到注册机器
go n.writeGetNodeResult(msg.CmdId)
}()
n.hostName = cert.RSADecrypterStr(regmsg.Hostname)
n.mainIp = cert.RSADecrypterStr(regmsg.MainIp)
if n.port, err = strconv.Atoi(cert.RSADecrypterStr(regmsg.Port)); err != nil {
n.port = -1
}
n.goos = cert.RSADecrypterStr(regmsg.Goos)
n.addr = n.conn.nodeConn.RemoteAddr().String()
if i := strings.Index(n.addr, ":"); i > -1 {
n.addr = n.addr[:i]
}
resultMsg := regmsg
resultMsg.UUID = currentNode.uuid
resultMsg.Hostname = cert.RSAEncrypterStr(currentNode.hostName)
resultMsg.MainIp = cert.RSAEncrypterStr(currentNode.mainIp)
resultMsg.Port = cert.RSAEncrypterStr(strconv.Itoa(currentNode.port))
resultMsg.Goos = cert.RSAEncrypterStr(currentNode.goos)
b, _ := json.Marshal(resultMsg)
//返回成功结果
n.Write(common.CMD_REG_RESULT, 0, b)
//储存节点
n.uuid = uuid
if v, ok := nodeMap.Load(uuid); !ok || v.(*node).conn.closeTag > 0 {
n.conn.node = n
nodeMap.Store(regmsg.UUID, n)
}
currentNode.broadcastNode()
case common.CMD_REG_RESULT:
var regmsg RegMsg
err = json.Unmarshal(msg.CmdData, &regmsg)
@@ -501,7 +533,7 @@ func (n *node) do(msg *common.Msg) {
}
//交换节点
go n.writeGetNodeResult(msg.CmdId)
n.writeGetNodeResult(msg.CmdId)
case common.CMD_REMOTE_REG:
var regmsg RegMsg
@@ -549,10 +581,9 @@ func (n *node) do(msg *common.Msg) {
v <- errors.New(regmsg.Err)
return
}
l := clientLock.Lock()
if n.uuid != regmsg.UUID {
var targetNode *node
if targetNode, ok = nodeMap[regmsg.UUID]; !ok {
if _v, ok := nodeMap.Load(regmsg.UUID); !ok {
targetNode = getNewNode(nodeInfo{
UUID: regmsg.UUID,
HostName: cert.RSADecrypterStr(regmsg.Hostname),
@@ -561,8 +592,9 @@ func (n *node) do(msg *common.Msg) {
Goos: cert.RSADecrypterStr(regmsg.Goos),
}, n)
nodeMap[regmsg.UUID] = targetNode
nodeMap.Store(regmsg.UUID, targetNode)
} else {
targetNode = _v.(*node)
targetNode.updateNode(nodeInfo{
UUID: regmsg.UUID,
HostName: cert.RSADecrypterStr(regmsg.Hostname),
@@ -575,8 +607,6 @@ func (n *node) do(msg *common.Msg) {
} else {
v <- n
}
l.Unlock()
n.writeGetNodeResult(msg.CmdId)
case common.CMD_PING:
@@ -769,9 +799,6 @@ func (n *node) do(msg *common.Msg) {
case common.CMD_GET_NODE:
n.writeGetNodeResult(msg.CmdId)
case common.CMD_GET_NODE_RESULT:
l := clientLock.Lock()
defer l.Unlock()
var s []nodeInfo
err = json.Unmarshal(msg.CmdData, &s)
if err == nil {
@@ -784,12 +811,12 @@ func (n *node) do(msg *common.Msg) {
Goos: cert.RSADecrypterStr(_n.Goos),
}
if _n.UUID != currentNode.uuid {
if v, ok := nodeMap[_n.UUID]; !ok {
nodeMap[_n.UUID] = getNewNode(_n, n)
if v, ok := nodeMap.Load(_n.UUID); !ok {
nodeMap.Store(_n.UUID, getNewNode(_n, n))
} else {
v.hostName = _n.HostName
v.mainIp = _n.MainIp
v.port, _ = strconv.Atoi(_n.Port)
v.(*node).hostName = _n.HostName
v.(*node).mainIp = _n.MainIp
v.(*node).port, _ = strconv.Atoi(_n.Port)
}
}
@@ -817,31 +844,27 @@ func (n *node) do(msg *common.Msg) {
case common.CMD_ADD_NODE:
var nmsg nodeInfo
err = json.Unmarshal(msg.CmdData, &nmsg)
if err != nil {
return
}
l := clientLock.Lock()
defer l.Unlock()
if v, ok := nodeMap[nmsg.UUID]; !ok {
if v, ok := nodeMap.Load(nmsg.UUID); !ok {
newNode := getNewNode(nmsg, n)
nodeMap[nmsg.UUID] = newNode
nodeMap.Store(nmsg.UUID, newNode)
} else if nmsg.UUID != currentNode.uuid {
n := v.(*node)
port, err := strconv.Atoi(cert.RSADecrypterStr(nmsg.Port))
if err == nil {
v.port = port
n.port = port
} else {
v.port = -1
n.port = -1
}
v.mainIp = cert.RSADecrypterStr(nmsg.MainIp)
v.hostName = cert.RSADecrypterStr(nmsg.HostName)
v.goos = cert.RSADecrypterStr(nmsg.Goos)
v.uuid = nmsg.UUID
nodeMap[nmsg.UUID] = v
n.mainIp = cert.RSADecrypterStr(nmsg.MainIp)
n.hostName = cert.RSADecrypterStr(nmsg.HostName)
n.goos = cert.RSADecrypterStr(nmsg.Goos)
n.uuid = nmsg.UUID
nodeMap.Store(nmsg.UUID, n)
}
case common.CMD_DIR:
@@ -1206,34 +1229,21 @@ func getNewNode(m nodeInfo, n *node) *node {
func allNodesDo(f func(*node) (bool, error)) (err error) {
var ok bool
l := clientLock.RLock()
defer l.RUnlock()
for _, n := range nodeMap {
nodeMap.Range(func(key, value interface{}) bool {
n := value.(*node)
if n.uuid != currentNode.uuid {
func() {
l.RUnlock()
defer clientLock.RLock(l)
ok, err = f(n)
}()
if err != nil {
return err
}
if !ok {
break
ok, err = f(n)
if err != nil || !ok {
return false
}
}
}
return nil
return true
})
return err
}
func (n *node) ping(id uint32) {
l := clientLock.Lock()
defer func() {
l.Unlock()
}()
l := clientLock1.Lock()
defer l.Unlock()
now := time.Now()
if n.pingTime > n.pongTime {
@@ -1293,12 +1303,7 @@ func (n *node) ping(id uint32) {
func (n *node) Delete(reason string) {
go func() {
if atomic.CompareAndSwapInt32(&n.isClose, 0, 1) {
l := clientLock.Lock()
_, ok := nodeMap[n.uuid]
if ok {
delete(nodeMap, n.uuid)
}
l.Unlock()
n.connMap.Range(func(key, value interface{}) bool {
if v, ok := value.(common.Conn); ok {
v.Close(reason)
@@ -1329,6 +1334,7 @@ func (n *node) Delete(reason string) {
n.shellMap.Delete(key)
return true
})
nodeMap.Delete(n.uuid)
}
}()
@@ -1424,27 +1430,32 @@ func (n *node) storeConn(v common.Conn) (newID uint32) {
}
func (n *node) writeGetNodeResult(id uint32) {
l := clientLock.RLock()
go func() {
defer func() {
if err := recover(); err != nil {
fmt.Println(err)
debug.PrintStack()
}
}()
var s []*nodeInfo
nodeMap.Range(func(key, value interface{}) bool {
_n := value.(*node)
if _n.uuid != currentNode.uuid {
s = append(s, &nodeInfo{
UUID: _n.uuid,
HostName: cert.RSAEncrypterStr(_n.hostName),
MainIp: cert.RSAEncrypterStr(_n.mainIp),
Port: cert.RSAEncrypterStr(strconv.Itoa(_n.port)),
Goos: cert.RSAEncrypterStr(_n.goos),
})
}
return true
})
defer l.RUnlock()
b, _ := json.Marshal(s)
n.Write(common.CMD_GET_NODE_RESULT, id, b)
}()
var s []*nodeInfo
for _, _n := range nodeMap {
if _n.uuid != currentNode.uuid {
s = append(s, &nodeInfo{
UUID: _n.uuid,
HostName: cert.RSAEncrypterStr(_n.hostName),
MainIp: cert.RSAEncrypterStr(_n.mainIp),
Port: cert.RSAEncrypterStr(strconv.Itoa(_n.port)),
Goos: cert.RSAEncrypterStr(_n.goos),
})
}
}
b, _ := json.Marshal(s)
n.Write(common.CMD_GET_NODE_RESULT, id, b)
}
func (n *node) updateNode(msg nodeInfo) {
n.hostName = msg.HostName