http bandwidth throttling

main.go

package main

import (
	"log"
	"net"
	"net/http"
	"time"
)

var port = ":8080"

func main() {
	http.Handle("/public/", http.StripPrefix("/public/", http.FileServer(http.Dir("public"))))

	l, e := net.Listen("tcp", port)
	if e != nil {
		log.Fatal(e)
	}

	th := Listener{
		l,
		100 * time.Millisecond, // filling frequency for Token Bucket algorithm
		1024 * 1024,            // client download speed is not more than 1 MB per second
		1024 * 1024,            // client upload speed is not more than 1 MB per second
	}

	if err := http.Serve(&th, nil); err != nil {
		log.Fatal(err)
	}

	log.Printf("listeing http on port %d\n", port)

	select {}
}

throttle.go

package main

import (
	"net"
	"sync"
	"time"

	"github.com/tsenart/tb"
)

type Rate int64

type Listener struct {
	net.Listener
	Fillingfreq time.Duration

	Down Rate
	Up   Rate
}

func (ln *Listener) Accept() (net.Conn, error) {
	c, err := ln.Listener.Accept()
	if err != nil {
		return nil, err
	}

	tc := &conn{Conn: c,
		Down:      ln.Down,
		Up:        ln.Up,
		Throttler: tb.NewThrottler(ln.Fillingfreq)}
	return tc, nil
}

type conn struct {
	net.Conn

	Down      Rate
	Up        Rate
	Throttler *tb.Throttler

	closeOnce sync.Once
	closeErr  error
}

func (c *conn) Close() error {
	c.closeOnce.Do(func() {
		err := c.Conn.Close()
		c.closeErr = err

		c.Throttler.Close()
	})
	return c.closeErr
}

func (c *conn) Write(p []byte) (n int, err error) {
	c.Throttler.Wait("down", int64(len(p)), int64(c.Down))

	return c.Conn.Write(p)
}

func (c *conn) Read(p []byte) (n int, err error) {
	c.Throttler.Wait("up", int64(len(p)), int64(c.Up))

	return c.Conn.Read(p)
}