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"`
 }