/**
* Filename: udp.go
* Description: the PortForward udp layer implement.
* Author: knownsec404
* Time: 2020.09.23
 */

package main

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
}