proxytest.go

package main

import (
	"flag"
	log "github.com/Sirupsen/logrus"
	"io"
	"net"
)

var (
	localAddr = flag.String("local_server_addr", "localhost:9999", "local address")
	remoteAddr = flag.String("remote_server_addr", "localhost:10000", "remote address")
	debug     = flag.Bool("debug", false, "Turns on debug mode")
)

func proxyConn(conn *net.TCPConn, remoteConn *net.Conn) {
	errc := make(chan error, 2)
	copy := func(dst io.Writer, src io.Reader) {
		log.Debugf("Sending data from src %v to dst %v", src, dst)
		_, err := io.Copy(dst, src)
		log.Debugf("Error: %v", err)
		errc <- err
	}
	go copy(conn, *remoteConn)
	go copy(*remoteConn, conn)
	<-errc
}

func handleConn(in <-chan *net.TCPConn, out chan<- *net.TCPConn, remoteConn *net.Conn) {
	for conn := range in {
		log.Printf("Handling connection %v...", conn)
		proxyConn(conn, remoteConn)
		out <- conn
	}
}

func closeConn(in <-chan *net.TCPConn) {
	for conn := range in {
		conn.Close()
	}
}

func main() {
	flag.Parse()
	if *debug {
		log.SetLevel(log.DebugLevel)
	}
	log.Infof("Listening on port: %v", *localAddr)
	addr, err := net.ResolveTCPAddr("tcp", *localAddr)
	if err != nil {
		log.Fatalf("Error resolving TCP address: %v", err)
	}
	listener, err := net.ListenTCP("tcp", addr)
	if err != nil {
		log.Fatalf("Error creating TCP listener: %v", err)
	}
	log.Infof("Connecting to remote server %v", *remoteAddr)
	remoteServiceConn, err := net.Dial("tcp", *remoteAddr)
	if err != nil {
		log.Fatalf("net.Dial(%v) = _, %v", *remoteAddr, err)
	}
	defer remoteServiceConn.Close()
	pending, complete := make(chan *net.TCPConn), make(chan *net.TCPConn)
	for i := 0; i < 5; i++ {
		go handleConn(pending, complete, &remoteServiceConn)
	}
	go closeConn(complete)
	log.Info("Waiting for incoming connections...")
	for {
		conn, err := listener.AcceptTCP()
		log.Infof("Received incoming connection! %v", conn)
		if err != nil {
			log.Fatalf("Error accepting TCP connetion: %v", err)
		}
		pending <- conn
	}
}