feat: add enable/disable capability

.pre-commit-config.yaml

@@ -0,0 +1 @@
+repos: []

go.mod

@@ -1,19 +1,22 @@
-module github.com/kffl/speedbump
+module github.com/seadiaz/speedbump
 
-go 1.17
+go 1.22
+
+require (
+	github.com/hashicorp/go-hclog v1.2.1
+	github.com/kffl/speedbump v1.1.0
+	github.com/stretchr/testify v1.8.0
+	gopkg.in/alecthomas/kingpin.v2 v2.2.6
+)
 
 require (
 	github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 // indirect
 	github.com/alecthomas/units v0.0.0-20211218093645-b94a6e3cc137 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/fatih/color v1.13.0 // indirect
-	github.com/hashicorp/go-hclog v1.2.1 // indirect
 	github.com/mattn/go-colorable v0.1.12 // indirect
 	github.com/mattn/go-isatty v0.0.14 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/stretchr/objx v0.4.0 // indirect
-	github.com/stretchr/testify v1.8.0 // indirect
 	golang.org/x/sys v0.0.0-20220503163025-988cb79eb6c6 // indirect
-	gopkg.in/alecthomas/kingpin.v2 v2.2.6 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -9,6 +9,8 @@ github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w=
 github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
 github.com/hashicorp/go-hclog v1.2.1 h1:YQsLlGDJgwhXFpucSPyVbCBviQtjlHv3jLTlp8YmtEw=
 github.com/hashicorp/go-hclog v1.2.1/go.mod h1:W4Qnvbt70Wk/zYJryRzDRU/4r0kIg0PVHBcfoyhpF5M=
+github.com/kffl/speedbump v1.1.0 h1:mTLW9ZzWP/1FQCmkZgHhKbphhqJmzzajKKuGXvjibHE=
+github.com/kffl/speedbump v1.1.0/go.mod h1:6nNWIwc8zM0l41fIArBiVdvcomulEd8v5RX9YBjJoQ4=
 github.com/mattn/go-colorable v0.1.9/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
 github.com/mattn/go-colorable v0.1.12 h1:jF+Du6AlPIjs2BiUiQlKOX0rt3SujHxPnksPKZbaA40=
 github.com/mattn/go-colorable v0.1.12/go.mod h1:u5H1YNBxpqRaxsYJYSkiCWKzEfiAb1Gb520KVy5xxl4=
@@ -18,7 +20,6 @@ github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27k
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
-github.com/stretchr/objx v0.4.0 h1:M2gUjqZET1qApGOWNSnZ49BAIMX4F/1plDv3+l31EJ4=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
@@ -33,6 +34,7 @@ golang.org/x/sys v0.0.0-20220503163025-988cb79eb6c6 h1:nonptSpoQ4vQjyraW20DXPAgl
 golang.org/x/sys v0.0.0-20220503163025-988cb79eb6c6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6 h1:jMFz6MfLP0/4fUyZle81rXUoxOBFi19VUFKVDOQfozc=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

lib/connection.go

@@ -24,6 +24,7 @@ type connection struct {
 	done              chan error
 	ctx               context.Context
 	log               hclog.Logger
+	enabled           bool
 }
 
 func (c *connection) readFromSrc() {
@@ -75,7 +76,9 @@ func (c *connection) readFromDelayQueue() {
 
 		c.log.Trace("Read from delay queue", "bytes", len(t.data))
 
-		time.Sleep(time.Until(t.delayUntil))
+		if c.enabled {
+			time.Sleep(time.Until(t.delayUntil))
+		}
 
 		_, err := c.destConn.Write(t.data)
 		if err != nil {
@@ -124,6 +127,14 @@ func (c *connection) closeProxyConnections() {
 	c.destConn.Close()
 }
 
+func (c *connection) Enable() {
+	c.enabled = true
+}
+
+func (c *connection) Disable() {
+	c.enabled = false
+}
+
 func newProxyConnection(
 	ctx context.Context,
 	clientConn io.ReadWriteCloser,
@@ -133,6 +144,7 @@ func newProxyConnection(
 	queueSize int,
 	latencyGen LatencyGenerator,
 	logger hclog.Logger,
+	enabled bool,
 ) (*connection, error) {
 	destConn, err := net.DialTCP("tcp", nil, destAddr)
 	if err != nil {
@@ -147,6 +159,7 @@ func newProxyConnection(
 		done:       make(chan error, 3),
 		ctx:        ctx,
 		log:        logger,
+		enabled:    enabled,
 	}
 
 	return c, nil

lib/connection_test.go

@@ -299,6 +299,7 @@ func TestNewProxyConnectionError(t *testing.T) {
 		100,
 		&mockLatencyGenerator{time.Millisecond * 10},
 		hclog.Default(),
+		false,
 	)
 
 	assert.NotNil(t, err)

lib/speedbump.go

@@ -26,6 +26,9 @@ type Speedbump struct {
 	ctx       context.Context
 	ctxCancel context.CancelFunc
 	log       hclog.Logger
+	disabled  bool
+
+	activeConnections []*connection
 }
 
 // SpeedbumpCfg contains Spedbump instance configuration
@@ -44,6 +47,8 @@ type SpeedbumpCfg struct {
 	Latency *LatencyCfg
 	// LogLevel can be one of: DEBUG, TRACE, INFO, WARN, ERROR
 	LogLevel string
+	// Disabled allow to start the proxy without injecting latancy until Enalbe method is called
+	Disabled bool
 }
 
 // NewSpeedbump creates a Speedbump instance based on a provided config
@@ -72,6 +77,7 @@ func NewSpeedbump(cfg *SpeedbumpCfg) (*Speedbump, error) {
 		destAddr:   *destTCPAddr,
 		latencyGen: newSimpleLatencyGenerator(time.Now(), cfg.Latency),
 		log:        l,
+		disabled:   cfg.Disabled,
 	}
 	return s, nil
 }
@@ -98,6 +104,7 @@ func (s *Speedbump) startAcceptLoop() {
 			s.queueSize,
 			s.latencyGen,
 			l,
+			!s.disabled,
 		)
 		if err != nil {
 			s.log.Warn("Creating new proxy conn failed", "err", err)
@@ -106,6 +113,7 @@ func (s *Speedbump) startAcceptLoop() {
 		}
 		s.nextConnId++
 		s.active.Add(1)
+		s.activeConnections = append(s.activeConnections, p)
 		go s.startProxyConnection(p)
 	}
 }
@@ -121,7 +129,7 @@ func (s *Speedbump) startProxyConnection(p *connection) {
 func (s *Speedbump) Start() error {
 	listener, err := net.ListenTCP("tcp", &s.srcAddr)
 	if err != nil {
-		return fmt.Errorf("Error starting TCP listener: %s", err)
+		return fmt.Errorf("starting TCP listener: %w", err)
 	}
 	s.listener = listener
 
@@ -148,3 +156,17 @@ func (s *Speedbump) Stop() {
 	s.active.Wait()
 	s.log.Info("Speedbump stopped")
 }
+
+func (s *Speedbump) Enable() {
+	s.disabled = false
+	for _, conn := range s.activeConnections {
+		conn.Enable()
+	}
+}
+
+func (s *Speedbump) Disable() {
+	s.disabled = true
+	for _, conn := range s.activeConnections {
+		conn.Disable()
+	}
+}

lib/speedbump_test.go

@@ -44,6 +44,7 @@ func TestNewSpeedbump(t *testing.T) {
 		100,
 		defaultLatencyCfg,
 		"WARN",
+		false,
 	}
 	s, err := NewSpeedbump(&cfg)
 	assert.Nil(t, err)
@@ -59,6 +60,7 @@ func TestNewSpeedbumpInvalidHost(t *testing.T) {
 		100,
 		defaultLatencyCfg,
 		"WARN",
+		false,
 	}
 	s, err := NewSpeedbump(&cfg)
 	assert.Nil(t, s)
@@ -74,6 +76,7 @@ func TestNewSpeedbumpErrorResolvingLocal(t *testing.T) {
 		100,
 		defaultLatencyCfg,
 		"WARN",
+		false,
 	}
 	s, err := NewSpeedbump(&cfg)
 	assert.Nil(t, s)
@@ -89,6 +92,7 @@ func TestNewSpeedbumpErrorResolvingDest(t *testing.T) {
 		100,
 		defaultLatencyCfg,
 		"WARN",
+		false,
 	}
 	s, err := NewSpeedbump(&cfg)
 	assert.Nil(t, s)
@@ -118,6 +122,7 @@ func TestStartListenError(t *testing.T) {
 		100,
 		defaultLatencyCfg,
 		"WARN",
+		false,
 	}
 	s, _ := NewSpeedbump(&cfg)
 
@@ -153,6 +158,7 @@ func TestSpeedbumpWithEchoServer(t *testing.T) {
 			SinePeriod:    time.Millisecond * 400,
 		},
 		"WARN",
+		false,
 	}
 	s, err := NewSpeedbump(&cfg)
 	s.Start()
@@ -195,3 +201,93 @@ func TestSpeedbumpWithEchoServer(t *testing.T) {
 	assert.Equal(t, []byte("another-test"), trimmedRes)
 	assert.True(t, isDurationCloseTo(time.Millisecond*200, secondOpElapsed, 20))
 }
+
+func TestShouldDontDelayWhenDisabled(t *testing.T) {
+	port := 9006
+	testSrvAddr := fmt.Sprintf("localhost:%d", port)
+
+	go startEchoSrv(port)
+
+	cfg := SpeedbumpCfg{
+		"localhost",
+		8000,
+		testSrvAddr,
+		0xffff,
+		100,
+		&LatencyCfg{
+			Base: time.Millisecond * 250,
+		},
+		"WARN",
+		true,
+	}
+	s, err := NewSpeedbump(&cfg)
+	s.Start()
+
+	assert.Nil(t, err)
+
+	tcpAddr, _ := net.ResolveTCPAddr("tcp", "localhost:8000")
+
+	conn, err := net.DialTCP("tcp", nil, tcpAddr)
+	if err != nil {
+		panic(err)
+	}
+
+	firstOpStart := time.Now()
+
+	conn.Write([]byte("test-string"))
+	res := make([]byte, 1024)
+	bytes, _ := conn.Read(res)
+
+	firstOpElapsed := time.Since(firstOpStart)
+
+	trimmedRes := res[:bytes]
+
+	assert.Equal(t, []byte("test-string"), trimmedRes)
+	assert.Less(t, firstOpElapsed, 1*time.Millisecond)
+}
+
+func TestShouldDelayWhenDisabledAndthenSwitchToEnable(t *testing.T) {
+	port := 9006
+	testSrvAddr := fmt.Sprintf("localhost:%d", port)
+
+	go startEchoSrv(port)
+
+	cfg := SpeedbumpCfg{
+		"localhost",
+		8000,
+		testSrvAddr,
+		0xffff,
+		100,
+		&LatencyCfg{
+			Base: time.Millisecond * 250,
+		},
+		"WARN",
+		true,
+	}
+	s, err := NewSpeedbump(&cfg)
+	s.Start()
+	s.Enable()
+
+	assert.Nil(t, err)
+
+	tcpAddr, _ := net.ResolveTCPAddr("tcp", "localhost:8000")
+
+	conn, err := net.DialTCP("tcp", nil, tcpAddr)
+	if err != nil {
+		panic(err)
+	}
+
+	firstOpStart := time.Now()
+
+	conn.Write([]byte("test-string"))
+	res := make([]byte, 1024)
+	bytes, _ := conn.Read(res)
+
+	firstOpElapsed := time.Since(firstOpStart)
+
+	trimmedRes := res[:bytes]
+
+	fmt.Printf("*** %+v\n", firstOpElapsed)
+	assert.Equal(t, []byte("test-string"), trimmedRes)
+	assert.True(t, isDurationCloseTo(time.Millisecond*250, firstOpElapsed, 20))
+}