Zoraxy/mod/streamproxy/streamproxy.go

package streamproxy

import (
	"errors"
	"log"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/uuid"
	"imuslab.com/zoraxy/mod/database"
)

/*
	TCP Proxy

	Forward port from one port to another
	Also accept active connection and passive
	connection
*/

type ProxyRelayOptions struct {
	Name          string
	ListeningAddr string
	ProxyAddr     string
	Timeout       int
	UseTCP        bool
	UseUDP        bool
}

type ProxyRelayConfig struct {
	UUID                        string       //A UUIDv4 representing this config
	Name                        string       //Name of the config
	Running                     bool         //Status, read only
	AutoStart                   bool         //If the service suppose to started automatically
	ListeningAddress            string       //Listening Address, usually 127.0.0.1:port
	ProxyTargetAddr             string       //Proxy target address
	UseTCP                      bool         //Enable TCP proxy
	UseUDP                      bool         //Enable UDP proxy
	Timeout                     int          //Timeout for connection in sec
	tcpStopChan                 chan bool    //Stop channel for TCP listener
	udpStopChan                 chan bool    //Stop channel for UDP listener
	aTobAccumulatedByteTransfer atomic.Int64 //Accumulated byte transfer from A to B
	bToaAccumulatedByteTransfer atomic.Int64 //Accumulated byte transfer from B to A
	udpClientMap                sync.Map     //map storing the UDP client-server connections
	parent                      *Manager     `json:"-"`
}

type Options struct {
	Database             *database.Database
	DefaultTimeout       int
	AccessControlHandler func(net.Conn) bool
}

type Manager struct {
	//Config and stores
	Options *Options
	Configs []*ProxyRelayConfig

	//Realtime Statistics
	Connections int //currently connected connect counts

}

func NewStreamProxy(options *Options) *Manager {
	options.Database.NewTable("tcprox")

	//Load relay configs from db
	previousRules := []*ProxyRelayConfig{}
	if options.Database.KeyExists("tcprox", "rules") {
		options.Database.Read("tcprox", "rules", &previousRules)
	}

	//Check if the AccessControlHandler is empty. If yes, set it to always allow access
	if options.AccessControlHandler == nil {
		options.AccessControlHandler = func(conn net.Conn) bool {
			//Always allow access
			return true
		}
	}

	//Create a new proxy manager for TCP
	thisManager := Manager{
		Options:     options,
		Connections: 0,
	}

	//Inject manager into the rules
	for _, rule := range previousRules {
		rule.parent = &thisManager
		if rule.Running {
			//This was previously running. Start it again
			log.Println("[Stream Proxy] Resuming stream proxy rule " + rule.Name)
			rule.Start()
		}
	}

	thisManager.Configs = previousRules

	return &thisManager
}

func (m *Manager) NewConfig(config *ProxyRelayOptions) string {
	//Generate two zero value for atomic int64
	aAcc := atomic.Int64{}
	bAcc := atomic.Int64{}
	aAcc.Store(0)
	bAcc.Store(0)
	//Generate a new config from options
	configUUID := uuid.New().String()
	thisConfig := ProxyRelayConfig{
		UUID:                        configUUID,
		Name:                        config.Name,
		ListeningAddress:            config.ListeningAddr,
		ProxyTargetAddr:             config.ProxyAddr,
		UseTCP:                      config.UseTCP,
		UseUDP:                      config.UseUDP,
		Timeout:                     config.Timeout,
		tcpStopChan:                 nil,
		udpStopChan:                 nil,
		aTobAccumulatedByteTransfer: aAcc,
		bToaAccumulatedByteTransfer: bAcc,
		udpClientMap:                sync.Map{},
		parent:                      m,
	}
	m.Configs = append(m.Configs, &thisConfig)
	m.SaveConfigToDatabase()
	return configUUID
}

func (m *Manager) GetConfigByUUID(configUUID string) (*ProxyRelayConfig, error) {
	// Find and return the config with the specified UUID
	for _, config := range m.Configs {
		if config.UUID == configUUID {
			return config, nil
		}
	}
	return nil, errors.New("config not found")
}

