added unlinked controller; change create forwarded to createOrUpdate; added global proxy mutex
This commit is contained in:
+143
-87
@@ -2,11 +2,12 @@ package rproxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"fmt"
|
||||
"math/rand"
|
||||
"net"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"helmet/app/logger"
|
||||
)
|
||||
@@ -14,20 +15,27 @@ import (
|
||||
const (
|
||||
TCP = "tcp"
|
||||
UDP = "udp"
|
||||
|
||||
ForwStarted = "started"
|
||||
ForwStopped = "stopped"
|
||||
)
|
||||
|
||||
type Forwarder struct {
|
||||
Type string `json:"type" yaml:"type"`
|
||||
listen net.Listener `json:"-" yaml:"-"`
|
||||
ctx context.Context `json:"-" yaml:"-"`
|
||||
cancel context.CancelFunc `json:"-" yaml:"-"`
|
||||
Lport uint32 `json:"lport" yaml:"lport"`
|
||||
Dport uint32 `json:"dport" yaml:"dport"`
|
||||
Dests []*Destination `json:"dests" yaml:"dests"`
|
||||
log *logger.Logger
|
||||
State string `json:"state" yaml:"state"`
|
||||
Type string `json:"type" yaml:"type"`
|
||||
Lport uint32 `json:"lport" yaml:"lport"`
|
||||
Dport uint32 `json:"dport" yaml:"dport"`
|
||||
Dests []*Destination `json:"dests" yaml:"dests"`
|
||||
|
||||
listenTCP *net.TCPListener `json:"-" yaml:"-"`
|
||||
listenUDP *net.UDPConn `json:"-" yaml:"-"`
|
||||
ctx context.Context `json:"-" yaml:"-"`
|
||||
cancel context.CancelFunc `json:"-" yaml:"-"`
|
||||
log *logger.Logger
|
||||
}
|
||||
|
||||
func NewForwarder(ctx context.Context, typ string, lport, dport uint32, addrs ...string) (*Forwarder, error) {
|
||||
func NewForwarder(ctx context.Context, proto string, lport, dport uint32, addrs ...string) (*Forwarder, error) {
|
||||
var err error
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
forw := &Forwarder{
|
||||
Dests: make([]*Destination, 0),
|
||||
@@ -35,33 +43,67 @@ func NewForwarder(ctx context.Context, typ string, lport, dport uint32, addrs ..
|
||||
Dport: dport,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
Type: typ,
|
||||
Type: proto,
|
||||
}
|
||||
id := strconv.FormatUint(uint64(lport), 10)
|
||||
id := forw.Type + strconv.FormatUint(uint64(lport), 10)
|
||||
forw.log = logger.NewLogger("forwarder:" + id)
|
||||
for _, addr := range addrs {
|
||||
dest := NewDestination(addr)
|
||||
forw.Dests = append(forw.Dests, dest)
|
||||
}
|
||||
|
||||
portinfo := ":" + strconv.FormatUint(uint64(forw.Lport), 10)
|
||||
laddr, err := net.ResolveTCPAddr("tcp", portinfo)
|
||||
if err != nil {
|
||||
switch proto {
|
||||
case TCP:
|
||||
portinfo := ":" + strconv.FormatUint(uint64(forw.Lport), 10)
|
||||
laddr, err := net.ResolveTCPAddr("tcp", portinfo)
|
||||
if err != nil {
|
||||
return forw, err
|
||||
}
|
||||
listen, err := net.ListenTCP("tcp", laddr)
|
||||
if err != nil {
|
||||
return forw, err
|
||||
}
|
||||
forw.listenTCP = listen
|
||||
case UDP:
|
||||
portinfo := ":" + strconv.FormatUint(uint64(forw.Lport), 10)
|
||||
laddr, err := net.ResolveUDPAddr("udp", portinfo)
|
||||
if err != nil {
|
||||
return forw, err
|
||||
}
|
||||
listen, err := net.ListenUDP("udp", laddr)
|
||||
if err != nil {
|
||||
return forw, err
|
||||
}
|
||||
forw.listenUDP = listen
|
||||
default:
|
||||
err = fmt.Errorf("Unknown net type: %s", proto)
|
||||
return forw, err
|
||||
}
|
||||
listen, err := net.ListenTCP("tcp", laddr)
|
||||
if err != nil {
|
||||
return forw, err
|
||||
}
|
||||
forw.listen = listen
|
||||
return forw, err
|
||||
}
|
||||
|
||||
func (forw *Forwarder) Listen(wg *sync.WaitGroup) {
|
||||
switch forw.Type {
|
||||
case TCP:
|
||||
if forw.listenTCP != nil {
|
||||
forw.ListenTCP(wg)
|
||||
}
|
||||
case UDP:
|
||||
if forw.listenUDP != nil {
|
||||
forw.ListenUDP(wg)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (forw *Forwarder) ListenTCP(wg *sync.WaitGroup) {
|
||||
forw.log.Debugf("Start listening on %d", forw.Lport)
|
||||
forw.log.Debugf("Start listening on %s:%d", forw.Type, forw.Lport)
|
||||
forw.State = ForwStarted
|
||||
defer wg.Done()
|
||||
stater := func() {
|
||||
forw.State = ForwStopped
|
||||
}
|
||||
defer stater()
|
||||
for {
|
||||
conn, err := forw.listen.Accept()
|
||||
conn, err := forw.listenTCP.Accept()
|
||||
if err != nil {
|
||||
forw.log.Errorf("Listen err: %v", err)
|
||||
return
|
||||
@@ -70,8 +112,84 @@ func (forw *Forwarder) ListenTCP(wg *sync.WaitGroup) {
|
||||
}
|
||||
}
|
||||
|
||||
func (forw *Forwarder) ListenUDP(wg *sync.WaitGroup) {
|
||||
forw.log.Debugf("Start listening on %s:%d", forw.Type, forw.Lport)
|
||||
forw.State = ForwStarted
|
||||
defer wg.Done()
|
||||
stater := func() {
|
||||
forw.State = ForwStopped
|
||||
}
|
||||
defer stater()
|
||||
for {
|
||||
buffer := make([]byte, 2048)
|
||||
size, srcAddr, err := forw.listenUDP.ReadFromUDP(buffer)
|
||||
if err != nil {
|
||||
forw.log.Errorf("Error reading: %v", err)
|
||||
continue
|
||||
}
|
||||
go forw.handleUDP(forw.listenUDP, srcAddr, buffer[:size])
|
||||
}
|
||||
}
|
||||
|
||||
func (forw *Forwarder) handleUDP(listConn *net.UDPConn, srcAddr *net.UDPAddr, data []byte) {
|
||||
forw.log.Debugf("Handle on %d started", forw.Lport)
|
||||
if len(forw.Dests) == 0 {
|
||||
return
|
||||
}
|
||||
// Select dest address
|
||||
addrnum := rand.Uint32() % uint32(len(forw.Dests))
|
||||
ipaddr := forw.Dests[addrnum].Address
|
||||
destInfo := ipaddr + ":" + strconv.FormatUint(uint64(forw.Dport), 10)
|
||||
|
||||
destAddr, err := net.ResolveUDPAddr("udp", destInfo)
|
||||
if err != nil {
|
||||
forw.log.Debugf("Error resolving server address: %v", err)
|
||||
return
|
||||
}
|
||||
// Write to destination
|
||||
destConn, err := net.DialUDP("udp", nil, destAddr)
|
||||
if err != nil {
|
||||
forw.log.Debugf("Error dialing: %v", err)
|
||||
return
|
||||
}
|
||||
defer destConn.Close()
|
||||
_, err = destConn.Write(data)
|
||||
if err != nil {
|
||||
forw.log.Debugf("Error sending message: %v", err)
|
||||
return
|
||||
}
|
||||
const deadlinePeriod = 5 * time.Second
|
||||
destConn.SetReadDeadline(time.Now().Add(deadlinePeriod))
|
||||
// Read from destination and resend to initiator
|
||||
const readCount = 1
|
||||
for i := 0; i < readCount; i++ {
|
||||
buffer := make([]byte, 1024*2)
|
||||
size, _, err := destConn.ReadFromUDP(buffer)
|
||||
if err != nil {
|
||||
forw.log.Debugf("Error reading response: %v", err)
|
||||
return
|
||||
}
|
||||
_, err = listConn.WriteToUDP(buffer[:size], srcAddr)
|
||||
if err != nil {
|
||||
forw.log.Errorf("Error writing to back: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (forw *Forwarder) Stop() error {
|
||||
return forw.listen.Close()
|
||||
var err error
|
||||
switch forw.Type {
|
||||
case TCP:
|
||||
if forw.listenTCP != nil {
|
||||
return forw.listenTCP.Close()
|
||||
}
|
||||
case UDP:
|
||||
if forw.listenUDP != nil {
|
||||
return forw.listenUDP.Close()
|
||||
}
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (forw *Forwarder) handleTCP(ctx context.Context, inconn net.Conn) {
|
||||
@@ -80,23 +198,10 @@ func (forw *Forwarder) handleTCP(ctx context.Context, inconn net.Conn) {
|
||||
if len(forw.Dests) == 0 {
|
||||
return
|
||||
}
|
||||
// Select dest address
|
||||
addrnum := rand.Uint32() % uint32(len(forw.Dests))
|
||||
ipaddr := forw.Dests[addrnum].Address
|
||||
|
||||
/*
|
||||
dstaddr := ipaddr + ":" + strconv.FormatUint(uint64(forw.Dport), 10)
|
||||
outconn, err := net.Dial("tcp", dstaddr)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer outconn.Close()
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go forw.stream(&wg, inconn, outconn)
|
||||
wg.Add(1)
|
||||
go forw.stream(&wg, outconn, inconn)
|
||||
wg.Wait()
|
||||
*/
|
||||
str := NewStreamer(forw.ctx)
|
||||
err := str.Stream(inconn, ipaddr, forw.Dport)
|
||||
if err != nil {
|
||||
@@ -105,55 +210,6 @@ func (forw *Forwarder) handleTCP(ctx context.Context, inconn net.Conn) {
|
||||
forw.log.Debugf("Handler on %d stopped", forw.Lport)
|
||||
}
|
||||
|
||||
func (forw *Forwarder) stream(wg *sync.WaitGroup, inconn io.Reader, outconn io.Writer) {
|
||||
defer wg.Done()
|
||||
_, err := iocopy(forw.ctx, outconn, inconn)
|
||||
if err != nil {
|
||||
forw.log.Errorf("Copy err: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
type Streamer struct {
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
log *logger.Logger
|
||||
}
|
||||
|
||||
func NewStreamer(ctx context.Context) *Streamer {
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
log := logger.NewLogger("streamer")
|
||||
return &Streamer{
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
log: log,
|
||||
}
|
||||
}
|
||||
|
||||
func (str *Streamer) Stream(inconn net.Conn, dipaddr string, dport uint32) error {
|
||||
var err error
|
||||
dstaddr := dipaddr + ":" + strconv.FormatUint(uint64(dport), 10)
|
||||
outconn, err := net.Dial("tcp", dstaddr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer outconn.Close()
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go str.stream(&wg, inconn, outconn)
|
||||
wg.Add(1)
|
||||
go str.stream(&wg, outconn, inconn)
|
||||
wg.Wait()
|
||||
return err
|
||||
}
|
||||
|
||||
func (str *Streamer) stream(wg *sync.WaitGroup, inconn io.Reader, outconn io.Writer) {
|
||||
defer wg.Done()
|
||||
_, err := iocopy(str.ctx, outconn, inconn)
|
||||
if err != nil {
|
||||
str.log.Errorf("Copy err: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
type Destination struct {
|
||||
Address string `json:"address" yaml:"address"`
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user