[ Cmii ] [ Octopus ] - add tcp portforwarding part
This commit is contained in:
zeaslity
356
port_forwarding/forward.go
Normal file
356
port_forwarding/forward.go
Normal file
@@ -0,0 +1,356 @@
|
||||
/**
|
||||
* Filename: forward.go
|
||||
* Description: the PortForward main logic implement, Three working modes are
|
||||
* provided:
|
||||
* 1.Conn<=>Conn: dial A remote server, dial B remote server, connected
|
||||
* 2.Listen<=>Conn: listen local server, dial remote server, connected
|
||||
* 3.Listen<=>Listen: listen A local server, listen B local server, connected
|
||||
* Author: knownsec404
|
||||
* Time: 2020.09.23
|
||||
*/
|
||||
|
||||
package port_forwarding
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net"
|
||||
"time"
|
||||
)
|
||||
|
||||
const PORTFORWARD_PROTO_NIL uint8 = 0x00
|
||||
const PORTFORWARD_PROTO_TCP uint8 = 0x10
|
||||
const PORTFORWARD_PROTO_UDP uint8 = 0x20
|
||||
|
||||
const PORTFORWARD_SOCK_NIL uint8 = 0x00
|
||||
const PORTFORWARD_SOCK_LISTEN uint8 = 0x01
|
||||
const PORTFORWARD_SOCK_CONN uint8 = 0x02
|
||||
|
||||
// the PortForward network interface
|
||||
type Conn interface {
|
||||
// Read reads data from the connection.
|
||||
Read(b []byte) (n int, err error)
|
||||
// Write writes data to the connection.
|
||||
Write(b []byte) (n int, err error)
|
||||
// Close closes the connection.
|
||||
Close() error
|
||||
// RemoteAddr returns the remote network address.
|
||||
RemoteAddr() net.Addr
|
||||
}
|
||||
|
||||
// the PortForward launch arguemnt
|
||||
type Args struct {
|
||||
Protocol uint8
|
||||
// sock1
|
||||
Method1 uint8
|
||||
Addr1 string
|
||||
// sock2
|
||||
Method2 uint8
|
||||
Addr2 string
|
||||
}
|
||||
|
||||
var stop chan bool = nil
|
||||
|
||||
/**********************************************************************
|
||||
* @Function: Launch(args Args)
|
||||
* @Description: launch PortForward working mode by arguments
|
||||
* @Parameter: args Args, the launch arguments
|
||||
* @Return: nil
|
||||
**********************************************************************/
|
||||
func Launch(args Args) {
|
||||
// initialize stop channel, the maximum number of coroutines that
|
||||
// need to be managed is 2 (listen-listen)
|
||||
stop = make(chan bool, 3)
|
||||
|
||||
//
|
||||
if args.Method1 == PORTFORWARD_SOCK_CONN &&
|
||||
args.Method2 == PORTFORWARD_SOCK_CONN {
|
||||
// sock1 conn, sock2 conn
|
||||
ConnConn(args.Protocol, args.Addr1, args.Addr2)
|
||||
} else if args.Method1 == PORTFORWARD_SOCK_CONN &&
|
||||
args.Method2 == PORTFORWARD_SOCK_LISTEN {
|
||||
// sock1 conn, sock2 listen
|
||||
ListenConn(args.Protocol, args.Addr2, args.Addr1)
|
||||
} else if args.Method1 == PORTFORWARD_SOCK_LISTEN &&
|
||||
args.Method2 == PORTFORWARD_SOCK_CONN {
|
||||
// sock1 listen, sock2 conn
|
||||
ListenConn(args.Protocol, args.Addr1, args.Addr2)
|
||||
} else if args.Method1 == PORTFORWARD_SOCK_LISTEN &&
|
||||
args.Method2 == PORTFORWARD_SOCK_LISTEN {
|
||||
// sock1 listen , sock2 listen
|
||||
ListenListen(args.Protocol, args.Addr1, args.Addr2)
|
||||
} else {
|
||||
LogError("unknown forward method")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
/**********************************************************************
|
||||
* @Function: Shutdown()
|
||||
* @Description: the shutdown PortForward
|
||||
* @Parameter: nil
|
||||
* @Return: nil
|
||||
**********************************************************************/
|
||||
func Shutdown() {
|
||||
stop <- true
|
||||
}
|
||||
|
||||
/**********************************************************************
|
||||
* @Function: ListenConn(proto uint8, addr1 string, addr2 string)
|
||||
* @Description: "Listen<=>Conn" working mode
|
||||
* @Parameter: proto uint8, the tcp or udp protocol setting
|
||||
* @Parameter: addr1 string, the address1 "ip:port" string
|
||||
* @Parameter: addr2 string, the address2 "ip:port" string
|
||||
* @Return: nil
|
||||
**********************************************************************/
|
||||
func ListenConn(proto uint8, addr1 string, addr2 string) {
|
||||
// get sock launch function by protocol
|
||||
sockfoo1 := ListenTCP
|
||||
if proto == PORTFORWARD_PROTO_UDP {
|
||||
sockfoo1 = ListenUDP
|
||||
}
|
||||
sockfoo2 := ConnTCP
|
||||
if proto == PORTFORWARD_PROTO_UDP {
|
||||
sockfoo2 = ConnUDP
|
||||
}
|
||||
|
||||
// launch socket1 listen
|
||||
clientc := make(chan Conn)
|
||||
quit := make(chan bool, 1)
|
||||
LogInfo("listen A point with sock1 [%s]", addr1)
|
||||
go sockfoo1(addr1, clientc, quit)
|
||||
|
||||
var count int = 1
|
||||
for {
|
||||
// socket1 listen & quit signal
|
||||
var sock1 Conn = nil
|
||||
select {
|
||||
case <-stop:
|
||||
quit <- true
|
||||
return
|
||||
case sock1 = <-clientc:
|
||||
if sock1 == nil {
|
||||
// set stop flag when error happend
|
||||
stop <- true
|
||||
continue
|
||||
}
|
||||
}
|
||||
LogInfo("A point(link%d) [%s] is ready", count, sock1.RemoteAddr())
|
||||
|
||||
// socket2 dial
|
||||
LogInfo("dial B point with sock2 [%s]", addr2)
|
||||
sock2, err := sockfoo2(addr2)
|
||||
if err != nil {
|
||||
sock1.Close()
|
||||
LogError("%s", err)
|
||||
continue
|
||||
}
|
||||
LogInfo("B point(sock2) is ready")
|
||||
|
||||
// connect with sockets
|
||||
go ConnectSock(count, sock1, sock2)
|
||||
count += 1
|
||||
} // end for
|
||||
}
|
||||
|
||||
/**********************************************************************
|
||||
* @Function: ListenListen(proto uint8, addr1 string, addr2 string)
|
||||
* @Description: the "Listen<=>Listen" working mode
|
||||
* @Parameter: proto uint8, the tcp or udp protocol setting
|
||||
* @Parameter: addr1 string, the address1 "ip:port" string
|
||||
* @Parameter: addr2 string, the address2 "ip:port" string
|
||||
* @Return: nil
|
||||
**********************************************************************/
|
||||
func ListenListen(proto uint8, addr1 string, addr2 string) {
|
||||
release := func(s1 Conn, s2 Conn) {
|
||||
if s1 != nil {
|
||||
s1.Close()
|
||||
}
|
||||
if s2 != nil {
|
||||
s2.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// get sock launch function by protocol
|
||||
sockfoo := ListenTCP
|
||||
if proto == PORTFORWARD_PROTO_UDP {
|
||||
sockfoo = ListenUDP
|
||||
}
|
||||
|
||||
// launch socket1 listen
|
||||
clientc1 := make(chan Conn)
|
||||
quit1 := make(chan bool, 1)
|
||||
LogInfo("listen A point with sock1 [%s]", addr1)
|
||||
go sockfoo(addr1, clientc1, quit1)
|
||||
// launch socket2 listen
|
||||
clientc2 := make(chan Conn)
|
||||
quit2 := make(chan bool, 1)
|
||||
LogInfo("listen B point with sock2 [%s]", addr2)
|
||||
go sockfoo(addr2, clientc2, quit2)
|
||||
|
||||
var sock1 Conn = nil
|
||||
var sock2 Conn = nil
|
||||
var count int = 1
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
quit1 <- true
|
||||
quit2 <- true
|
||||
release(sock1, sock2)
|
||||
return
|
||||
case c1 := <-clientc1:
|
||||
if c1 == nil {
|
||||
// set stop flag when error happend
|
||||
stop <- true
|
||||
continue
|
||||
}
|
||||
// close the last pending sock1
|
||||
if sock1 != nil {
|
||||
sock1.Close()
|
||||
}
|
||||
sock1 = c1
|
||||
LogInfo("A point(link%d) [%s] is ready", count, sock1.RemoteAddr())
|
||||
case c2 := <-clientc2:
|
||||
if c2 == nil {
|
||||
// set stop flag when error happend
|
||||
stop <- true
|
||||
continue
|
||||
}
|
||||
// close the last pending sock2
|
||||
if sock2 != nil {
|
||||
sock2.Close()
|
||||
}
|
||||
sock2 = c2
|
||||
LogInfo("B point(link%d) [%s] is ready", count, sock2.RemoteAddr())
|
||||
case <-time.After(120 * time.Second):
|
||||
if sock1 != nil {
|
||||
LogWarn("A point(%s) socket wait timeout, reset", sock1.RemoteAddr())
|
||||
}
|
||||
if sock2 != nil {
|
||||
LogWarn("B point(%s) socket wait timeout, reset", sock2.RemoteAddr())
|
||||
}
|
||||
release(sock1, sock2)
|
||||
continue
|
||||
}
|
||||
|
||||
// wait another socket ready
|
||||
if sock1 == nil || sock2 == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// the two socket is ready, connect with sockets
|
||||
go ConnectSock(count, sock1, sock2)
|
||||
count += 1
|
||||
// reset sock1 & sock2
|
||||
sock1 = nil
|
||||
sock2 = nil
|
||||
} // end for
|
||||
}
|
||||
|
||||
/**********************************************************************
|
||||
* @Function: ConnConn(proto uint8, addr1 string, addr2 string)
|
||||
* @Description: the "Conn<=>Conn" working mode
|
||||
* @Parameter: proto uint8, the tcp or udp protocol setting
|
||||
* @Parameter: addr1 string, the address1 "ip:port" string
|
||||
* @Parameter: addr2 string, the address2 "ip:port" string
|
||||
* @Return: nil
|
||||
**********************************************************************/
|
||||
func ConnConn(proto uint8, addr1 string, addr2 string) {
|
||||
// get sock launch function by protocol
|
||||
sockfoo := ConnTCP
|
||||
if proto == PORTFORWARD_PROTO_UDP {
|
||||
sockfoo = ConnUDP
|
||||
}
|
||||
|
||||
var count int = 1
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
// socket1 dial
|
||||
LogInfo("dial A point with sock1 [%s]", addr1)
|
||||
sock1, err := sockfoo(addr1)
|
||||
if err != nil {
|
||||
LogError("%s", err)
|
||||
time.Sleep(16 * time.Second)
|
||||
continue
|
||||
}
|
||||
LogInfo("A point(sock1) is ready")
|
||||
|
||||
// waiting for the first message sent by the A point(sock1)
|
||||
buf := make([]byte, 32*1024)
|
||||
n, err := sock1.Read(buf)
|
||||
if err != nil {
|
||||
LogError("A point: %s", err)
|
||||
time.Sleep(16 * time.Second)
|
||||
continue
|
||||
}
|
||||
buf = buf[:n]
|
||||
|
||||
// socket2 dial
|
||||
LogInfo("dial B point with sock2 [%s]", addr2)
|
||||
sock2, err := sockfoo(addr2)
|
||||
if err != nil {
|
||||
sock1.Close()
|
||||
LogError("%s", err)
|
||||
time.Sleep(16 * time.Second)
|
||||
continue
|
||||
}
|
||||
LogInfo("B point(sock2) is ready")
|
||||
|
||||
// first pass in the first message above
|
||||
_, err = sock2.Write(buf)
|
||||
if err != nil {
|
||||
LogError("B point: %s", err)
|
||||
time.Sleep(16 * time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
// connect with sockets
|
||||
go ConnectSock(count, sock1, sock2)
|
||||
count += 1
|
||||
} // end for
|
||||
}
|
||||
|
||||
/**********************************************************************
|
||||
* @Function: ConnectSock(id int, sock1 Conn, sock2 Conn)
|
||||
* @Description: connect two sockets, if an error occurs, the socket will
|
||||
* be closed so that the coroutine can exit normally
|
||||
* @Parameter: id int, the communication link id
|
||||
* @Parameter: sock1 Conn, the first socket object
|
||||
* @Parameter: sock2 Conn, the second socket object
|
||||
* @Return: nil
|
||||
**********************************************************************/
|
||||
func ConnectSock(id int, sock1 Conn, sock2 Conn) {
|
||||
exit := make(chan bool, 1)
|
||||
|
||||
//
|
||||
go func() {
|
||||
_, err := io.Copy(sock1, sock2)
|
||||
if err != nil {
|
||||
LogError("ConnectSock%d(A=>B): %s", id, err)
|
||||
} else {
|
||||
LogInfo("ConnectSock%d(A=>B) exited", id)
|
||||
}
|
||||
exit <- true
|
||||
}()
|
||||
|
||||
//
|
||||
go func() {
|
||||
_, err := io.Copy(sock2, sock1)
|
||||
if err != nil {
|
||||
LogError("ConnectSock%d(B=>A): %s", id, err)
|
||||
} else {
|
||||
LogInfo("ConnectSock%d(B=>A) exited", id)
|
||||
}
|
||||
exit <- true
|
||||
}()
|
||||
|
||||
// exit when close either end
|
||||
<-exit
|
||||
// close all socket, so that "io.Copy" can exit
|
||||
sock1.Close()
|
||||
sock2.Close()
|
||||
}
|
||||
Reference in New Issue
Block a user