Add websocket function partly & Reduce duplicate code

ph4ntonn · Mar 30, 2024 · e48c50b · e48c50b
1 parent de126db
commit e48c50b
Show file tree

Hide file tree

Showing 36 changed files with 319 additions and 143 deletions.
diff --git a/admin/admin.go b/admin/admin.go
@@ -1,5 +1,4 @@
 //go:build !windows
-// +build !windows
 
 package main
 
@@ -34,7 +33,7 @@ func main() {
 
 	share.GeneratePreAuthToken(options.Secret)
 
-	protocol.DecideType("raw", options.Downstream)
+	protocol.SetUpDownStream("raw", options.Downstream)
 
 	cli.Banner()
 

diff --git a/admin/admin_win.go b/admin/admin_win.go
@@ -1,5 +1,4 @@
 //go:build windows
-// +build windows
 
 package main
 
@@ -29,7 +28,7 @@ func main() {
 
 	share.GeneratePreAuthToken(options.Secret)
 
-	protocol.DecideType("raw", options.Downstream)
+	protocol.SetUpDownStream("raw", options.Downstream)
 
 	cli.Banner()
 

diff --git a/admin/cli/interactive.go b/admin/cli/interactive.go
@@ -1,5 +1,4 @@
 //go:build !windows
-// +build !windows
 
 package cli
 
@@ -1043,7 +1042,7 @@ func (console *Console) handleNodePanelCommand(uuidNum int) {
 }
 
 func (console *Console) handleShellPanelCommand(route string, uuid string) {
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,
@@ -1071,7 +1070,7 @@ func (console *Console) handleShellPanelCommand(route string, uuid string) {
 }
 
 func (console *Console) handleSSHPanelCommand(route string, uuid string) {
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,

diff --git a/admin/cli/interactive_win.go b/admin/cli/interactive_win.go
@@ -1,4 +1,4 @@
-// +build windows
+//go:build windows
 
 package cli
 
@@ -1038,7 +1038,7 @@ func (console *Console) handleNodePanelCommand(uuidNum int) {
 }
 
 func (console *Console) handleShellPanelCommand(route string, uuid string) {
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,
@@ -1066,7 +1066,7 @@ func (console *Console) handleShellPanelCommand(route string, uuid string) {
 }
 
 func (console *Console) handleSSHPanelCommand(route string, uuid string) {
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,

diff --git a/admin/handler/backward.go b/admin/handler/backward.go
@@ -24,7 +24,7 @@ func NewBackward(lPort, rPort string) *Backward {
 }
 
 func (backward *Backward) LetBackward(mgr *manager.Manager, route string, uuid string) error {
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 	// test if node can listen on assigned port
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,
@@ -63,7 +63,7 @@ func (backward *Backward) LetBackward(mgr *manager.Manager, route string, uuid s
 }
 
 func (backward *Backward) start(mgr *manager.Manager, topo *topology.Topology, uuid string) {
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 	// first , admin need to know the route to target node,so ask topo for the answer
 	topoTask := &topology.TopoTask{
 		Mode: topology.GETROUTE,
@@ -214,7 +214,7 @@ func GetBackwardInfo(mgr *manager.Manager, uuid string) (int, bool) {
 }
 
 func StopBackward(mgr *manager.Manager, uuid, route string, choice int) {
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,

diff --git a/admin/handler/connect.go b/admin/handler/connect.go
@@ -14,7 +14,7 @@ func LetConnect(mgr *manager.Manager, route, uuid, addr string) error {
 		return err
 	}
 
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,

diff --git a/admin/handler/forward.go b/admin/handler/forward.go
@@ -28,7 +28,7 @@ func (forward *Forward) LetForward(mgr *manager.Manager, route string, uuid stri
 		return err
 	}
 
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,
@@ -103,7 +103,7 @@ func (forward *Forward) handleForwardListener(mgr *manager.Manager, listener net
 }
 
 func (forward *Forward) handleForward(mgr *manager.Manager, conn net.Conn, route string, uuid string, seq uint64) {
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 	// tell agent to start
 	startHeader := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,

diff --git a/admin/handler/info.go b/admin/handler/info.go
@@ -15,7 +15,7 @@ func AddMemo(taskChan chan *topology.TopoTask, info []string, uuid string, route
 		memo = memo + " " + i
 	}
 
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	topoTask := &topology.TopoTask{
 		Mode: topology.UPDATEMEMO,
@@ -51,7 +51,7 @@ func DelMemo(taskChan chan *topology.TopoTask, uuid string, route string) {
 	}
 	taskChan <- topoTask
 
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,

diff --git a/admin/handler/listen.go b/admin/handler/listen.go
@@ -35,7 +35,7 @@ func (listen *Listen) LetListen(mgr *manager.Manager, route, uuid string) error
 		}
 	}
 
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,
@@ -99,7 +99,7 @@ func dispatchChildUUID(mgr *manager.Manager, topo *topology.Topology, parentUUID
 	topoResult = <-topo.ResultChan
 	route := topoResult.Route
 
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,

diff --git a/admin/handler/offline.go b/admin/handler/offline.go
@@ -6,7 +6,7 @@ import (
 )
 
 func LetShutdown(route string, uuid string) {
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,

diff --git a/admin/handler/shell.go b/admin/handler/shell.go
@@ -9,7 +9,7 @@ import (
 )
 
 func LetShellStart(route string, uuid string) {
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,

diff --git a/admin/handler/socks.go b/admin/handler/socks.go
@@ -66,7 +66,7 @@ func (socks *Socks) LetSocks(mgr *manager.Manager, route string, uuid string) er
 		return err
 	}
 
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,
@@ -134,7 +134,7 @@ func (socks *Socks) handleSocksListener(mgr *manager.Manager, listener net.Liste
 }
 
 func (socks *Socks) handleSocks(mgr *manager.Manager, conn net.Conn, route string, uuid string, seq uint64) {
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,
@@ -278,7 +278,7 @@ func startUDPAss(mgr *manager.Manager, topo *topology.Topology, seq uint64) {
 		OK:  0,
 	}
 
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	defer func() {
 		if err != nil {
@@ -329,7 +329,7 @@ func startUDPAss(mgr *manager.Manager, topo *topology.Topology, seq uint64) {
 }
 
 func handleUDPAss(mgr *manager.Manager, listener *net.UDPConn, route string, uuid string, seq uint64) {
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	dataHeader := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,

diff --git a/admin/handler/ssh.go b/admin/handler/ssh.go
@@ -42,7 +42,7 @@ func (ssh *SSH) LetSSH(route string, uuid string) error {
 		}
 	}
 
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,

diff --git a/admin/handler/sshtunnel.go b/admin/handler/sshtunnel.go
@@ -38,7 +38,7 @@ func (sshTunnel *SSHTunnel) LetSSHTunnel(route, uuid string) error {
 		}
 	}
 
-	sMessage := protocol.PrepareAndDecideWhichSProtoToLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      protocol.ADMIN_UUID,

diff --git a/admin/initial/method.go b/admin/initial/method.go
@@ -30,7 +30,7 @@ func dispatchUUID(conn net.Conn, secret string) string {
 		Route:       protocol.TEMP_ROUTE,
 	}
 
-	sMessage = protocol.PrepareAndDecideWhichSProtoToLower(conn, secret, protocol.ADMIN_UUID)
+	sMessage = protocol.NewDownMsg(conn, secret, protocol.ADMIN_UUID)
 
 	protocol.ConstructMessage(sMessage, header, uuidMess, false)
 	sMessage.SendMessage()
@@ -95,12 +95,12 @@ func NormalActive(userOptions *Options, topo *topology.Topology, proxy share.Pro
 			os.Exit(0)
 		}
 
-		sMessage = protocol.PrepareAndDecideWhichSProtoToLower(conn, userOptions.Secret, protocol.ADMIN_UUID)
+		sMessage = protocol.NewDownMsg(conn, userOptions.Secret, protocol.ADMIN_UUID)
 
 		protocol.ConstructMessage(sMessage, header, hiMess, false)
 		sMessage.SendMessage()
 
-		rMessage = protocol.PrepareAndDecideWhichRProtoFromLower(conn, userOptions.Secret, protocol.ADMIN_UUID)
+		rMessage = protocol.NewDownMsg(conn, userOptions.Secret, protocol.ADMIN_UUID)
 		fHeader, fMessage, err := protocol.DestructMessage(rMessage)
 
 		if err != nil {
@@ -213,7 +213,7 @@ func NormalPassive(userOptions *Options, topo *topology.Topology) net.Conn {
 			continue
 		}
 
-		rMessage = protocol.PrepareAndDecideWhichRProtoFromLower(conn, userOptions.Secret, protocol.ADMIN_UUID)
+		rMessage = protocol.NewDownMsg(conn, userOptions.Secret, protocol.ADMIN_UUID)
 		fHeader, fMessage, err := protocol.DestructMessage(rMessage)
 
 		if err != nil {
@@ -225,7 +225,7 @@ func NormalPassive(userOptions *Options, topo *topology.Topology) net.Conn {
 		if fHeader.MessageType == protocol.HI {
 			mmess := fMessage.(*protocol.HIMess)
 			if mmess.Greeting == "Shhh..." && mmess.IsAdmin == 0 {
-				sMessage = protocol.PrepareAndDecideWhichSProtoToLower(conn, userOptions.Secret, protocol.ADMIN_UUID)
+				sMessage = protocol.NewDownMsg(conn, userOptions.Secret, protocol.ADMIN_UUID)
 				protocol.ConstructMessage(sMessage, header, hiMess, false)
 				sMessage.SendMessage()
 

diff --git a/admin/initial/parser.go b/admin/initial/parser.go
@@ -1,5 +1,4 @@
 //go:build !windows
-// +build !windows
 
 package initial
 
@@ -49,7 +48,7 @@ func init() {
 	flag.StringVar(&args.Socks5ProxyP, "socks5-proxyp", "", "socks5 password")
 	flag.StringVar(&args.HttpProxy, "http-proxy", "", "The http proxy server ip:port you want to use")
 	flag.StringVar(&args.Downstream, "down", "raw", "Downstream data type you want to use")
-	flag.StringVar(&args.Domain, "domain", "", "Domain name for TLS SNI")
+	flag.StringVar(&args.Domain, "domain", "", "Domain name for TLS SNI/WS HOST")
 	flag.BoolVar(&args.TlsEnable, "tls-enable", false, "Encrypt connection by TLS")
 
 	flag.Usage = newUsage

diff --git a/admin/initial/parser_win.go b/admin/initial/parser_win.go
@@ -1,5 +1,4 @@
 //go:build windows
-// +build windows
 
 package initial
 
@@ -47,7 +46,7 @@ func init() {
 	flag.StringVar(&args.Socks5ProxyP, "socks5-proxyp", "", "socks5 password")
 	flag.StringVar(&args.HttpProxy, "http-proxy", "", "The http proxy server ip:port you want to use")
 	flag.StringVar(&args.Downstream, "down", "raw", "")
-	flag.StringVar(&args.Domain, "domain", "", "Domain name for TLS SNI")
+	flag.StringVar(&args.Domain, "domain", "", "Domain name for TLS SNI/WS HOST")
 	flag.BoolVar(&args.TlsEnable, "tls-enable", false, "Encrypt connection by TLS")
 
 	flag.Usage = newUsage

diff --git a/admin/process/process.go b/admin/process/process.go
@@ -1,4 +1,4 @@
-// +build !windows
+//go:build !windows
 
 package process
 
@@ -51,7 +51,7 @@ func (admin *Admin) Run() {
 }
 
 func (admin *Admin) handleMessFromDownstream(console *cli.Console) {
-	rMessage := protocol.PrepareAndDecideWhichRProtoFromLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	rMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	for {
 		header, message, err := protocol.DestructMessage(rMessage)

diff --git a/admin/process/process_win.go b/admin/process/process_win.go
@@ -1,4 +1,4 @@
-// +build windows
+//go:build windows
 
 package process
 
@@ -49,7 +49,7 @@ func (admin *Admin) Run() {
 }
 
 func (admin *Admin) handleMessFromDownstream(console *cli.Console) {
-	rMessage := protocol.PrepareAndDecideWhichRProtoFromLower(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	rMessage := protocol.NewDownMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	for {
 		header, message, err := protocol.DestructMessage(rMessage)

diff --git a/agent/agent.go b/agent/agent.go
@@ -23,7 +23,7 @@ func main() {
 
 	agent := process.NewAgent(options)
 
-	protocol.DecideType(options.Upstream, options.Downstream)
+	protocol.SetUpDownStream(options.Upstream, options.Downstream)
 
 	var conn net.Conn
 	switch options.Mode {

diff --git a/agent/handler/backward.go b/agent/handler/backward.go
@@ -33,7 +33,7 @@ func (backward *Backward) start(mgr *manager.Manager) {
 	mgr.BackwardManager.TaskChan <- mgrTask
 	<-mgr.BackwardManager.ResultChan
 
-	sMessage := protocol.PrepareAndDecideWhichSProtoToUpper(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewUpMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	for {
 		conn, err := backward.Listener.Accept()
@@ -84,7 +84,7 @@ func (backward *Backward) start(mgr *manager.Manager) {
 }
 
 func (backward *Backward) handleBackward(mgr *manager.Manager, conn net.Conn, seq uint64) {
-	sMessage := protocol.PrepareAndDecideWhichSProtoToUpper(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewUpMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	defer func() {
 		finHeader := &protocol.Header{
@@ -170,7 +170,7 @@ func (backward *Backward) handleBackward(mgr *manager.Manager, conn net.Conn, se
 }
 
 func testBackward(mgr *manager.Manager, lPort, rPort string) {
-	sMessage := protocol.PrepareAndDecideWhichSProtoToUpper(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewUpMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      global.G_Component.UUID,
@@ -207,7 +207,7 @@ func testBackward(mgr *manager.Manager, lPort, rPort string) {
 func sendDoneMess(all uint16, rPort string) {
 	// here is a problem,if some of the backward conns cannot send FIN before DONE,then the FIN they send cannot be processed by admin
 	// but it's not a really big problem,because users must know some data maybe lost since they choose to close backward
-	sMessage := protocol.PrepareAndDecideWhichSProtoToUpper(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+	sMessage := protocol.NewUpMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 	header := &protocol.Header{
 		Sender:      global.G_Component.UUID,
@@ -249,7 +249,7 @@ func DispatchBackwardMess(mgr *manager.Manager) {
 				result.SeqChan <- mess.Seq
 				<-mgr.BackwardManager.SeqReady
 			} else {
-				sMessage := protocol.PrepareAndDecideWhichSProtoToUpper(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
+				sMessage := protocol.NewUpMsg(global.G_Component.Conn, global.G_Component.Secret, global.G_Component.UUID)
 
 				finHeader := &protocol.Header{
 					Sender:      global.G_Component.UUID,