rproxy: added streamer object

app/rproxy/forwarder.go

@@ -66,7 +66,7 @@ func (forw *Forwarder) ListenTCP(wg *sync.WaitGroup) {
 			forw.log.Errorf("Listen err: %v", err)
 			return
 		}
-		go forw.handle(forw.ctx, conn)
+		go forw.handleTCP(forw.ctx, conn)
 	}
 }
 
@@ -74,22 +74,7 @@ func (forw *Forwarder) Stop() error {
 	return forw.listen.Close()
 }
 
-type Streamer struct {
-	source string
-	dest   string
-	ctx    context.Context
-	cancel context.CancelFunc
-}
-
-func NewStreamer(ctx context.Context) *Streamer {
-	ctx, cancel := context.WithCancel(ctx)
-	return &Streamer{
-		ctx:    ctx,
-		cancel: cancel,
-	}
-}
-
-func (forw *Forwarder) handle(ctx context.Context, inconn net.Conn) {
+func (forw *Forwarder) handleTCP(ctx context.Context, inconn net.Conn) {
 	forw.log.Debugf("%s: Handle on %d started", inconn.RemoteAddr(), forw.Lport)
 	defer inconn.Close()
 	if len(forw.Dests) == 0 {
@@ -97,29 +82,78 @@ func (forw *Forwarder) handle(ctx context.Context, inconn net.Conn) {
 	}
 	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)
+
+	/*
+		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 {
-		return
+		forw.log.Errorf("Handler on %d error: %v", forw.Lport, err)
 	}
-	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()
 	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 := copy(forw.ctx, outconn, inconn)
+	_, 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"`
 }

app/rproxy/rproxy.go

@@ -3,7 +3,9 @@ package rproxy
 import (
 	"context"
 	"errors"
+	"fmt"
 	"io"
+	"strings"
 	"sync"
 
 	"helmet/app/logger"
@@ -11,7 +13,7 @@ import (
 
 type Proxy struct {
 	Forwarders []*Forwarder       `json:"forwarders"  yaml:"forwarders"`
-	ctx        context.Context    `json:"-"           yaml:"-"`
+	ctx        context.Context    `jsreturn errors.New("Zero dport")on:"-"           yaml:"-"`
 	cancel     context.CancelFunc `json:"-"           yaml:"-"`
 	wg         sync.WaitGroup     `json:"-"           yaml:"-"`
 	log        *logger.Logger     `json:"-"           yaml:"-"`
@@ -35,6 +37,13 @@ func (bal *Proxy) AddForwarder(ctx context.Context, typ string, lport, dport uin
 	if dport == 0 {
 		return errors.New("Zero dport")
 	}
+	if typ == "" {
+		return errors.New("Empty type")
+	}
+	typ = strings.ToLower(typ)
+	if typ != TCP {
+		return fmt.Errorf("Unknown type %s", typ)
+	}
 	forw, err := NewForwarder(ctx, typ, lport, dport, addrs...)
 	if err != nil {
 		return err
@@ -92,7 +101,7 @@ func (bal *Proxy) Stop() error {
 	return err
 }
 
-func copy(ctx context.Context, writer io.Writer, reader io.Reader) (int64, error) {
+func iocopy(ctx context.Context, writer io.Writer, reader io.Reader) (int64, error) {
 	var err error
 	var size int64
 	var halt bool