[+] #28 Support remote port frowarding

lib/context/context.go

@@ -70,6 +70,8 @@ func CreateContext() {
 			Interacting:        new(sync.Mutex),
 			PullTunnelConfig:   make(map[string]PullTunnelConfig),
 			PullTunnelInstance: make(map[string]PullTunnelInstance),
+			PushTunnelConfig:   make(map[string]PushTunnelConfig),
+			PushTunnelInstance: make(map[string]PushTunnelInstance),
 		}
 	}
 	// Signal Handler
@@ -195,18 +197,29 @@ func Shutdown() {
 }
 
 func AddPushTunnelConfig(termite *TermiteClient, local_address string, remote_address string) {
-	// termite.AtomLock.Lock()
-	// defer func() { termite.AtomLock.Unlock() }()
-
-	// termite.EncoderLock.Lock()
-	// err := termite.Encoder.Encode(message.Message{
-	// 	Type: message.PUSH_PULL_TUNNEL_CREATE,
-	// 	Body: message.BodyPushTunnelCreate{
-	// 		Token: token,
-	// 		Data:  data,
-	// 	},
-	// })
-	// termite.EncoderLock.Unlock()
+	token := str.RandomString(0x10)
+
+	termite.AtomLock.Lock()
+	defer func() { termite.AtomLock.Unlock() }()
+
+	termite.EncoderLock.Lock()
+	err := termite.Encoder.Encode(message.Message{
+		Type: message.PUSH_TUNNEL_CREATE,
+		Body: message.BodyPushTunnelCreate{
+			Token:   token,
+			Address: remote_address,
+		},
+	})
+	termite.EncoderLock.Unlock()
+
+	if err != nil {
+		log.Error(err.Error())
+	} else {
+		Ctx.PushTunnelConfig[token] = PushTunnelConfig{
+			Termite: termite,
+			Address: local_address,
+		}
+	}
 }
 
 func AddPullTunnelConfig(termite *TermiteClient, local_address string, remote_address string) {

lib/context/server.go

@@ -507,6 +507,7 @@ func TermiteMessageDispatcher(client *TermiteClient) {
 		var key string
 		switch msg.Type {
 		case message.STDIO:
+			log.Debug("case message.STDIO")
 			key = msg.Body.(*message.BodyStdio).Key
 			if process, exists := client.Processes[key]; exists {
 				if process.WebSocket != nil {
@@ -518,6 +519,7 @@ func TermiteMessageDispatcher(client *TermiteClient) {
 				log.Error("No such key")
 			}
 		case message.PROCESS_STARTED:
+			log.Debug("case message.PROCESS_STARTED")
 			key = msg.Body.(*message.BodyProcessStarted).Key
 			if process, exists := client.Processes[key]; exists {
 				process.Pid = msg.Body.(*message.BodyProcessStarted).Pid
@@ -530,6 +532,7 @@ func TermiteMessageDispatcher(client *TermiteClient) {
 				log.Error("No such key")
 			}
 		case message.PROCESS_STOPED:
+			log.Debug("case message.PROCESS_STOPED")
 			key = msg.Body.(*message.BodyProcessStoped).Key
 			if process, exists := client.Processes[key]; exists {
 				code := msg.Body.(*message.BodyProcessStoped).Code
@@ -543,6 +546,7 @@ func TermiteMessageDispatcher(client *TermiteClient) {
 				log.Error("No such key")
 			}
 		case message.PULL_TUNNEL_CONNECTED:
+			log.Debug("case message.PULL_TUNNEL_CONNECTED")
 			token := msg.Body.(*message.BodyPullTunnelConnected).Token
 			log.Success("Tunnel (%s) connected", token)
 			if ti, exists := Ctx.PullTunnelInstance[token]; exists {
@@ -574,6 +578,7 @@ func TermiteMessageDispatcher(client *TermiteClient) {
 				log.Error("No such connection")
 			}
 		case message.PULL_TUNNEL_CONNECT_FAILED:
+			log.Debug("case message.PULL_TUNNEL_CONNECT_FAILED")
 			token := msg.Body.(*message.BodyPullTunnelConnectFailed).Token
 			reason := msg.Body.(*message.BodyPullTunnelConnectFailed).Reason
 			if ti, exists := Ctx.PullTunnelInstance[token]; exists {
@@ -584,6 +589,7 @@ func TermiteMessageDispatcher(client *TermiteClient) {
 				log.Error("No such connection")
 			}
 		case message.PULL_TUNNEL_DISCONNECTED:
+			log.Debug("case message.PULL_TUNNEL_DISCONNECTED")
 			token := msg.Body.(*message.BodyPullTunnelDisconnected).Token
 			if ti, exists := Ctx.PullTunnelInstance[token]; exists {
 				log.Error("%s disconnected", token)
@@ -593,13 +599,131 @@ func TermiteMessageDispatcher(client *TermiteClient) {
 				log.Error("No such connection")
 			}
 		case message.PULL_TUNNEL_DATA:
+			log.Debug("case message.PULL_TUNNEL_DATA")
 			token := msg.Body.(*message.BodyPullTunnelData).Token
 			data := msg.Body.(*message.BodyPullTunnelData).Data
 			if ti, exists := Ctx.PullTunnelInstance[token]; exists {
 				(*ti.Conn).Write(data)
 			} else {
 				log.Error("No such connection")
 			}
+		case message.PUSH_TUNNEL_CONNECT:
+			log.Debug("case message.PUSH_TUNNEL_CONNECT")
+			token := msg.Body.(*message.BodyPushTunnelConnect).Token
+			if tc, exists := Ctx.PushTunnelConfig[token]; exists {
+				conn, err := net.Dial("tcp", tc.Address)
+				if err != nil {
+					tc.Termite.Encoder.Encode(message.Message{
+						Type: message.PUSH_TUNNEL_CONNECT_FAILED,
+						Body: message.BodyPushTunnelConnectFailed{
+							Token:  token,
+							Reason: err.Error(),
+						},
+					})
+				} else {
+					Ctx.PushTunnelInstance[token] = PushTunnelInstance{
+						Termite: tc.Termite,
+						Conn:    &conn,
+					}
+					tc.Termite.Encoder.Encode(message.Message{
+						Type: message.PUSH_TUNNEL_CONNECTED,
+						Body: message.BodyPushTunnelConnected{
+							Token: token,
+						},
+					})
+					go func() {
+						for {
+							buf := make([]byte, 1024)
+							n, err := conn.Read(buf)
+							if err != nil {
+								tc.Termite.Encoder.Encode(message.Message{
+									Type: message.PUSH_TUNNEL_CONNECT_FAILED,
+									Body: message.BodyPushTunnelConnectFailed{
+										Token:  token,
+										Reason: err.Error(),
+									},
+								})
+								conn.Close()
+								delete(Ctx.PushTunnelInstance, token)
+								break
+							} else {
+								tc.Termite.Encoder.Encode(message.Message{
+									Type: message.PUSH_TUNNEL_DATA,
+									Body: message.BodyPushTunnelData{
+										Token: token,
+										Data:  buf[0:n],
+									},
+								})
+							}
+						}
+					}()
+				}
+			} else {
+				log.Error("No such tunnel")
+			}
+		case message.PUSH_TUNNEL_DISCONNECT:
+			log.Debug("case message.PUSH_TUNNEL_DISCONNECT")
+			token := msg.Body.(*message.BodyPushTunnelDisonnect).Token
+			if ti, exists := Ctx.PushTunnelInstance[token]; exists {
+				log.Success("Tunnel %v closed", (*ti.Conn).RemoteAddr().String())
+				(*ti.Conn).Close()
+			} else {
+				log.Error("No such tunnel")
+			}
+		case message.PUSH_TUNNEL_DISCONNECTED:
+			log.Debug("case message.PUSH_TUNNEL_DISCONNECTED")
+			token := msg.Body.(*message.BodyPushTunnelDisonnected).Token
+			if ti, exists := Ctx.PushTunnelInstance[token]; exists {
+				(*ti.Conn).Close()
+				delete(Ctx.PushTunnelInstance, token)
+			} else {
+				log.Error("No such tunnel")
+			}
+		case message.PUSH_TUNNEL_CREATED:
+			log.Debug("case message.PUSH_TUNNEL_CREATED")
+			token := msg.Body.(*message.BodyPushTunnelCreated).Token
+			if tc, exists := Ctx.PushTunnelConfig[token]; exists {
+				log.Success("Tunnel created: %s", tc.Address)
+			} else {
+				log.Error("No such tunnel")
+			}
+		case message.PUSH_TUNNEL_CREATE_FAILED:
+			log.Debug("case message.PUSH_TUNNEL_CREATE_FAILED")
+			token := msg.Body.(*message.BodyPushTunnelCreateFailed).Token
+			reason := msg.Body.(*message.BodyPushTunnelCreateFailed).Reason
+			if tc, exists := Ctx.PushTunnelConfig[token]; exists {
+				log.Success("Tunnel create failed: %s", tc.Address, reason)
+			} else {
+				log.Error("No such tunnel")
+			}
+		case message.PUSH_TUNNEL_DELETED:
+			log.Debug("case message.PUSH_TUNNEL_DELETED")
+			// TODO
+			log.Info("PUSH_TUNNEL_DELETED")
+		case message.PUSH_TUNNEL_DELETE_FAILED:
+			log.Debug("case message.PUSH_TUNNEL_DELETE_FAILED")
+			// TODO
+			log.Info("PUSH_TUNNEL_DELETE_FAILED")
+		case message.PUSH_TUNNEL_DATA:
+			log.Debug("case message.PUSH_TUNNEL_DATA")
+			token := msg.Body.(*message.BodyPushTunnelData).Token
+			data := msg.Body.(*message.BodyPushTunnelData).Data
+			if ti, exists := Ctx.PushTunnelInstance[token]; exists {
+				_, err := (*ti.Conn).Write(data)
+				if err != nil {
+					ti.Termite.Encoder.Encode(message.Message{
+						Type: message.PUSH_TUNNEL_CONNECT_FAILED,
+						Body: message.BodyPushTunnelConnectFailed{
+							Token:  token,
+							Reason: err.Error(),
+						},
+					})
+					(*ti.Conn).Close()
+					delete(Ctx.PushTunnelInstance, token)
+				}
+			} else {
+				log.Error("No such tunnel")
+			}
 		}
 	}
 }

lib/context/termite.go

@@ -277,14 +277,25 @@ func (c *TermiteClient) System(command string) string {
 
 func (c *TermiteClient) Close() {
 	log.Info("Closing client: %s", c.FullDesc())
+	for k, ti := range Ctx.PushTunnelInstance {
+		if ti.Termite == c && ti.Conn != nil {
+			delete(Ctx.PushTunnelInstance, k)
+		}
+	}
+	for k, tc := range Ctx.PushTunnelConfig {
+		if tc.Termite == c {
+			delete(Ctx.PushTunnelConfig, k)
+		}
+	}
+
 	for k, ti := range Ctx.PullTunnelInstance {
 		if ti.Termite == c && ti.Conn != nil {
 			delete(Ctx.PullTunnelInstance, k)
 		}
 	}
 	for k, tc := range Ctx.PullTunnelConfig {
 		if tc.Termite == c {
-			log.Info("Removing tunnel config from %s to %s", (*tc.Server).Addr().String(), tc.Address)
+			log.Info("Removing pull tunnel config from %s to %s", (*tc.Server).Addr().String(), tc.Address)
 			(*tc.Server).Close()
 			delete(Ctx.PullTunnelConfig, k)
 		}

lib/util/message/message.go

@@ -10,16 +10,22 @@ const (
 	// Platypus <-> Termite
 	STDIO MessageType = iota
 	PULL_TUNNEL_DATA
+	PUSH_TUNNEL_DATA
 
 	// Platypus -> Termite
 	WINDOW_SIZE
 	GET_CLIENT_INFO
 	DUPLICATED_CLIENT
 	PROCESS_START
 	PROCESS_TERMINATE
-	Pull_TUNNEL_CREATE
 	PULL_TUNNEL_CONNECT
 	PULL_TUNNEL_DISCONNECT
+	PUSH_TUNNEL_CREATE
+	PUSH_TUNNEL_DELETE
+	PUSH_TUNNEL_CONNECTED
+	PUSH_TUNNEL_CONNECT_FAILED
+	PUSH_TUNNEL_DISCONNECTED
+	PUSH_TUNNEL_DISCONNECT_FAILED
 
 	// Termite -> Platypus
 	PROCESS_STARTED
@@ -28,6 +34,12 @@ const (
 	PULL_TUNNEL_CONNECTED
 	PULL_TUNNEL_CONNECT_FAILED
 	PULL_TUNNEL_DISCONNECTED
+	PUSH_TUNNEL_CONNECT
+	PUSH_TUNNEL_DISCONNECT
+	PUSH_TUNNEL_CREATED
+	PUSH_TUNNEL_CREATE_FAILED
+	PUSH_TUNNEL_DELETED
+	PUSH_TUNNEL_DELETE_FAILED
 )
 
 type Message struct {
@@ -106,20 +118,94 @@ type BodyPullTunnelData struct {
 	Data  []byte
 }
 
+type BodyPushTunnelData struct {
+	Token string
+	Data  []byte
+}
+
+type BodyPushTunnelCreate struct {
+	Token   string
+	Address string
+}
+
+type BodyPushTunnelCreated struct {
+	Token string
+}
+
+type BodyPushTunnelCreateFailed struct {
+	Token  string
+	Reason string
+}
+
+type BodyPushTunnelDelete struct {
+	Token string
+}
+
+type BodyPushTunnelDeleted struct {
+	Token string
+}
+
+type BodyPushTunnelDeleteFailed struct {
+	Token  string
+	Reason string
+}
+
+type BodyPushTunnelConnect struct {
+	Token string
+}
+
+type BodyPushTunnelConnected struct {
+	Token string
+}
+type BodyPushTunnelConnectFailed struct {
+	Token  string
+	Reason string
+}
+
+type BodyPushTunnelDisonnect struct {
+	Token string
+}
+
+type BodyPushTunnelDisonnected struct {
+	Token  string
+	Reason string
+}
+
+type BodyPushTunnelDisonnectFailed struct {
+	Token string
+}
+
 func RegisterGob() {
+	// Client Management
+	gob.Register(&BodyClientInfo{})
+	gob.Register(&BodyGetClientInfo{})
+	gob.Register(&BodyDuplicateClient{})
+	// Process management
 	gob.Register(&BodyStdio{})
-	gob.Register(&BodyWindowSize{})
 	gob.Register(&BodyStartProcess{})
 	gob.Register(&BodyProcessStarted{})
 	gob.Register(&BodyProcessStoped{})
-	gob.Register(&BodyGetClientInfo{})
-	gob.Register(&BodyDuplicateClient{})
-	gob.Register(&BodyClientInfo{})
 	gob.Register(&BodyTerminateProcess{})
+	gob.Register(&BodyWindowSize{})
+	// Local port forwarding
 	gob.Register(&BodyPullTunnelConnect{})
 	gob.Register(&BodyPullTunnelConnected{})
 	gob.Register(&BodyPullTunnelConnectFailed{})
 	gob.Register(&BodyPullTunnelDisconnect{})
 	gob.Register(&BodyPullTunnelDisconnected{})
 	gob.Register(&BodyPullTunnelData{})
+	// Remote port forwarding
+	gob.Register(&BodyPushTunnelData{})
+	gob.Register(&BodyPushTunnelCreate{})
+	gob.Register(&BodyPushTunnelCreated{})
+	gob.Register(&BodyPushTunnelCreateFailed{})
+	gob.Register(&BodyPushTunnelDelete{})
+	gob.Register(&BodyPushTunnelDeleted{})
+	gob.Register(&BodyPushTunnelDeleteFailed{})
+	gob.Register(&BodyPushTunnelConnect{})
+	gob.Register(&BodyPushTunnelConnected{})
+	gob.Register(&BodyPushTunnelConnectFailed{})
+	gob.Register(&BodyPushTunnelDisonnect{})
+	gob.Register(&BodyPushTunnelDisonnected{})
+	gob.Register(&BodyPushTunnelDisonnectFailed{})
 }