TCP proxy

tcpproxy.go

package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
)

var flagLocalAddr *string = flag.String("l", "localhost:9999", "local address")
var flagRemoteAddr *string = flag.String("r", "localhost:80", "remote address")

var remoteAddress, localAddress *net.TCPAddr

func init() {
	var err error

	flag.Parse()

	localAddress, err = net.ResolveTCPAddr("tcp", *flagLocalAddr)
	if err != nil {
		panic(err)
	}

	remoteAddress, err = net.ResolveTCPAddr("tcp", *flagRemoteAddr)
	if err != nil {
		panic(err)
	}
}

func newProxyConnection(ctx context.Context, cancelFunc context.CancelFunc, r io.Reader, w io.Writer) {
	dataPipe := make(chan []byte)
	maxBytes := 0
	go func() {
		for {
			select {
			case <-ctx.Done():
				log.Print("done, closing reader")
				return
			default:
				data := make([]byte, 256)
				n, err := r.Read(data)
				if err != nil {
					log.Printf("aborting connection - error reading: %v\n", err)
					cancelFunc()
				}
				if n > 0 {
					if n > maxBytes {
						maxBytes = n
						log.Printf("max %v bytes", maxBytes)
					}

					dataPipe <- data[:n]
				}
			}
		}
	}()

	go func() {
		for {
			select {
			case <-ctx.Done():
				log.Print("done, closing writer")
				return
			case data := <-dataPipe:
				_, err := w.Write(data)
				if err != nil {
					log.Printf("aborting connection - error writing: %v", err)
					cancelFunc()
				}
			}
		}
	}()
}

func handleConn(ctx context.Context, newConnections <-chan *net.TCPConn) {
	for clientConnection := range newConnections {
		log.Printf("starting new proxy connection for: %v", clientConnection.RemoteAddr().String())
		targetConnection, err := net.DialTCP("tcp", nil, remoteAddress)
		if err != nil {
			panic(err)
		}
		newProxy(ctx, clientConnection, targetConnection)
	}
}

func newProxy(ctx context.Context, clientConnection *net.TCPConn, targetConnection *net.TCPConn) {
	proxyCtx, cancelFunc := context.WithCancel(ctx)
	newProxyConnection(proxyCtx, cancelFunc, clientConnection, targetConnection)
	newProxyConnection(proxyCtx, cancelFunc, targetConnection, clientConnection)
}

func main() {
	ctx := newSignalContextWithSignals(os.Interrupt)

	fmt.Printf("Listening: %v\nProxying: %v\n\n", localAddress.String(), remoteAddress.String())

	listener, err := net.ListenTCP("tcp", localAddress)
	if err != nil {
		panic(err)
	}

	newConnections := make(chan *net.TCPConn)

	go handleConn(ctx, newConnections)

	go func() {
		for {
			log.Print("waiting for new connection")
			conn, err := listener.AcceptTCP()
			if err != nil {
				panic(err)
			}
			log.Print("accepted new connection")
			newConnections <- conn

		}
	}()

	<-ctx.Done()
	log.Printf("stopping listener: %v", listener.Close())
}

func newSignalContextWithSignals(signals ...os.Signal) context.Context {

	ctx, cancelFunc := context.WithCancel(context.Background())

	go func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, signals...)

		for {
			select {
			case <-c:
				cancelFunc()
				return
			}
		}
	}()

	return ctx
}