优化代码结构

This commit is contained in:
lwch
2021-11-05 11:40:19 +08:00
parent a62c190ccc
commit 268af0b63e
4 changed files with 129 additions and 108 deletions
+1 -1
View File
@@ -50,4 +50,4 @@ jobs:
go run contrib/bindata/main.go -pkg shell -o code/client/tunnel/shell/assets.go -prefix html/shell html/shell/...
go run contrib/bindata/main.go -pkg vnc -o code/client/tunnel/vnc/assets.go -prefix html/vnc html/vnc/...
go run contrib/bindata/main.go -pkg dashboard -o code/client/dashboard/assets.go -prefix html/dashboard html/dashboard/...
go build -v code/client/main.go code/client/connect.go
go build -v code/client/main.go
+120
View File
@@ -0,0 +1,120 @@
package app
import (
"natpass/code/client/dashboard"
"natpass/code/client/global"
"natpass/code/client/pool"
"natpass/code/client/tunnel"
"natpass/code/client/tunnel/reverse"
"natpass/code/client/tunnel/shell"
"natpass/code/client/tunnel/vnc"
"natpass/code/network"
rt "runtime"
"time"
"github.com/kardianos/service"
"github.com/lwch/logging"
"github.com/lwch/runtime"
)
// App application
type App struct {
confDir string
cfg *global.Configure
version string
}
// New create application
func New(ver, dir string, cfg *global.Configure) *App {
return &App{version: ver, confDir: dir, cfg: cfg}
}
// Start start application
func (a *App) Start(s service.Service) error {
go a.run()
return nil
}
// Stop stop application
func (a *App) Stop(s service.Service) error {
return nil
}
func (a *App) run() {
// go func() {
// http.ListenAndServe(":9000", nil)
// }()
stdout := true
if rt.GOOS == "windows" {
stdout = false
}
logging.SetSizeRotate(a.cfg.LogDir, "np-cli", int(a.cfg.LogSize.Bytes()), a.cfg.LogRotate, stdout)
defer logging.Flush()
pl := pool.New(a.cfg)
mgr := tunnel.New()
for _, t := range a.cfg.Tunnels {
switch t.Type {
case "tcp", "udp":
tn := reverse.New(t)
mgr.Add(tn)
go tn.Handle(pl)
case "shell":
sh := shell.New(t)
mgr.Add(sh)
go sh.Handle(pl)
case "vnc":
v := vnc.New(t)
mgr.Add(v)
go v.Handle(pl)
}
}
for i := 0; i < a.cfg.Links-pl.Size(); i++ {
go func() {
for {
conn := pl.Get()
if conn == nil {
time.Sleep(time.Second)
continue
}
for {
msg := <-conn.ChanUnknown()
if msg == nil {
break
}
var linkID string
switch msg.GetXType() {
case network.Msg_connect_req:
switch msg.GetCreq().GetXType() {
case network.ConnectRequest_tcp, network.ConnectRequest_udp:
a.connect(mgr, conn, msg)
case network.ConnectRequest_shell:
a.shellCreate(mgr, conn, msg)
case network.ConnectRequest_vnc:
a.vncCreate(a.confDir, mgr, conn, msg)
}
default:
linkID = msg.GetLinkId()
}
if len(linkID) > 0 {
logging.Error("link of %s on connection %d not found, type=%s",
linkID, conn.Idx, msg.GetXType().String())
continue
}
}
logging.Info("connection %s-%d exited", a.cfg.ID, conn.Idx)
time.Sleep(time.Second)
}
}()
}
if a.cfg.DashboardEnabled {
db := dashboard.New(a.cfg, pl, mgr, a.version)
runtime.Assert(db.ListenAndServe(a.cfg.DashboardListen, a.cfg.DashboardPort))
} else {
select {}
}
}
@@ -1,4 +1,4 @@
package main
package app
import (
"fmt"
@@ -15,7 +15,7 @@ import (
"github.com/lwch/logging"
)
func connect(mgr *tunnel.Mgr, conn *pool.Conn, msg *network.Msg) {
func (a *App) connect(mgr *tunnel.Mgr, conn *pool.Conn, msg *network.Msg) {
req := msg.GetCreq()
// TODO: 创建tcp连接移到reverse包中实现
dial := "tcp"
@@ -50,7 +50,7 @@ func connect(mgr *tunnel.Mgr, conn *pool.Conn, msg *network.Msg) {
lk.OnWork <- struct{}{}
}
func shellCreate(mgr *tunnel.Mgr, conn *pool.Conn, msg *network.Msg) {
func (a *App) shellCreate(mgr *tunnel.Mgr, conn *pool.Conn, msg *network.Msg) {
create := msg.GetCreq()
tn := mgr.Get(create.GetName(), msg.GetFrom())
if tn == nil {
@@ -74,7 +74,7 @@ func shellCreate(mgr *tunnel.Mgr, conn *pool.Conn, msg *network.Msg) {
lk.Forward()
}
func vncCreate(confDir string, mgr *tunnel.Mgr, conn *pool.Conn, msg *network.Msg) {
func (a *App) vncCreate(confDir string, mgr *tunnel.Mgr, conn *pool.Conn, msg *network.Msg) {
create := msg.GetCreq()
tn := mgr.Get(create.GetName(), msg.GetFrom())
if tn == nil {
+4 -103
View File
@@ -3,19 +3,13 @@ package main
import (
"flag"
"fmt"
"natpass/code/client/dashboard"
"natpass/code/client/app"
"natpass/code/client/global"
"natpass/code/client/pool"
"natpass/code/client/tunnel"
"natpass/code/client/tunnel/reverse"
"natpass/code/client/tunnel/shell"
"natpass/code/client/tunnel/vnc"
"natpass/code/network"
"natpass/code/utils"
"os"
"path/filepath"
rt "runtime"
"time"
_ "net/http/pprof"
@@ -39,110 +33,17 @@ func showVersion() {
os.Exit(0)
}
type app struct {
confDir string
cfg *global.Configure
}
func (a *app) Start(s service.Service) error {
go a.run()
return nil
}
func (a *app) run() {
// go func() {
// http.ListenAndServe(":9000", nil)
// }()
stdout := true
if rt.GOOS == "windows" {
stdout = false
}
logging.SetSizeRotate(a.cfg.LogDir, "np-cli", int(a.cfg.LogSize.Bytes()), a.cfg.LogRotate, stdout)
defer logging.Flush()
pl := pool.New(a.cfg)
mgr := tunnel.New()
for _, t := range a.cfg.Tunnels {
switch t.Type {
case "tcp", "udp":
tn := reverse.New(t)
mgr.Add(tn)
go tn.Handle(pl)
case "shell":
sh := shell.New(t)
mgr.Add(sh)
go sh.Handle(pl)
case "vnc":
v := vnc.New(t)
mgr.Add(v)
go v.Handle(pl)
}
}
for i := 0; i < a.cfg.Links-pl.Size(); i++ {
go func() {
for {
conn := pl.Get()
if conn == nil {
time.Sleep(time.Second)
continue
}
for {
msg := <-conn.ChanUnknown()
if msg == nil {
break
}
var linkID string
switch msg.GetXType() {
case network.Msg_connect_req:
switch msg.GetCreq().GetXType() {
case network.ConnectRequest_tcp, network.ConnectRequest_udp:
connect(mgr, conn, msg)
case network.ConnectRequest_shell:
shellCreate(mgr, conn, msg)
case network.ConnectRequest_vnc:
vncCreate(a.confDir, mgr, conn, msg)
}
default:
linkID = msg.GetLinkId()
}
if len(linkID) > 0 {
logging.Error("link of %s on connection %d not found, type=%s",
linkID, conn.Idx, msg.GetXType().String())
continue
}
}
logging.Info("connection %s-%d exited", a.cfg.ID, conn.Idx)
time.Sleep(time.Second)
}
}()
}
if a.cfg.DashboardEnabled {
db := dashboard.New(a.cfg, pl, mgr, version)
runtime.Assert(db.ListenAndServe(a.cfg.DashboardListen, a.cfg.DashboardPort))
} else {
select {}
}
}
func (a *app) Stop(s service.Service) error {
return nil
}
func main() {
user := flag.String("user", "", "service user")
conf := flag.String("conf", "", "configure file path")
version := flag.Bool("version", false, "show version info")
ver := flag.Bool("version", false, "show version info")
act := flag.String("action", "", "install or uninstall")
name := flag.String("name", "", "tunnel name")
vport := flag.Uint("vport", 6155, "vnc worker listen port")
vcursor := flag.Bool("vcursor", false, "vnc show cursor")
flag.Parse()
if *version {
if *ver {
showVersion()
os.Exit(0)
}
@@ -191,7 +92,7 @@ func main() {
return
}
app := &app{confDir: *conf, cfg: cfg}
app := app.New(version, *conf, cfg)
sv, err := service.New(app, appCfg)
runtime.Assert(err)