Backgrounder

backgrounder.go

package main

import (
	"flag"
	"io"
	"net"
	"time"

	"golang.org/x/net/proxy"

	"github.com/Sirupsen/logrus"
	"github.com/juju/ratelimit"
)

var (
	mode = flag.String("mode", "client", "Mode")

	bind            = flag.String("bind", ":7999", "[Server] Bind address")
	proxyAddr       = flag.String("proxy", "", "[Client] socks5 Address")
	connect         = flag.String("connect", "127.0.0.1:7999", "[Client] Connect address")
	bandwidthInByte = flag.Int("bandwidth", 1024, "Bandwidth")
)

func main() {
	flag.Parse()

	switch *mode {
	case "client":
		for {
			send()
			time.Sleep(time.Second * 5)
		}
	case "server":
		l, err := net.Listen("tcp", *bind)
		if err != nil {
			logrus.WithError(err).Fatal("Can't bind")
		}
		defer l.Close()
		logrus.Info("Listening ", l.Addr)

		for {
			conn, err := l.Accept()
			if err != nil {
				logrus.WithError(err).Error("Error Accepting")
				continue
			}
			go handle(conn)
		}
	default:
		logrus.Panic("Unsupported mode")
	}
}

func handle(conn net.Conn) error {
	logrus.Infof("connection %s -> %s", conn.LocalAddr(), conn.RemoteAddr())
	errChan := make(chan error, 2)

	go func() {
		errChan <- alwaysRead(conn)
	}()
	go func() {
		errChan <- alwaysWrite(conn)
	}()
	err := <-errChan
	logrus.WithError(err).Errorf("stop connection %s -> %s", conn.LocalAddr(), conn.RemoteAddr())
	conn.Close()
	return err
}

func alwaysWrite(w io.Writer) error {
	bucket := ratelimit.NewBucketWithRate(float64(*bandwidthInByte), int64(*bandwidthInByte*2))
	r := ratelimit.Reader(&MyReader{}, bucket)
	buf := make([]byte, *bandwidthInByte)
	for {
		r.Read(buf)
		n, err := w.Write(buf)
		if err != nil {
			return err
		}
		logrus.Info("write", n)
	}
}

func alwaysRead(r io.Reader) error {
	buf := make([]byte, *bandwidthInByte)
	bucket := ratelimit.NewBucketWithRate(float64(*bandwidthInByte), int64(*bandwidthInByte*2))
	r2 := ratelimit.Reader(r, bucket)
	for {
		n, err := r2.Read(buf)
		logrus.Info("read", n)
		if err != nil {
			logrus.WithError(err).Error("reading")
			return err
		}
	}
}

type MyReader struct{}

func (m MyReader) Read(b []byte) (i int, e error) {
	for x := range b {
		b[x] = 'A'
	}
	return len(b), nil
}

func send() error {
	var dialer func(string, string) (net.Conn, error)

	if *proxyAddr != "" {
		p, err := proxy.SOCKS5("tcp", *proxyAddr, nil, proxy.Direct)
		if err != nil {
			return err
		}
		logrus.Infof("Setup proxy@%s", *proxyAddr)
		dialer = p.Dial
	} else {
		dialer = net.Dial
	}

	logrus.Infof("Dialing %s", *connect)
	conn, err := dialer("tcp", *connect)
	if err != nil {
		return err
	}
	defer conn.Close()
	logrus.Infof("Start backgrounding %s", *connect)
	err = handle(conn)
	return err
}