Teach proxies to use new ratelimiter wrapper; updated version#

Author: opencoff

Makefile

@@ -1,10 +1,6 @@
 # Simple makefile to build this go program
 
-all: depend
-	./build
-
-
-static: depend
+all static: depend
 	./build -s
 
 

go.mod

@@ -1,10 +1,10 @@
 module github.com/opencoff/go-proxies
 
-go 1.12
+go 1.14
 
 require (
 	github.com/opencoff/go-logger v0.0.0-20190612060632-bf4528b7367d
-	github.com/opencoff/go-ratelimit v0.6.0
-	github.com/opencoff/pflag v0.3.3
-	gopkg.in/yaml.v2 v2.2.2
+	github.com/opencoff/go-ratelimit v0.7.0
+	github.com/opencoff/pflag v0.5.0
+	gopkg.in/yaml.v2 v2.3.0
 )

go.sum

@@ -1,11 +1,13 @@
 github.com/opencoff/go-logger v0.0.0-20190612060632-bf4528b7367d h1:kBo3CACJRG/TO3VzmSSmzyHbEivF8V20hXuXtdAh0dU=
 github.com/opencoff/go-logger v0.0.0-20190612060632-bf4528b7367d/go.mod h1:0uZokzKt+uCJkbz12vSoChasSJoLc2aNuCS0A/U7Dqs=
-github.com/opencoff/go-ratelimit v0.6.0 h1:u+OUXaHtwJ3J9Yd+hGyU+JSafaoPBXvXrlEbWuHl8RQ=
-github.com/opencoff/go-ratelimit v0.6.0/go.mod h1:MlK6FlcsSUqs9xz3r3ZVa7E02OgYpYB7U6JAMNx1ZXg=
+github.com/opencoff/go-ratelimit v0.7.0 h1:hXadrYOPFlS6l+jmrok87BX0Oh+oeWuLelNYXAUgGqA=
+github.com/opencoff/go-ratelimit v0.7.0/go.mod h1:CZOjkRlhRo07XJt81kMF0NfOP7cYTfhZG1zPU5AAK78=
 github.com/opencoff/golang-lru v0.6.0 h1:e5jyAHA4AJbohh8mmPB6JpTvZMVrnh3z5GFAqTADVm8=
 github.com/opencoff/golang-lru v0.6.0/go.mod h1:Ll98eBFICVmenoj+uJfH+ReFgDMD+nuK9VshgMwDs80=
-github.com/opencoff/pflag v0.3.3 h1:yohZkwYGPkB34WXvUQzU5GyLhImnjfePDARUaE8me3U=
-github.com/opencoff/pflag v0.3.3/go.mod h1:mTLzGGUGda1Av3d34iAJlh0JIlRxmFZtmc6qoWPspK0=
+github.com/opencoff/pflag v0.5.0 h1:kK3cSTlGj0fHby/PoFzHkf+Jx3PdiACJwzYDWEWlEKQ=
+github.com/opencoff/pflag v0.5.0/go.mod h1:mTLzGGUGda1Av3d34iAJlh0JIlRxmFZtmc6qoWPspK0=
+golang.org/x/time v0.0.0-20191024005414-555d28b269f0 h1:/5xXl8Y5W96D+TtHSlonuFqGHIWVuyCkGJLwGh9JJFs=
+golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
-gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.3.0 h1:clyUAQHOM3G0M3f5vQj7LuJrETvjVot3Z5el9nffUtU=
+gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=

src/copy.go

@@ -10,21 +10,20 @@ package main
 
 import (
 	"context"
-	"net"
 	"io"
+	"net"
 	"sync"
 	"time"
 )
 
-
 type CancellableCopier struct {
 	Lhs *net.TCPConn
 	Rhs *net.TCPConn
 
-	ReadTimeout int
+	ReadTimeout  int
 	WriteTimeout int
 
-	IOBufsize  int
+	IOBufsize int
 }
 
 // CancellableCopy does bi-directional I/O between two connections d & s. It is cancellable