// Edit the config based on config UUID, leave empty for unchange fields
func (m *Manager) EditConfig(configUUID string, newName string, newListeningAddr string, newProxyAddr string, useTCP bool, useUDP bool, newTimeout int) error {
	// Find the config with the specified UUID
	foundConfig, err := m.GetConfigByUUID(configUUID)
	if err != nil {
		return err
	}

	// Validate and update the fields
	if newName != "" {
		foundConfig.Name = newName
	}
	if newListeningAddr != "" {
		foundConfig.ListeningAddress = newListeningAddr
	}
	if newProxyAddr != "" {
		foundConfig.ProxyTargetAddr = newProxyAddr
	}

	foundConfig.UseTCP = useTCP
	foundConfig.UseUDP = useUDP

	if newTimeout != -1 {
		if newTimeout < 0 {
			return errors.New("invalid timeout value given")
		}
		foundConfig.Timeout = newTimeout
	}

	m.SaveConfigToDatabase()

	//Check if config is running. If yes, restart it
	if foundConfig.IsRunning() {
		foundConfig.Restart()
	}

	return nil
}

func (m *Manager) RemoveConfig(configUUID string) error {
	// Find and remove the config with the specified UUID
	for i, config := range m.Configs {
		if config.UUID == configUUID {
			m.Configs = append(m.Configs[:i], m.Configs[i+1:]...)
			m.SaveConfigToDatabase()
			return nil
		}
	}
	return errors.New("config not found")
}

func (m *Manager) SaveConfigToDatabase() {
	m.Options.Database.Write("tcprox", "rules", m.Configs)
}

/*
	Config Functions
*/

// Start a proxy if stopped
func (c *ProxyRelayConfig) Start() error {
	if c.IsRunning() {
		c.Running = true
		return errors.New("proxy already running")
	}

	// Create a stopChan to control the loop
	tcpStopChan := make(chan bool)
	udpStopChan := make(chan bool)

	//Start the proxy service
	if c.UseUDP {
		c.udpStopChan = udpStopChan
		go func() {
			err := c.ForwardUDP(c.ListeningAddress, c.ProxyTargetAddr, udpStopChan)
			if err != nil {
				if !c.UseTCP {
					c.Running = false
					c.parent.SaveConfigToDatabase()
				}
				log.Println("[TCP] Error starting stream proxy " + c.Name + "(" + c.UUID + "): " + err.Error())
			}
		}()
	}

	if c.UseTCP {
		c.tcpStopChan = tcpStopChan
		go func() {
			//Default to transport mode
			err := c.Port2host(c.ListeningAddress, c.ProxyTargetAddr, tcpStopChan)
			if err != nil {
				c.Running = false
				c.parent.SaveConfigToDatabase()
				log.Println("[TCP] Error starting stream proxy " + c.Name + "(" + c.UUID + "): " + err.Error())
			}
		}()
	}

	//Successfully spawned off the proxy routine
	c.Running = true
	c.parent.SaveConfigToDatabase()
	return nil
}

// Return if a proxy config is running
func (c *ProxyRelayConfig) IsRunning() bool {
	return c.tcpStopChan != nil || c.udpStopChan != nil
}

// Restart a proxy config
func (c *ProxyRelayConfig) Restart() {
	if c.IsRunning() {
		c.Stop()
	}
	time.Sleep(300 * time.Millisecond)
	c.Start()
}

// Stop a running proxy if running
func (c *ProxyRelayConfig) Stop() {
	log.Println("[STREAM PROXY] Stopping Stream Proxy " + c.Name)

	if c.udpStopChan != nil {
		log.Println("[STREAM PROXY] Stopping UDP for " + c.Name)
		c.udpStopChan <- true
		c.udpStopChan = nil
	}

	if c.tcpStopChan != nil {
		log.Println("[STREAM PROXY] Stopping TCP for " + c.Name)
		c.tcpStopChan <- true
		c.tcpStopChan = nil
	}

	log.Println("[STREAM PROXY] Stopped Stream Proxy " + c.Name)
	c.Running = false

	//Update the running status
	c.parent.SaveConfigToDatabase()
}