[+] #28 Support local port frowarding

Author: WangYihang

lib/cli/dispatcher/tunnel.go

@@ -16,27 +16,63 @@ func (dispatcher Dispatcher) Tunnel(args []string) {
 	}
 
 	if context.Ctx.CurrentTermite != nil {
-		if len(args) != 3 {
+		if len(args) != 6 {
 			log.Error("Arguments error, use `Help Tunnel` to get more information")
 			dispatcher.TunnelHelp([]string{})
 			return
 		}
 
-		mode := args[0]
-		host := args[1]
-		port, err := strconv.ParseUint(args[2], 10, 16)
+		action := args[0]
+		mode := args[1]
+		src_host := args[2]
+		src_port, err := strconv.ParseUint(args[3], 10, 16)
 
 		if err != nil {
 			log.Error("Invalid port: %s, use `Help Tunnel` to get more information", args[1])
 			dispatcher.TunnelHelp([]string{})
 			return
 		}
 
-		switch strings.ToLower(mode) {
+		dst_host := args[4]
+		dst_port, err := strconv.ParseUint(args[5], 10, 16)
+
+		if err != nil {
+			log.Error("Invalid port: %s, use `Help Tunnel` to get more information", args[1])
+			dispatcher.TunnelHelp([]string{})
+			return
+		}
+
+		switch strings.ToLower(action) {
 		case "create":
-			context.Ctx.CurrentTermite.CreateTunnel(host, uint16(port))
+			switch strings.ToLower(mode) {
+			case "pull":
+				local_address := fmt.Sprintf("%s:%d", dst_host, dst_port)
+				remote_address := fmt.Sprintf("%s:%d", src_host, src_port)
+				log.Info("Mapping remote (%s) to local (%s)", remote_address, local_address)
+				context.AddTunnelConfig(context.Ctx.CurrentTermite, local_address, remote_address)
+			case "push":
+				// context.Ctx.CurrentTermite.CreatePushTunnel(src_host, uint16(src_port), dst_host, uint16(dst_port))
+				log.Error("TBD")
+			case "dynamic":
+				log.Error("TBD")
+			default:
+				log.Error("Invalid mode: %s, should be in {'Pull', 'Push', 'Dynamic'}", mode)
+			}
 		case "delete":
-			context.Ctx.CurrentTermite.DeleteTunnel(host, uint16(port))
+			switch strings.ToLower(mode) {
+			case "pull":
+				// context.Ctx.CurrentTermite.DeletePullTunnel(dst_host, uint16(dst_port), src_host, uint16(src_port))
+				log.Error("TBD")
+			case "push":
+				// context.Ctx.CurrentTermite.DeleteTunnel(src_host, uint16(src_port), dst_host, uint16(dst_port))
+				log.Error("TBD")
+			case "dynamic":
+				log.Error("TBD")
+			default:
+				log.Error("Invalid mode: %s, should be in {'Pull', 'Push', 'Dynamic'}", mode)
+			}
+		default:
+			log.Error("Invalid action: %s, should be in {'Create', 'Delete'}", action)
 		}
 	}
 
@@ -47,7 +83,7 @@ func (dispatcher Dispatcher) Tunnel(args []string) {
 
 func (dispatcher Dispatcher) TunnelHelp(args []string) {
 	fmt.Println("Usage of Tunnel")
-	fmt.Println("\tTunnel [Create|Delete] [Host] [Port]")
+	fmt.Println("\tTunnel [Create|Delete] [Mode] [Src Host] [Src Port] [Dst Host] [Dst Port]")
 }
 
 func (dispatcher Dispatcher) TunnelDesc(args []string) {

lib/context/client.go

@@ -77,7 +77,6 @@ func CreateTCPClient(conn net.Conn, server *TCPServer) *TCPClient {
 }
 
 func (c *TCPClient) Close() {
-	log.Debug("Closing client: %s", c.FullDesc())
 	c.conn.Close()
 }
 

lib/context/context.go

@@ -1,19 +1,33 @@
 package context
 
 import (
+	"net"
 	"os"
 	"strings"
 	"sync"
 
 	"github.com/WangYihang/Platypus/lib/util/config"
+	"github.com/WangYihang/Platypus/lib/util/log"
 	"github.com/WangYihang/Platypus/lib/util/message"
+	"github.com/WangYihang/Platypus/lib/util/str"
 	"github.com/WangYihang/Platypus/lib/util/ui"
 	"github.com/WangYihang/readline"
 	"github.com/fatih/color"
 	"github.com/gin-gonic/gin"
 	"gopkg.in/olahol/melody.v1"
 )
 
+type TunnelConfig struct {
+	Termite *TermiteClient
+	Address string
+	Server  *net.Listener
+}
+
+type TunnelInstance struct {
+	Termite *TermiteClient
+	Conn    *net.Conn
+}
+
 type Context struct {
 	Servers         map[string](*TCPServer)
 	NotifyWebSocket *melody.Melody
@@ -22,6 +36,8 @@ type Context struct {
 	CommandPrompt   string
 	RLInstance      *readline.Instance
 	Interacting     *sync.Mutex
+	TunnelConfig    map[string]TunnelConfig
+	TunnelInstance  map[string]TunnelInstance
 	// Set later in platypus.go
 	Distributor *Distributor
 	RESTful     *gin.Engine
@@ -40,6 +56,8 @@ func CreateContext() {
 			CommandPrompt:   color.CyanString("» "),
 			RLInstance:      nil,
 			Interacting:     new(sync.Mutex),
+			TunnelConfig:    make(map[string]TunnelConfig),
+			TunnelInstance:  make(map[string]TunnelInstance),
 		}
 	}
 	// Signal Handler
@@ -163,3 +181,91 @@ func Shutdown() {
 	}
 	os.Exit(0)
 }
+
+func AddTunnelConfig(termite *TermiteClient, local_address string, remote_address string) {
+	tunnel, err := net.Listen("tcp", local_address)
+	if err != nil {
+		log.Error(err.Error())
+		return
+	} else {
+		Ctx.TunnelConfig[local_address] = TunnelConfig{
+			Termite: termite,
+			Address: remote_address,
+			Server:  &tunnel,
+		}
+	}
+
+	go func() {
+		for {
+			conn, _ := tunnel.Accept()
+
+			token := str.RandomString(0x10)
+
+			termite.EncoderLock.Lock()
+			err := termite.Encoder.Encode(message.Message{
+				Type: message.TUNNEL_CONNECT,
+				Body: message.BodyTunnelConnect{
+					Token:   token,
+					Address: remote_address,
+				},
+			})
+			termite.EncoderLock.Unlock()
+
+			if err == nil {
+				Ctx.TunnelInstance[token] = TunnelInstance{
+					Conn:    &conn,
+					Termite: termite,
+				}
+			}
+		}
+	}()
+}
+
+func WriteTunnel(termite *TermiteClient, token string, data []byte) {
+	termite.AtomLock.Lock()
+	defer func() { termite.AtomLock.Unlock() }()
+
+	termite.EncoderLock.Lock()
+	err := termite.Encoder.Encode(message.Message{
+		Type: message.TUNNEL_DATA,
+		Body: message.BodyTunnelData{
+			Token: token,
+			Data:  data,
+		},
+	})
+	termite.EncoderLock.Unlock()
+
+	if err != nil {
+		log.Error("Network error: %s", err)
+	}
+}
+
+// func DeleteTunnelConfig(local_host string, local_port uint16, remote_host string, remote_port uint16) {
+// 	local_address := fmt.Sprintf("%s:%d", local_host, local_port)
+// 	remote_address := fmt.Sprintf("%s:%d", remote_host, remote_port)
+
+// 	log.Info("Unmapping from remote %s to local %s", remote_address, local_address)
+
+// 	if tc, exists := Ctx.TunnelConfig[local_address]; exists {
+// 		c.AtomLock.Lock()
+// 		defer func() { c.AtomLock.Unlock() }()
+
+// 		c.EncoderLock.Lock()
+// 		err := c.Encoder.Encode(message.Message{
+// 			Type: message.TUNNEL_DELETE,
+// 			Body: message.BodyTunnelDelete{
+// 				Key:         key,
+// 				TermiteHash: c.Hash,
+// 			},
+// 		})
+// 		c.EncoderLock.Unlock()
+
+// 		if err != nil {
+// 			log.Error("Network error: %s", err)
+// 		} else {
+// 			delete(Ctx.TunnelConfig, local_address)
+// 		}
+// 	} else {
+// 		log.Info("No such tunnel from remote %s to local %s", remote_address, local_address)
+// 	}
+// }

lib/context/server.go

@@ -20,7 +20,6 @@ import (
 	"github.com/WangYihang/Platypus/lib/util/str"
 	humanize "github.com/dustin/go-humanize"
 	"github.com/jedib0t/go-pretty/table"
-	"github.com/phayes/freeport"
 )
 
 type WebSocketMessage struct {
@@ -544,43 +543,63 @@ func TermiteMessageDispatcher(client *TermiteClient) {
 				log.Error("No such key")
 			}
 		case message.TUNNEL_CONNECTED:
-			target := msg.Body.(*message.BodyTunnelConnected).Target
-			port, _ := freeport.GetFreePort()
-			localAddress := fmt.Sprintf("0.0.0.0:%d", port)
-			tunnel, err := net.Listen("tcp", localAddress)
-			if err != nil {
-				log.Error(err.Error())
-				break
-			}
-			log.Info("%s -> %s", localAddress, target)
-			go func(target string, port int) {
-				conn, _ := tunnel.Accept()
-				Ctx.CurrentTermite.Tunnels[target] = &conn
-				log.Info("Server client tunnel connected: %v", conn)
-				for {
-					buf := make([]byte, 1024)
-					n, err := conn.Read(buf)
-					if err != nil {
-						log.Error(err.Error())
-						break
-					}
-					if n > 0 {
-						log.Info(">> %v", buf[0:n])
-						Ctx.CurrentTermite.WriteTunnel(target, buf[0:n])
+			token := msg.Body.(*message.BodyTunnelConnected).Token
+			log.Success("Tunnel (%s) connected", token)
+			if ti, exists := Ctx.TunnelInstance[token]; exists {
+				go func() {
+					for {
+						buf := make([]byte, 1024)
+						n, err := (*ti.Conn).Read(buf)
+						if err != nil {
+							log.Success("Tunnel (%s) disconnected: %s", token, err.Error())
+							ti.Termite.EncoderLock.Lock()
+							ti.Termite.Encoder.Encode(message.Message{
+								Type: message.TUNNEL_DISCONNECT,
+								Body: message.BodyTunnelDisconnect{
+									Token: token,
+								},
+							})
+							ti.Termite.EncoderLock.Unlock()
+							(*ti.Conn).Close()
+							break
+						} else {
+							if n > 0 {
+								data := buf[0:n]
+								WriteTunnel(ti.Termite, token, data)
+							}
+						}
 					}
-				}
-			}(target, port)
+				}()
+			} else {
+				log.Error("No such connection")
+			}
+		case message.TUNNEL_CONNECT_FAILED:
+			token := msg.Body.(*message.BodyTunnelConnectFailed).Token
+			reason := msg.Body.(*message.BodyTunnelConnectFailed).Reason
+			if ti, exists := Ctx.TunnelInstance[token]; exists {
+				log.Error("Connecting to %s failed: %s", token, reason)
+				(*ti.Conn).Close()
+				delete(Ctx.TunnelInstance, token)
+			} else {
+				log.Error("No such connection")
+			}
+		case message.TUNNEL_DISCONNECTED:
+			token := msg.Body.(*message.BodyTunnelDisconnected).Token
+			if ti, exists := Ctx.TunnelInstance[token]; exists {
+				log.Error("%s disconnected", token)
+				(*ti.Conn).Close()
+				delete(Ctx.TunnelInstance, token)
+			} else {
+				log.Error("No such connection")
+			}
 		case message.TUNNEL_DATA:
-			target := msg.Body.(*message.BodyTunnelData).Target
+			token := msg.Body.(*message.BodyTunnelData).Token
 			data := msg.Body.(*message.BodyTunnelData).Data
-			log.Info("%s, %v, connected", target, data)
-
-			if conn, exists := Ctx.CurrentTermite.Tunnels[target]; exists {
-				(*conn).Write(data)
+			if ti, exists := Ctx.TunnelInstance[token]; exists {
+				(*ti.Conn).Write(data)
 			} else {
-				log.Error("No such tunnel")
+				log.Error("No such connection")
 			}
-
 		}
 	}
 }