forward.go

/**
* 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 main

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()
}