retryable-bouncer.go

package main

import (
	"flag"
	"io"
	"log"
	"net"
	"sync"
	"time"

	"github.com/Songmu/retry"
)

func main() {
	var (
		local  = flag.String("local", "local addr", ":9999")
		remote = flag.String("remote", "remote addr", "127.0.0.1:5000")
	)
	flag.Parse()
	l, err := net.Listen("tcp", *local)
	if err != nil {
		panic(err)
	}
	for {
		conn, err := l.Accept()
		if err != nil {
			log.Println("accept error", err)
			continue
		}
		go handleConn(conn, *remote)
	}
}

func handleConn(conn net.Conn, remote string) {
	defer conn.Close()
	var upstream net.Conn
	err := retry.Retry(10, time.Second, func() error {
		var err error
		log.Println("connect to upstream")
		upstream, err = net.DialTimeout("tcp", remote, time.Second)
		return err
	})
	if err != nil {
		log.Println("upstream error", err)
		return
	}
	log.Println("upstream connected")

	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		io.Copy(upstream, conn)
	}()
	go func() {
		defer wg.Done()
		io.Copy(conn, upstream)
	}()
	wg.Wait()
}