tcppos

leakybuf.go

// Package leakybuf provides leaky buffer.
// It's based on the example in Effective Go.
package main

type LeakyBuf struct {
	bufSize  int // size of each buffer
	freeList chan []byte
}

// NewLeakyBuf creates a leaky buffer which can hold at most n buffer, each
// with bufSize bytes.
func NewLeakyBuf(n, bufSize int) *LeakyBuf {
	return &LeakyBuf{
		bufSize:  bufSize,
		freeList: make(chan []byte, n),
	}
}

// Get returns a buffer from the leaky buffer or create a new buffer.
func (lb *LeakyBuf) Get() (b []byte) {
	select {
	case b = <-lb.freeList:
	default:
		b = make([]byte, lb.bufSize)
	}
	return
}

// Put add the buffer into the free buffer pool for reuse. Panic if the buffer
// size is not the same with the leaky buffer's. This is intended to expose
// error usage of leaky buffer.
func (lb *LeakyBuf) Put(b []byte) {
	if len(b) != lb.bufSize {
		panic("invalid buffer size that's put into leaky buffer")
	}
	select {
	case lb.freeList <- b:
	default:
	}
	return
}

main.go

package main

import (
	"encoding/hex"
	"flag"
	"fmt"
	"log"
	"net"
	"time"
)

const bufSize = 4096
const nBuf = 2048

const (
	NO_TIMEOUT = iota
	SET_TIMEOUT
	TYPE_OUT
	TYPE_IN
)

var (
	port   = flag.String("port", "1234", "listen port, default 1234")
	target = flag.String("target", "", "target host and port. eg: 192.168.1.100:9001")
	t      = flag.Int("timeout", 60, "timeout (second), default 60s")
	debug  = flag.Bool("debug", false, "print debug message")

	timeout time.Duration

	leakybuf = NewLeakyBuf(nBuf, bufSize)
)

func main() {

	fmt.Println("\n  tcpproxy @author yinheli (me@yinheli.com), version: 1.0.0\n\n**********")

	flag.Parse()

	if *target == "" {
		flag.Usage()
		return
	}

	timeout = time.Duration(*t) * time.Second

	lis, err := net.Listen("tcp", "0.0.0.0:"+*port)
	if err != nil {
		panic(err.Error())
	}

	log.Printf("start proxy: listen:%v, target:%v, debug:%v", *port, *target, *debug)

	for {
		conn, err := lis.Accept()
		if err != nil {
			panic(fmt.Sprintf("accept error: %v", err.Error()))
			continue
		}

		go func() {
			defer func() {
				if x := recover(); x != nil {
					panic(fmt.Sprintf("handle exception: %v", x))
				}
			}()

			tconn, err := net.Dial("tcp", *target)
			if err != nil {
				panic(fmt.Sprintf("Dial target error: %v", err.Error()))
			}

			log.Printf("start: %s<->%s", conn.RemoteAddr(), *target)
			go pipeThenClose(conn, tconn, SET_TIMEOUT, TYPE_OUT)
			pipeThenClose(tconn, conn, NO_TIMEOUT, TYPE_IN)
		}()
	}

}

func setReadTimeout(c net.Conn) {
	if timeout != 0 {
		c.SetReadDeadline(time.Now().Add(timeout))
	}
}

func pipeThenClose(src, dst net.Conn, timeoutOpt int, dt int) {
	defer dst.Close()
	buf := leakybuf.Get()
	defer leakybuf.Put(buf)

	for {
		if timeoutOpt == SET_TIMEOUT {
			setReadTimeout(src)
		}
		n, err := src.Read(buf)
		// read may return EOF with n > 0
		// should always process n > 0 bytes before handling error
		if n > 0 {
			data := buf[0:n]

			if _, err = dst.Write(data); err != nil {
				break
			} else {
				if *debug {
					if dt == TYPE_OUT {
						log.Printf("send: %s <-> %s \n%v", src.RemoteAddr(), dst.RemoteAddr(), hex.Dump(data))
					} else {
						log.Printf("recive: %s <-> %s \n%v", src.RemoteAddr(), dst.RemoteAddr(), hex.Dump(data))
					}
				}
			}
		}
		if err != nil {
			break
		}
	}
}