Skip to content

Commit ccd1927

Browse files
WangYihangWangYihang
committed
[*] #28 Fix bug in remote port frowarding
1 parent b667e8c commit ccd1927

File tree

4 files changed

+66
-61
lines changed

4 files changed

+66
-61
lines changed

‎lib/context/context.go

+1-4
Original file line numberDiff line numberDiff line change
@@ -197,16 +197,13 @@ func Shutdown() {
197197
}
198198

199199
func AddPushTunnelConfig(termite *TermiteClient, local_address string, remote_address string) {
200-
token := str.RandomString(0x10)
201-
202200
termite.AtomLock.Lock()
203201
defer func() { termite.AtomLock.Unlock() }()
204202

205203
termite.EncoderLock.Lock()
206204
err := termite.Encoder.Encode(message.Message{
207205
Type: message.PUSH_TUNNEL_CREATE,
208206
Body: message.BodyPushTunnelCreate{
209-
Token: token,
210207
Address: remote_address,
211208
},
212209
})
@@ -215,7 +212,7 @@ func AddPushTunnelConfig(termite *TermiteClient, local_address string, remote_ad
215212
if err != nil {
216213
log.Error(err.Error())
217214
} else {
218-
Ctx.PushTunnelConfig[token] = PushTunnelConfig{
215+
Ctx.PushTunnelConfig[remote_address] = PushTunnelConfig{
219216
Termite: termite,
220217
Address: local_address,
221218
}

‎lib/context/server.go

+28-35
Original file line numberDiff line numberDiff line change
@@ -507,7 +507,6 @@ func TermiteMessageDispatcher(client *TermiteClient) {
507507
var key string
508508
switch msg.Type {
509509
case message.STDIO:
510-
log.Debug("case message.STDIO")
511510
key = msg.Body.(*message.BodyStdio).Key
512511
if process, exists := client.Processes[key]; exists {
513512
if process.WebSocket != nil {
@@ -516,10 +515,9 @@ func TermiteMessageDispatcher(client *TermiteClient) {
516515
os.Stdout.Write(msg.Body.(*message.BodyStdio).Data)
517516
}
518517
} else {
519-
log.Error("No such key")
518+
log.Debug("No such key: %s", key)
520519
}
521520
case message.PROCESS_STARTED:
522-
log.Debug("case message.PROCESS_STARTED")
523521
key = msg.Body.(*message.BodyProcessStarted).Key
524522
if process, exists := client.Processes[key]; exists {
525523
process.Pid = msg.Body.(*message.BodyProcessStarted).Pid
@@ -529,10 +527,9 @@ func TermiteMessageDispatcher(client *TermiteClient) {
529527
client.CurrentProcessKey = key
530528
}
531529
} else {
532-
log.Error("No such key")
530+
log.Debug("No such key: %s", key)
533531
}
534532
case message.PROCESS_STOPED:
535-
log.Debug("case message.PROCESS_STOPED")
536533
key = msg.Body.(*message.BodyProcessStoped).Key
537534
if process, exists := client.Processes[key]; exists {
538535
code := msg.Body.(*message.BodyProcessStoped).Code
@@ -543,10 +540,9 @@ func TermiteMessageDispatcher(client *TermiteClient) {
543540
client.CurrentProcessKey = ""
544541
}
545542
} else {
546-
log.Error("No such key")
543+
log.Debug("No such key: %s", key)
547544
}
548545
case message.PULL_TUNNEL_CONNECTED:
549-
log.Debug("case message.PULL_TUNNEL_CONNECTED")
550546
token := msg.Body.(*message.BodyPullTunnelConnected).Token
551547
log.Success("Tunnel (%s) connected", token)
552548
if ti, exists := Ctx.PullTunnelInstance[token]; exists {
@@ -574,44 +570,43 @@ func TermiteMessageDispatcher(client *TermiteClient) {
574570
}
575571
}()
576572
} else {
577-
log.Error("No such connection")
573+
log.Debug("No such connection")
578574
}
579575
case message.PULL_TUNNEL_CONNECT_FAILED:
580-
log.Debug("case message.PULL_TUNNEL_CONNECT_FAILED")
581576
token := msg.Body.(*message.BodyPullTunnelConnectFailed).Token
582577
reason := msg.Body.(*message.BodyPullTunnelConnectFailed).Reason
583578
if ti, exists := Ctx.PullTunnelInstance[token]; exists {
584579
log.Error("Connecting to %s failed: %s", token, reason)
585580
(*ti.Conn).Close()
586581
delete(Ctx.PullTunnelInstance, token)
587582
} else {
588-
log.Error("No such connection")
583+
log.Debug("No such connection")
589584
}
590585
case message.PULL_TUNNEL_DISCONNECTED:
591-
log.Debug("case message.PULL_TUNNEL_DISCONNECTED")
592586
token := msg.Body.(*message.BodyPullTunnelDisconnected).Token
593587
if ti, exists := Ctx.PullTunnelInstance[token]; exists {
594588
log.Error("%s disconnected", token)
595589
(*ti.Conn).Close()
596590
delete(Ctx.PullTunnelInstance, token)
597591
} else {
598-
log.Error("No such connection")
592+
log.Debug("No such connection")
599593
}
600594
case message.PULL_TUNNEL_DATA:
601-
log.Debug("case message.PULL_TUNNEL_DATA")
602595
token := msg.Body.(*message.BodyPullTunnelData).Token
603596
data := msg.Body.(*message.BodyPullTunnelData).Data
604597
if ti, exists := Ctx.PullTunnelInstance[token]; exists {
605598
(*ti.Conn).Write(data)
606599
} else {
607-
log.Error("No such connection")
600+
log.Debug("No such connection")
608601
}
609602
case message.PUSH_TUNNEL_CONNECT:
610-
log.Debug("case message.PUSH_TUNNEL_CONNECT")
611603
token := msg.Body.(*message.BodyPushTunnelConnect).Token
612-
if tc, exists := Ctx.PushTunnelConfig[token]; exists {
604+
address := msg.Body.(*message.BodyPushTunnelConnect).Address
605+
if tc, exists := Ctx.PushTunnelConfig[address]; exists {
606+
log.Info("Connecting to %s", tc.Address)
613607
conn, err := net.Dial("tcp", tc.Address)
614608
if err != nil {
609+
log.Error("Connecting to %s failed: %s", tc.Address, err.Error())
615610
tc.Termite.Encoder.Encode(message.Message{
616611
Type: message.PUSH_TUNNEL_CONNECT_FAILED,
617612
Body: message.BodyPushTunnelConnectFailed{
@@ -620,6 +615,7 @@ func TermiteMessageDispatcher(client *TermiteClient) {
620615
},
621616
})
622617
} else {
618+
log.Success("Connecting to %s succeed", tc.Address)
623619
Ctx.PushTunnelInstance[token] = PushTunnelInstance{
624620
Termite: tc.Termite,
625621
Conn: &conn,
@@ -635,9 +631,10 @@ func TermiteMessageDispatcher(client *TermiteClient) {
635631
buffer := make([]byte, 0x400)
636632
n, err := conn.Read(buffer)
637633
if err != nil {
634+
log.Debug("Reading from %s failed: %s", tc.Address, err.Error())
638635
tc.Termite.Encoder.Encode(message.Message{
639-
Type: message.PUSH_TUNNEL_CONNECT_FAILED,
640-
Body: message.BodyPushTunnelConnectFailed{
636+
Type: message.PUSH_TUNNEL_DISCONNECTED,
637+
Body: message.BodyPushTunnelDisonnected{
641638
Token: token,
642639
Reason: err.Error(),
643640
},
@@ -646,6 +643,7 @@ func TermiteMessageDispatcher(client *TermiteClient) {
646643
delete(Ctx.PushTunnelInstance, token)
647644
break
648645
} else {
646+
log.Debug("%d bytes read from %s", n, tc.Address)
649647
tc.Termite.Encoder.Encode(message.Message{
650648
Type: message.PUSH_TUNNEL_DATA,
651649
Body: message.BodyPushTunnelData{
@@ -658,53 +656,48 @@ func TermiteMessageDispatcher(client *TermiteClient) {
658656
}()
659657
}
660658
} else {
661-
log.Error("No such tunnel")
659+
log.Debug("No such tunnel: %s", token)
662660
}
663661
case message.PUSH_TUNNEL_DISCONNECT:
664-
log.Debug("case message.PUSH_TUNNEL_DISCONNECT")
665662
token := msg.Body.(*message.BodyPushTunnelDisonnect).Token
666663
if ti, exists := Ctx.PushTunnelInstance[token]; exists {
667664
log.Success("Tunnel %v closed", (*ti.Conn).RemoteAddr().String())
668665
(*ti.Conn).Close()
666+
delete(Ctx.PushTunnelInstance, token)
669667
} else {
670-
log.Error("No such tunnel")
668+
log.Debug("No such tunnel: %s", token)
671669
}
672670
case message.PUSH_TUNNEL_DISCONNECTED:
673-
log.Debug("case message.PUSH_TUNNEL_DISCONNECTED")
674671
token := msg.Body.(*message.BodyPushTunnelDisonnected).Token
675672
if ti, exists := Ctx.PushTunnelInstance[token]; exists {
676673
(*ti.Conn).Close()
677674
delete(Ctx.PushTunnelInstance, token)
678675
} else {
679-
log.Error("No such tunnel")
676+
log.Debug("No such tunnel: %s", token)
680677
}
681678
case message.PUSH_TUNNEL_CREATED:
682-
log.Debug("case message.PUSH_TUNNEL_CREATED")
683-
token := msg.Body.(*message.BodyPushTunnelCreated).Token
684-
if tc, exists := Ctx.PushTunnelConfig[token]; exists {
679+
address := msg.Body.(*message.BodyPushTunnelCreated).Address
680+
if tc, exists := Ctx.PushTunnelConfig[address]; exists {
685681
log.Success("Tunnel created: %s", tc.Address)
686682
} else {
687-
log.Error("No such tunnel")
683+
log.Debug("No such tunnel: %s", address)
688684
}
689685
case message.PUSH_TUNNEL_CREATE_FAILED:
690-
log.Debug("case message.PUSH_TUNNEL_CREATE_FAILED")
691-
token := msg.Body.(*message.BodyPushTunnelCreateFailed).Token
686+
address := msg.Body.(*message.BodyPushTunnelCreateFailed).Address
692687
reason := msg.Body.(*message.BodyPushTunnelCreateFailed).Reason
693-
if tc, exists := Ctx.PushTunnelConfig[token]; exists {
688+
if tc, exists := Ctx.PushTunnelConfig[address]; exists {
694689
log.Success("Tunnel create failed: %s", tc.Address, reason)
690+
delete(Ctx.PushTunnelConfig, address)
695691
} else {
696-
log.Error("No such tunnel")
692+
log.Debug("No such tunnel: %s", address)
697693
}
698694
case message.PUSH_TUNNEL_DELETED:
699-
log.Debug("case message.PUSH_TUNNEL_DELETED")
700695
// TODO
701696
log.Info("PUSH_TUNNEL_DELETED")
702697
case message.PUSH_TUNNEL_DELETE_FAILED:
703-
log.Debug("case message.PUSH_TUNNEL_DELETE_FAILED")
704698
// TODO
705699
log.Info("PUSH_TUNNEL_DELETE_FAILED")
706700
case message.PUSH_TUNNEL_DATA:
707-
log.Debug("case message.PUSH_TUNNEL_DATA")
708701
token := msg.Body.(*message.BodyPushTunnelData).Token
709702
data := msg.Body.(*message.BodyPushTunnelData).Data
710703
if ti, exists := Ctx.PushTunnelInstance[token]; exists {
@@ -721,7 +714,7 @@ func TermiteMessageDispatcher(client *TermiteClient) {
721714
delete(Ctx.PushTunnelInstance, token)
722715
}
723716
} else {
724-
log.Error("No such tunnel")
717+
log.Debug("No such tunnel: %s", token)
725718
}
726719
}
727720
}

‎lib/util/message/message.go

+5-5
Original file line numberDiff line numberDiff line change
@@ -124,17 +124,16 @@ type BodyPushTunnelData struct {
124124
}
125125

126126
type BodyPushTunnelCreate struct {
127-
Token string
128127
Address string
129128
}
130129

131130
type BodyPushTunnelCreated struct {
132-
Token string
131+
Address string
133132
}
134133

135134
type BodyPushTunnelCreateFailed struct {
136-
Token string
137-
Reason string
135+
Address string
136+
Reason string
138137
}
139138

140139
type BodyPushTunnelDelete struct {
@@ -151,7 +150,8 @@ type BodyPushTunnelDeleteFailed struct {
151150
}
152151

153152
type BodyPushTunnelConnect struct {
154-
Token string
153+
Token string
154+
Address string
155155
}
156156

157157
type BodyPushTunnelConnected struct {

0 commit comments

Comments
 (0)