Skip to content

Commit

Permalink
[+] Prepare for Push Tunnel
Browse files Browse the repository at this point in the history
  • Loading branch information
@WangYihang
WangYihang committed Jul 11, 2021
1 parent c134085 commit 29e04ac
Show file tree
Hide file tree
Showing 5 changed files with 105 additions and 77 deletions.
8 changes: 5 additions & 3 deletions lib/cli/dispatcher/tunnel.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,10 +49,12 @@ func (dispatcher Dispatcher) Tunnel(args []string) {
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)
context.AddPullTunnelConfig(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")
local_address := fmt.Sprintf("%s:%d", src_host, src_port)
remote_address := fmt.Sprintf("%s:%d", dst_host, dst_port)
log.Info("Mapping local (%s) to remote (%s)", local_address, remote_address)
context.AddPushTunnelConfig(context.Ctx.CurrentTermite, local_address, remote_address)
case "dynamic":
log.Error("TBD")
default:
Expand Down
91 changes: 59 additions & 32 deletions lib/context/context.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,27 +17,39 @@ import (
"gopkg.in/olahol/melody.v1"
)

type TunnelConfig struct {
type PullTunnelConfig struct {
Termite *TermiteClient
Address string
Server *net.Listener
}

type TunnelInstance struct {
type PullTunnelInstance struct {
Termite *TermiteClient
Conn *net.Conn
}

type PushTunnelConfig struct {
Termite *TermiteClient
Address string
}

type PushTunnelInstance struct {
Termite *TermiteClient
Conn *net.Conn
}

type Context struct {
Servers map[string](*TCPServer)
NotifyWebSocket *melody.Melody
Current *TCPClient
CurrentTermite *TermiteClient
CommandPrompt string
RLInstance *readline.Instance
Interacting *sync.Mutex
TunnelConfig map[string]TunnelConfig
TunnelInstance map[string]TunnelInstance
Servers map[string](*TCPServer)
NotifyWebSocket *melody.Melody
Current *TCPClient
CurrentTermite *TermiteClient
CommandPrompt string
RLInstance *readline.Instance
Interacting *sync.Mutex
PullTunnelConfig map[string]PullTunnelConfig
PullTunnelInstance map[string]PullTunnelInstance
PushTunnelConfig map[string]PushTunnelConfig
PushTunnelInstance map[string]PushTunnelInstance
// Set later in platypus.go
Distributor *Distributor
RESTful *gin.Engine
Expand All @@ -49,15 +61,15 @@ var Ctx *Context
func CreateContext() {
if Ctx == nil {
Ctx = &Context{
Servers: make(map[string](*TCPServer)),
NotifyWebSocket: nil,
Current: nil,
CurrentTermite: nil,
CommandPrompt: color.CyanString("» "),
RLInstance: nil,
Interacting: new(sync.Mutex),
TunnelConfig: make(map[string]TunnelConfig),
TunnelInstance: make(map[string]TunnelInstance),
Servers: make(map[string](*TCPServer)),
NotifyWebSocket: nil,
Current: nil,
CurrentTermite: nil,
CommandPrompt: color.CyanString("» "),
RLInstance: nil,
Interacting: new(sync.Mutex),
PullTunnelConfig: make(map[string]PullTunnelConfig),
PullTunnelInstance: make(map[string]PullTunnelInstance),
}
}
// Signal Handler
Expand Down Expand Up @@ -182,13 +194,28 @@ func Shutdown() {
os.Exit(0)
}

func AddTunnelConfig(termite *TermiteClient, local_address string, remote_address string) {
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()
}

func AddPullTunnelConfig(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{
Ctx.PullTunnelConfig[local_address] = PullTunnelConfig{
Termite: termite,
Address: remote_address,
Server: &tunnel,
Expand All @@ -203,16 +230,16 @@ func AddTunnelConfig(termite *TermiteClient, local_address string, remote_addres

termite.EncoderLock.Lock()
err := termite.Encoder.Encode(message.Message{
Type: message.TUNNEL_CONNECT,
Body: message.BodyTunnelConnect{
Type: message.PULL_TUNNEL_CONNECT,
Body: message.BodyPullTunnelConnect{
Token: token,
Address: remote_address,
},
})
termite.EncoderLock.Unlock()

if err == nil {
Ctx.TunnelInstance[token] = TunnelInstance{
Ctx.PullTunnelInstance[token] = PullTunnelInstance{
Conn: &conn,
Termite: termite,
}
Expand All @@ -227,8 +254,8 @@ func WriteTunnel(termite *TermiteClient, token string, data []byte) {

termite.EncoderLock.Lock()
err := termite.Encoder.Encode(message.Message{
Type: message.TUNNEL_DATA,
Body: message.BodyTunnelData{
Type: message.PULL_TUNNEL_DATA,
Body: message.BodyPullTunnelData{
Token: token,
Data: data,
},
Expand All @@ -240,20 +267,20 @@ func WriteTunnel(termite *TermiteClient, token string, data []byte) {
}
}

// func DeleteTunnelConfig(local_host string, local_port uint16, remote_host string, remote_port uint16) {
// func DeletePullTunnelConfig(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 {
// if tc, exists := Ctx.PullTunnelConfig[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{
// Type: message.PULL_TUNNEL_DELETE,
// Body: message.BodyPullTunnelDelete{
// Key: key,
// TermiteHash: c.Hash,
// },
Expand All @@ -263,7 +290,7 @@ func WriteTunnel(termite *TermiteClient, token string, data []byte) {
// if err != nil {
// log.Error("Network error: %s", err)
// } else {
// delete(Ctx.TunnelConfig, local_address)
// delete(Ctx.PullTunnelConfig, local_address)
// }
// } else {
// log.Info("No such tunnel from remote %s to local %s", remote_address, local_address)
Expand Down
36 changes: 18 additions & 18 deletions lib/context/server.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -542,10 +542,10 @@ func TermiteMessageDispatcher(client *TermiteClient) {
} else {
log.Error("No such key")
}
case message.TUNNEL_CONNECTED:
token := msg.Body.(*message.BodyTunnelConnected).Token
case message.PULL_TUNNEL_CONNECTED:
token := msg.Body.(*message.BodyPullTunnelConnected).Token
log.Success("Tunnel (%s) connected", token)
if ti, exists := Ctx.TunnelInstance[token]; exists {
if ti, exists := Ctx.PullTunnelInstance[token]; exists {
go func() {
for {
buf := make([]byte, 1024)
Expand All @@ -554,8 +554,8 @@ func TermiteMessageDispatcher(client *TermiteClient) {
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{
Type: message.PULL_TUNNEL_DISCONNECT,
Body: message.BodyPullTunnelDisconnect{
Token: token,
},
})
Expand All @@ -573,29 +573,29 @@ func TermiteMessageDispatcher(client *TermiteClient) {
} 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 {
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 {
log.Error("Connecting to %s failed: %s", token, reason)
(*ti.Conn).Close()
delete(Ctx.TunnelInstance, token)
delete(Ctx.PullTunnelInstance, token)
} else {
log.Error("No such connection")
}
case message.TUNNEL_DISCONNECTED:
token := msg.Body.(*message.BodyTunnelDisconnected).Token
if ti, exists := Ctx.TunnelInstance[token]; exists {
case message.PULL_TUNNEL_DISCONNECTED:
token := msg.Body.(*message.BodyPullTunnelDisconnected).Token
if ti, exists := Ctx.PullTunnelInstance[token]; exists {
log.Error("%s disconnected", token)
(*ti.Conn).Close()
delete(Ctx.TunnelInstance, token)
delete(Ctx.PullTunnelInstance, token)
} else {
log.Error("No such connection")
}
case message.TUNNEL_DATA:
token := msg.Body.(*message.BodyTunnelData).Token
data := msg.Body.(*message.BodyTunnelData).Data
if ti, exists := Ctx.TunnelInstance[token]; exists {
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")
Expand Down
8 changes: 4 additions & 4 deletions lib/context/termite.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -277,16 +277,16 @@ func (c *TermiteClient) System(command string) string {

func (c *TermiteClient) Close() {
log.Info("Closing client: %s", c.FullDesc())
for k, ti := range Ctx.TunnelInstance {
for k, ti := range Ctx.PullTunnelInstance {
if ti.Termite == c && ti.Conn != nil {
delete(Ctx.TunnelInstance, k)
delete(Ctx.PullTunnelInstance, k)
}
}
for k, tc := range Ctx.TunnelConfig {
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)
(*tc.Server).Close()
delete(Ctx.TunnelConfig, k)
delete(Ctx.PullTunnelConfig, k)
}
}
c.conn.Close()
Expand Down
39 changes: 19 additions & 20 deletions lib/util/message/message.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,26 +9,25 @@ type MessageType int
const (
// Platypus <-> Termite
STDIO MessageType = iota
TUNNEL_DATA
PULL_TUNNEL_DATA

// Platypus -> Termite
WINDOW_SIZE
GET_CLIENT_INFO
DUPLICATED_CLIENT
PROCESS_START
PROCESS_TERMINATE
TUNNEL_CREATE
TUNNEL_DELETE
TUNNEL_CONNECT
TUNNEL_DISCONNECT
Pull_TUNNEL_CREATE
PULL_TUNNEL_CONNECT
PULL_TUNNEL_DISCONNECT

// Termite -> Platypus
PROCESS_STARTED
PROCESS_STOPED
CLIENT_INFO
TUNNEL_CONNECTED
TUNNEL_CONNECT_FAILED
TUNNEL_DISCONNECTED
PULL_TUNNEL_CONNECTED
PULL_TUNNEL_CONNECT_FAILED
PULL_TUNNEL_DISCONNECTED
)

type Message struct {
Expand Down Expand Up @@ -80,29 +79,29 @@ type BodyTerminateProcess struct {
Key string
}

type BodyTunnelConnect struct {
type BodyPullTunnelConnect struct {
Token string
Address string
}

type BodyTunnelConnected struct {
type BodyPullTunnelConnected struct {
Token string
}

type BodyTunnelConnectFailed struct {
type BodyPullTunnelConnectFailed struct {
Token string
Reason string
}

type BodyTunnelDisconnect struct {
type BodyPullTunnelDisconnect struct {
Token string
}

type BodyTunnelDisconnected struct {
type BodyPullTunnelDisconnected struct {
Token string
}

type BodyTunnelData struct {
type BodyPullTunnelData struct {
Token string
Data []byte
}
Expand All @@ -117,10 +116,10 @@ func RegisterGob() {
gob.Register(&BodyDuplicateClient{})
gob.Register(&BodyClientInfo{})
gob.Register(&BodyTerminateProcess{})
gob.Register(&BodyTunnelConnect{})
gob.Register(&BodyTunnelConnected{})
gob.Register(&BodyTunnelConnectFailed{})
gob.Register(&BodyTunnelDisconnect{})
gob.Register(&BodyTunnelDisconnected{})
gob.Register(&BodyTunnelData{})
gob.Register(&BodyPullTunnelConnect{})
gob.Register(&BodyPullTunnelConnected{})
gob.Register(&BodyPullTunnelConnectFailed{})
gob.Register(&BodyPullTunnelDisconnect{})
gob.Register(&BodyPullTunnelDisconnected{})
gob.Register(&BodyPullTunnelData{})
}

0 comments on commit 29e04ac

Please sign in to comment.