[ Cmii ] [ Octopus ] - add tcp portforwarding part
This commit is contained in:
zeaslity
215
port_forwarding/udp.go
Normal file
215
port_forwarding/udp.go
Normal file
@@ -0,0 +1,215 @@
|
||||
/**
|
||||
* Filename: udp.go
|
||||
* Description: the PortForward udp layer implement.
|
||||
* Author: knownsec404
|
||||
* Time: 2020.09.23
|
||||
*/
|
||||
|
||||
package port_forwarding
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net"
|
||||
"time"
|
||||
)
|
||||
|
||||
// as UDP client Conn
|
||||
type UDPDistribute struct {
|
||||
Established bool
|
||||
Conn *(net.UDPConn)
|
||||
RAddr net.Addr
|
||||
Cache chan []byte
|
||||
}
|
||||
|
||||
/**********************************************************************
|
||||
* @Function: NewUDPDistribute(conn *(net.UDPConn), addr net.Addr) (*UDPDistribute)
|
||||
* @Description: initialize UDPDistribute structure (as UDP client Conn)
|
||||
* @Parameter: conn *(net.UDPConn), the udp connection object
|
||||
* @Parameter: addr net.Addr, the udp client remote adddress
|
||||
* @Return: *UDPDistribute, the new UDPDistribute structure pointer
|
||||
**********************************************************************/
|
||||
func NewUDPDistribute(conn *(net.UDPConn), addr net.Addr) *UDPDistribute {
|
||||
return &UDPDistribute{
|
||||
Established: true,
|
||||
Conn: conn,
|
||||
RAddr: addr,
|
||||
Cache: make(chan []byte, 16),
|
||||
}
|
||||
}
|
||||
|
||||
/**********************************************************************
|
||||
* @Function: (this *UDPDistribute) Close() (error)
|
||||
* @Description: set "Established" flag is false, the UDP service will
|
||||
* cleaned up according to certain conditions.
|
||||
* @Parameter: nil
|
||||
* @Return: error, the error
|
||||
**********************************************************************/
|
||||
func (this *UDPDistribute) Close() error {
|
||||
this.Established = false
|
||||
return nil
|
||||
}
|
||||
|
||||
/**********************************************************************
|
||||
* @Function: (this *UDPDistribute) Read(b []byte) (n int, err error)
|
||||
* @Description: read data from connection, due to the udp implementation of
|
||||
* PortForward, read here will only produce a timeout error and closed error
|
||||
* (compared to the normal net.Conn object)
|
||||
* @Parameter: b []byte, the buffer for receive data
|
||||
* @Return: (n int, err error), the length of the data read and error
|
||||
**********************************************************************/
|
||||
func (this *UDPDistribute) Read(b []byte) (n int, err error) {
|
||||
if !this.Established {
|
||||
return 0, errors.New("udp distrubute has closed")
|
||||
}
|
||||
|
||||
select {
|
||||
case <-time.After(16 * time.Second):
|
||||
return 0, errors.New("udp distrubute read timeout")
|
||||
case data := <-this.Cache:
|
||||
n := len(data)
|
||||
copy(b, data)
|
||||
return n, nil
|
||||
}
|
||||
}
|
||||
|
||||
/**********************************************************************
|
||||
* @Function: (this *UDPDistribute) Write(b []byte) (n int, err error)
|
||||
* @Description: write data to connection by "WriteTo()"
|
||||
* @Parameter: b []byte, the data to be sent
|
||||
* @Return: (n int, err error), the length of the data write and error
|
||||
**********************************************************************/
|
||||
func (this *UDPDistribute) Write(b []byte) (n int, err error) {
|
||||
if !this.Established {
|
||||
return 0, errors.New("udp distrubute has closed")
|
||||
}
|
||||
return this.Conn.WriteTo(b, this.RAddr)
|
||||
}
|
||||
|
||||
/**********************************************************************
|
||||
* @Function: (this *UDPDistribute) RemoteAddr() (net.Addr)
|
||||
* @Description: get remote address
|
||||
* @Parameter: nil
|
||||
* @Return: net.Addr, the remote address
|
||||
**********************************************************************/
|
||||
func (this *UDPDistribute) RemoteAddr() net.Addr {
|
||||
return this.RAddr
|
||||
}
|
||||
|
||||
/**********************************************************************
|
||||
* @Function: ListenUDP(address string, clientc chan Conn, quit chan bool)
|
||||
* @Description: listen local udp service, and accept client connection,
|
||||
* initialize connection and return by channel.
|
||||
* since udp is running as a service, it only obtains remote data through
|
||||
* the Read* function cluster (different from tcp), so we need a temporary
|
||||
* table to record, so that we can use the temporary table to determine
|
||||
* whether to forward or create a new link
|
||||
* @Parameter: address string, the local listen address
|
||||
* @Parameter: clientc chan Conn, new client connection channel
|
||||
* @Parameter: quit chan bool, the quit signal channel
|
||||
* @Return: nil
|
||||
**********************************************************************/
|
||||
func ListenUDP(address string, clientc chan Conn, quit chan bool) {
|
||||
addr, err := net.ResolveUDPAddr("udp", address)
|
||||
if err != nil {
|
||||
LogError("udp listen error, %s", err)
|
||||
clientc <- nil
|
||||
return
|
||||
}
|
||||
serv, err := net.ListenUDP("udp", addr)
|
||||
if err != nil {
|
||||
LogError("udp listen error, %s", err)
|
||||
clientc <- nil
|
||||
return
|
||||
}
|
||||
defer serv.Close()
|
||||
|
||||
// the udp distrubute table
|
||||
table := make(map[string]*UDPDistribute)
|
||||
// NOTICE:
|
||||
// in the process of running, the table will generate invalid historical
|
||||
// data, we have not cleaned it up(These invalid historical data will not
|
||||
// affect our normal operation).
|
||||
//
|
||||
// if you want to deal with invalid historical data, the best way is to
|
||||
// launch a new coroutine to handle net.Read, and another coroutine to
|
||||
// handle the connection status, but additional logic to exit the
|
||||
// coroutine is needed. Currently we think it is unnecessary.
|
||||
//
|
||||
// under the current code logic: if a udp connection exits, it will timeout
|
||||
// to trigger "Close()", and finally set "Established" to false, but there
|
||||
// is no logic to check this state to clean up, so invalid historical data
|
||||
// is generated (of course, we can traverse to clean up? every packet? no)
|
||||
//
|
||||
// PS: you can see that we called "delete()" in the following logic, but
|
||||
// this is only for processing. when a new udp client hits invalid
|
||||
// historical data, we need to return a new connection so that PortForward
|
||||
// can create a communication link.
|
||||
|
||||
for {
|
||||
// check quit
|
||||
select {
|
||||
case <-quit:
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
// set timeout, for check "quit" signal
|
||||
serv.SetDeadline(time.Now().Add(16 * time.Second))
|
||||
|
||||
// just new 32*1024, in the outer layer we used "io.Copy()", which
|
||||
// can only handle the size of 32*1024
|
||||
buf := make([]byte, 32*1024)
|
||||
n, addr, err := serv.ReadFrom(buf)
|
||||
if err != nil {
|
||||
if err, ok := err.(net.Error); ok && err.Timeout() {
|
||||
continue
|
||||
}
|
||||
LogError("udp listen error, %s", err)
|
||||
clientc <- nil
|
||||
return
|
||||
}
|
||||
buf = buf[:n]
|
||||
|
||||
// if the address in table, we distrubute message
|
||||
if d, ok := table[addr.String()]; ok {
|
||||
if d.Established {
|
||||
// it is established, distrubute message
|
||||
d.Cache <- buf
|
||||
continue
|
||||
} else {
|
||||
// we remove it when the connnection has expired
|
||||
delete(table, addr.String())
|
||||
}
|
||||
}
|
||||
// if the address not in table, we create new connection object
|
||||
conn := NewUDPDistribute(serv, addr)
|
||||
table[addr.String()] = conn
|
||||
conn.Cache <- buf
|
||||
clientc <- conn
|
||||
} // end for
|
||||
}
|
||||
|
||||
/**********************************************************************
|
||||
* @Function: ConnUDP(address string) (Conn, error)
|
||||
* @Description: dial to remote server, and return udp connection
|
||||
* @Parameter: address string, the remote server address that needs to be dialed
|
||||
* @Return: (Conn, error), the udp connection and error
|
||||
**********************************************************************/
|
||||
func ConnUDP(address string) (Conn, error) {
|
||||
conn, err := net.DialTimeout("udp", address, 10*time.Second)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// send one byte(knock) to server, get "established" udp connection
|
||||
_, err = conn.Write([]byte("\x00"))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// due to the characteristics of udp, when the udp server exits, we will
|
||||
// not receive any signal, it will be blocked at conn.Read();
|
||||
// here we set a timeout for udp
|
||||
conn.SetDeadline(time.Now().Add(60 * time.Second))
|
||||
return conn, nil
|
||||
}
|
||||
Reference in New Issue
Block a user