@@ -38,7 +37,7 @@ func (c *CancellableCopier) Copy(ctx context.Context) (nLhs, nRhs int, err error
 	}
 
 	if c.ReadTimeout <= 0 {
-		c.ReadTimeout = 10  // seconds
+		c.ReadTimeout = 10 // seconds
 	}
 
 	if c.WriteTimeout <= 0 {
@@ -74,27 +73,23 @@ func (c *CancellableCopier) Copy(ctx context.Context) (nLhs, nRhs int, err error
 		_, nRhs, _ = c.copyBuf(c.Rhs, c.Lhs, b1)
 	}()
 
-
 	// Wait for parent to kill us or the copy routines to end.
 	// If parent kills us, we wait for copy-routines to end as well.
 	select {
 	case <-ctx.Done():
 		// close the sockets and force the i/o loop in copybuf to end.
 		c.Lhs.Close()
 		c.Rhs.Close()
-		<- ch
+		<-ch
 
 	case <-ch:
 	}
 
-
 	// XXX Gah which error do I report?
 	err = nil
 	return
 }
 
-
-
 func (c *CancellableCopier) copyBuf(d, s *net.TCPConn, b []byte) (nr, nw int, err error) {
 	rto := time.Duration(c.ReadTimeout) * time.Second
 	wto := time.Duration(c.WriteTimeout) * time.Second

src/http.go

@@ -29,8 +29,7 @@ type HTTPProxy struct {
 	// listen address
 	conf *ListenConf
 
-	grl *ratelimit.RateLimiter
-	prl *ratelimit.PerIPRateLimiter
+	rl *ratelimit.RateLimiter
 
 	log  *L.Logger
 	ulog *L.Logger
@@ -58,8 +57,10 @@ func NewHTTPProxy(lc *ListenConf, log, ulog *L.Logger) (Proxy, error) {
 	}
 
 	// Conf file specifies ratelimit as N conns/sec
-	grl, _ := ratelimit.New(lc.Ratelimit.Global, 1)
-	prl, _ := ratelimit.NewPerIP(lc.Ratelimit.PerHost, 1, 30000)
+	rl, err := ratelimit.New(lc.Ratelimit.Global, lc.Ratelimit.PerHost, 10000)
+	if err != nil {
+		die("%s: can't setup rate limiter: %s", addr, err)
+	}
 
 	ctx, cancel := context.WithCancel(context.Background())
 
@@ -73,8 +74,7 @@ func NewHTTPProxy(lc *ListenConf, log, ulog *L.Logger) (Proxy, error) {
 		conf:        lc,
 		log:         log.New("http-"+ln.Addr().String(), 0),
 		ulog:        ulog,
-		grl:         grl,
-		prl:         prl,
+		rl:          rl,
 		ctx:         ctx,
 		cancel:      cancel,
 
@@ -254,7 +254,6 @@ func (p *HTTPProxy) handleConnect(w http.ResponseWriter, r *http.Request) {
 
 	host := extractHost(r.URL)
 
-
 	ctx := r.Context()
 
 	dest, err := p.tr.DialContext(ctx, "tcp", host)
@@ -272,19 +271,17 @@ func (p *HTTPProxy) handleConnect(w http.ResponseWriter, r *http.Request) {
 
 	p.log.Debug("%s: CONNECT %s", s.RemoteAddr().String(), host)
 
-
 	cp := &CancellableCopier{
 		Lhs:          s,
 		Rhs:          d,
-		ReadTimeout:  10,	// XXX Config file
-		WriteTimeout: 15,	// XXX Config file
+		ReadTimeout:  10, // XXX Config file
+		WriteTimeout: 15, // XXX Config file
 		IOBufsize:    16384,
 	}
 
 	cp.Copy(ctx)
 }
 
-
 // Accept() new socket connections from the listener
 // Note:
 //   - HTTPProxy is also a TCPListener
@@ -311,13 +308,13 @@ func (p *HTTPProxy) Accept() (net.Conn, error) {
 			return nil, err
 		}
 
-		if p.grl.Limit() {
+		if !p.rl.Allow() {
 			nc.Close()
 			p.log.Debug("%s: globally ratelimited", nc.RemoteAddr().String())
 			continue
 		}
 
-		if p.prl.Limit(nc.RemoteAddr()) {
+		if !p.rl.AllowHost(nc.RemoteAddr()) {
 			nc.Close()
 			p.log.Debug("%s: per-IP ratelimited", nc.RemoteAddr().String())
 			continue

src/main.go

@@ -62,8 +62,8 @@ type ListenConf struct {
 }
 
 type RateLimit struct {
-	Global  uint `yaml:"global"`
-	PerHost uint `yaml:"perhost"`
+	Global  int `yaml:"global"`
+	PerHost int `yaml:"perhost"`
 }
 
 // An IP/Subnet

src/priv_unix.go

@@ -15,7 +15,6 @@ import (
 	"syscall"
 )
 
-
 // DropPrivilege changes the uid/gid. It dies if it cannot.
 func DropPrivilege(uids, gids string) {
 
@@ -43,7 +42,6 @@ func DropPrivilege(uids, gids string) {
 		}
 	}
 
-
 	if len(uids) > 0 {
 		ui, err := u.Lookup(uids)
 		if err != nil {
@@ -57,11 +55,8 @@ func DropPrivilege(uids, gids string) {
 			die("can't parse integer uid %s: %s", ui.Uid, err)
 		}
 
-
 		if err = syscall.Setuid(uid); err != nil {
 			die("can't change Uid to %d: %s", uid, err)
 		}
 	}
 }
-
-

src/socks.go

@@ -9,13 +9,13 @@
 package main
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
 	"net"
 	"sync"
 	"time"
-	"context"
 	//"encoding/hex"
 
 	L "github.com/opencoff/go-logger"
@@ -33,19 +33,18 @@ type Methods struct {
 type socksProxy struct {
 	*net.TCPListener
 
-	cfg  *ListenConf // config block
+	cfg *ListenConf // config block
 
-	bind net.Addr    // address to bind to when connect to remote
-	log  *L.Logger   // Shortcut to logger
-	ulog *L.Logger   // URL Logger
+	bind net.Addr  // address to bind to when connect to remote
+	log  *L.Logger // Shortcut to logger
+	ulog *L.Logger // URL Logger
 
-	grl  *ratelimit.RateLimiter
-	prl  *ratelimit.PerIPRateLimiter
+	rl *ratelimit.RateLimiter
 
-	ctx  context.Context
+	ctx    context.Context
 	cancel context.CancelFunc
 
-	wg   sync.WaitGroup
+	wg sync.WaitGroup
 }
 
 // Make a new proxy server
@@ -73,20 +72,21 @@ func NewSocksv5Proxy(cfg *ListenConf, log, ulog *L.Logger) (px *socksProxy, err
 
 	log = log.New("socks-"+ln.Addr().String(), 0)
 
-	grl, _ := ratelimit.New(cfg.Ratelimit.Global, 1)
-	prl, _ := ratelimit.NewPerIP(cfg.Ratelimit.PerHost, 1, 30000)
+	rl, err := ratelimit.New(cfg.Ratelimit.Global, cfg.Ratelimit.PerHost, 10000)
+	if err != nil {
+		die("%s: can't setup rate limiter: %s", addr, err)
+	}
 
 	ctx, cancel := context.WithCancel(context.Background())
 	px = &socksProxy{
-		TCPListener:  ln,
-		cfg:          cfg,
-		bind:         addr,
-		log:          log,
-		ulog:         ulog,
-		grl:          grl,
-		prl:          prl,
-		ctx:          ctx,
-		cancel:       cancel,
+		TCPListener: ln,
+		cfg:         cfg,
+		bind:        addr,
+		log:         log,
+		ulog:        ulog,
+		rl:          rl,
+		ctx:         ctx,
+		cancel:      cancel,
 	}
 
 	return
@@ -109,7 +109,6 @@ func (px *socksProxy) Stop() {
 	px.log.Info("SOCKS proxy shutdown")
 }
 
-
 // start the proxy
 // Caller is expected to kick this off as a go-routine
 // XXX Also need a global limit on total concurrent connections?
@@ -146,13 +145,13 @@ func (px *socksProxy) accept() {
 		rem := conn.RemoteAddr().String()
 
 		// Ratelimit before anything else we do
-		if px.grl.Limit() {
+		if !px.rl.Allow() {
 			conn.Close()
 			log.Debug("global ratelimit reached: %s", rem)
 			continue
 		}
 
-		if px.prl.Limit(conn.RemoteAddr()) {
+		if !px.rl.AllowHost(conn.RemoteAddr()) {
 			conn.Close()
 			log.Debug("per-host ratelimit reached: %s", rem)
 			continue
@@ -221,8 +220,8 @@ func (px *socksProxy) Proxy(lhs net.Conn) {
 	cp := &CancellableCopier{
 		Lhs:          lx,
 		Rhs:          rx,
-		ReadTimeout:  10,	// XXX Config file
-		WriteTimeout: 15,	// XXX Config file
+		ReadTimeout:  10, // XXX Config file
+		WriteTimeout: 15, // XXX Config file
 		IOBufsize:    16384,
 	}
 

version

@@ -1 +1 @@
-0.8.0
+0.8.1