github/natpass
4
0
Fork 0
mirror of https://github.com/lwch/natpass.git synced 2024-04-21 12:41:54 +00:00
Code Issues Projects Releases Wiki Activity

1. 增加代码注释

Browse Source 
2. 增加通用encoding层
This commit is contained in:
lwch
2023-02-01 17:08:09 +08:00
parent a4c397c5ef
commit bc5d709dcb
11 changed files with 334 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
code/client/app/cmd.go
+1 -1
View File
@@ -163,7 +163,7 @@ func (a *App) Status(*cobra.Command, []string) {
}
}
// Vnc handle vnc child process
// Vnc handle vnc child process handler
func (a *App) Vnc(*cobra.Command, []string) {
defer utils.Recover("vnc.worker")
code/client/app/callback.go → code/client/app/handler.go
+15
View File
@@ -11,6 +11,21 @@ import (
"github.com/lwch/natpass/code/network"
)
/*
this is the handler function file like http.HandlerFunc.
for linked rule:
1. get rule from manager
2. create rule if not exists by its type
3. create link call NewLink
4. do the initialize logic for the link
5. response connect ok message
6. loop forward
no linked rule:
TODO
*/
func (p *program) shellCreate(mgr *rule.Mgr, conn *conn.Conn, msg *network.Msg) {
create := msg.GetCreq()
tn := mgr.GetLinked(create.GetName(), msg.GetFrom())
code/client/conn/conn.go
+62 -34
View File
@@ -20,14 +20,14 @@ const dropBlockTimeout = 10 * time.Minute
// Conn connection
type Conn struct {
sync.RWMutex
cfg *global.Configure
conn *network.Conn
cfg *global.Configure // configure
conn *network.Conn // connection wrap, read write with timeout
read map[string]chan *network.Msg // link id => channel
unknownRead chan *network.Msg // read message without link
onDisconnect chan string
write chan *network.Msg
lockDrop sync.RWMutex
drop map[string]time.Time
onDisconnect chan string // on disconnect channel, the value is clientid
write chan *network.Msg // write queue
lockDrop sync.RWMutex // drop mutex
drop map[string]time.Time // drop cache, drop message when this link is closed
// runtime
ctx context.Context
cancel context.CancelFunc
@@ -52,11 +52,12 @@ func New(cfg *global.Configure) *Conn {
return conn
}
// connect connect server and write handshake packet
func (conn *Conn) connect() error {
var dial net.Conn
var err error
if conn.cfg.UseSSL {
if conn.cfg.SSLInsecure {
if conn.cfg.SSLInsecure { // disable sni
rawConn, err := net.Dial("tcp", conn.cfg.Server)
if err != nil {
logging.Error("raw dial: %v", err)
@@ -96,6 +97,7 @@ func (conn *Conn) close() {
}
}
// writeHandshake send handshake message, default timeout is 5 seconds
func writeHandshake(conn *network.Conn, cfg *global.Configure) error {
var msg network.Msg
msg.XType = network.Msg_handshake
@@ -109,6 +111,7 @@ func writeHandshake(conn *network.Conn, cfg *global.Configure) error {
return conn.WriteMessage(&msg, 5*time.Second)
}
// isDrop check the message is dropped by linkid
func (conn *Conn) isDrop(linkID string) bool {
conn.lockDrop.RLock()
defer conn.lockDrop.RUnlock()
@@ -116,6 +119,14 @@ func (conn *Conn) isDrop(linkID string) bool {
return ok
}
// addDrop add to drop queue
func (conn *Conn) addDrop(linkID string) {
conn.lockDrop.Lock()
defer conn.lockDrop.Unlock()
conn.drop[linkID] = time.Now().Add(dropBlockTimeout)
}
// getChan get read channel by linkid
func (conn *Conn) getChan(linkID string) chan *network.Msg {
conn.RLock()
ch := conn.read[linkID]
@@ -126,19 +137,48 @@ func (conn *Conn) getChan(linkID string) chan *network.Msg {
return ch
}
func (conn *Conn) hookDispatch(ch chan *network.Msg, msg *network.Msg) bool {
// hookDispatch hook message before dispatcher
func (conn *Conn) hookDispatch(msg *network.Msg) bool {
switch msg.GetXType() {
// if disconnected add linkid to drop list, and break the handle chain
case network.Msg_disconnect:
conn.lockDrop.Lock()
conn.drop[msg.GetLinkId()] = time.Now().Add(dropBlockTimeout)
conn.lockDrop.Unlock()
conn.onDisconnect <- msg.GetLinkId()
conn.addDrop(msg.GetLinkId())
// TODO: no need will block
// conn.onDisconnect <- msg.GetLinkId()
logging.Info("connection %s disconnected", msg.GetLinkId())
return false
}
return true
}
// handleLinkedMessage linked message handler, return false means break read loop
func (conn *Conn) handleLinkedMessage(msg *network.Msg) bool {
linkID := msg.GetLinkId()
if conn.isDrop(linkID) {
return true
}
if !conn.hookDispatch(msg) {
return true
}
ch := conn.getChan(linkID)
select {
case ch <- msg:
case <-time.After(conn.cfg.WriteTimeout):
logging.Error("drop message: %s", msg.GetXType().String())
conn.addDrop(linkID)
case <-conn.ctx.Done():
return false
}
return true
}
// handleUnlinkedMessage unlinked message handler, return false means break read loop
func (conn *Conn) handleUnlinkedMessage(msg *network.Msg) bool {
// TODO
return true
}
// loopRead loop read message
func (conn *Conn) loopRead() {
defer utils.Recover("loopRead")
defer conn.close()
@@ -146,30 +186,17 @@ func (conn *Conn) loopRead() {
var timeout int
run := func(msg *network.Msg) bool {
timeout = 0
// skip keepalive message
if msg.GetXType() == network.Msg_keepalive {
return true
}
logging.Debug("read message %s(%s) from %s",
msg.GetXType().String(), msg.GetLinkId(), msg.GetFrom())
linkID := msg.GetLinkId()
if conn.isDrop(linkID) {
return true
if len(linkID) > 0 {
return conn.handleLinkedMessage(msg)
}
ch := conn.getChan(linkID)
if !conn.hookDispatch(ch, msg) {
return true
}
select {
case ch <- msg:
case <-time.After(conn.cfg.WriteTimeout):
logging.Error("drop message: %s", msg.GetXType().String())
conn.lockDrop.Lock()
conn.drop[msg.GetLinkId()] = time.Now().Add(dropBlockTimeout)
conn.lockDrop.Unlock()
case <-conn.ctx.Done():
return false
}
return true
return conn.handleUnlinkedMessage(msg)
}
for {
msg, _, err := conn.conn.ReadMessage(conn.cfg.ReadTimeout)
@@ -191,6 +218,7 @@ func (conn *Conn) loopRead() {
}
}
// loopWrite loop write message
func (conn *Conn) loopWrite() {
defer utils.Recover("loopWrite")
defer conn.close()
@@ -212,6 +240,7 @@ func (conn *Conn) loopWrite() {
}
}
// keepalive loop send keepalive message
func (conn *Conn) keepalive() {
defer utils.Recover("keepalive")
defer conn.close()
@@ -237,8 +266,8 @@ func (conn *Conn) AddLink(id string) {
conn.Unlock()
}
// Reset reset message next read
func (conn *Conn) Reset(id string, msg *network.Msg) {
// Requeue requeue for next read
func (conn *Conn) Requeue(id string, msg *network.Msg) {
conn.RLock()
ch := conn.read[id]
conn.RUnlock()
@@ -262,6 +291,7 @@ func (conn *Conn) ChanDisconnect() <-chan string {
return conn.onDisconnect
}
// checkDrop clear timeouted drop queue
func (conn *Conn) checkDrop() {
for {
time.Sleep(time.Second)
@@ -298,7 +328,5 @@ func (conn *Conn) ChanClose(id string) {
delete(conn.read, id)
conn.Unlock()
conn.lockDrop.Lock()
conn.drop[id] = time.Now().Add(dropBlockTimeout)
conn.lockDrop.Unlock()
conn.addDrop(id)
}
code/client/rule/code/code.go
+1 -1
View File
@@ -143,7 +143,7 @@ func (code *Code) new(conn *conn.Conn) (string, error) {
return "", errWaitingTimeout
}
if msg.GetXType() != network.Msg_connect_rep {
conn.Reset(id, msg)
conn.Requeue(id, msg)
time.Sleep(code.readTimeout / 10)
continue
}
code/client/rule/shell/h_new.go
+1 -1
View File
@@ -35,7 +35,7 @@ func (shell *Shell) New(conn *conn.Conn, w http.ResponseWriter, r *http.Request)
return
}
if msg.GetXType() != network.Msg_connect_rep {
conn.Reset(id, msg)
conn.Requeue(id, msg)
time.Sleep(shell.readTimeout / 10)
continue
}
code/client/rule/vnc/h_new.go
+1 -1
View File
@@ -51,7 +51,7 @@ func (v *VNC) New(conn *conn.Conn, w http.ResponseWriter, r *http.Request) {
return
}
if msg.GetXType() != network.Msg_connect_rep {
conn.Reset(id, msg)
conn.Requeue(id, msg)
time.Sleep(v.readTimeout / 10)
continue
}