Add Bidirectional Toxics

CHANGELOG.md

@@ -1,3 +1,7 @@
+# 2.2.0 (unreleased)
+
+* Add support for bidirectional toxics (for protocol-aware toxics) #132
+
 # 2.1.3
 
 * Update `/version` endpoint to also return a charset of utf-8. #204

CREATING_TOXICS.md

@@ -118,9 +118,18 @@ and how much memory you are comfortable with using.
 ## Stateful toxics
 
 If a toxic needs to store extra information for a connection such as the number of bytes
-transferred (See `limit_data` toxic), a state object can be created by implementing the
-`StatefulToxic` interface. This interface defines the `NewState()` function that can create
-a new state object with default values set.
+transferred (See the [limit_data toxic](https://github.com/Shopify/toxiproxy/blob/master/toxics/limit_data.go)),
+a state object can be created by implementing the `StatefulToxic` interface. This interface
+defines the `NewState()` function that can create a new state object with default values set.
+
+```go
+func (t *ExampleToxic) NewState() interface{} {
+    return &ExampleToxicState{
+        BytesRemaining: t.BytesAllowed,
+        SomeOtherState: true,
+    }
+}
+```
 
 When a stateful toxic is created, the state object will be stored on the `ToxicStub` and
 can be accessed from `toxic.Pipe()`:
@@ -133,6 +142,78 @@ If necessary, some global state can be stored in the toxic struct, which will no
 instanced per-connection. These fields cannot have a custom default value set and will
 not be thread-safe, so proper locking or atomic operations will need to be used.
 
+## Bidirectional toxics
+
+Regular toxics are limited to data flowing in a single direction, so they can't make decisions
+for the `downstream` based on a request in the `upstream`. For things like protocol aware toxics
+this is a problem.
+
+Bidirectional toxics allow state to be shared for the `upstream` and `downstream` pipes in a single
+toxic implementation. They also ensure direction-specific code is always run on the correct pipe
+(a toxic that only works on the `upstream` can't be added to the `downstream`).
+
+Creating a bidirectional toxic is done by implementing a second `Pipe()` function called `PipeUpstream()`.
+The implementation is same as a regular toxic, and can be paired with other types such as a stateful toxic.
+
+One use case of a bidirectional toxic is to mock out the backend server entirely, which is shown below:
+
+```go
+type EchoToxic struct {}
+
+type EchoToxicState struct {
+    Request          chan *stream.StreamChunk
+}
+
+// PipeUpstream handles the upstream direction
+func (t *EchoToxic) PipeUpstream(stub *toxics.ToxicStub) {
+    state := stub.State.(*EchoToxicState)
+
+    for {
+        select {
+        case <-stub.Interrupt:
+            return
+        case c := <-stub.Input:
+            if c == nil {
+                // Close the downstream when the client closes
+                close(state.Request)
+                stub.Close()
+                return
+            }
+            // Send the data to the downstream through the state object
+            state.Request <- c
+        }
+    }
+}
+
+// Pipe() will only handle the downstream on a bidirectional toxic
+func (t *EchoToxic) Pipe(stub *toxics.ToxicStub) {
+    state := stub.State.(*EchoToxicState)
+
+    for {
+        select {
+        case <-stub.Interrupt:
+            return
+        case c := <-state.Request: // Read from the upstream instead of the server
+            if c == nil {
+                stub.Close()
+                return
+            }
+            stub.Output <- c
+        }
+    }
+}
+
+func (t *EchoToxic) NewState() interface{} {
+    return &EchoToxicState{
+        Request: make(chan *stream.StreamChunk),
+    }
+}
+```
+
+This example will loop back all data send to the server back to the client. Another use case seen
+within toxiproxy is to filter http response modifications based on the request URL (See the
+[http toxic](https://github.com/Shopify/toxiproxy/tree/master/toxics/http.go)).
+
 ## Using `io.Reader` and `io.Writer`
 
 If your toxic involves modifying the data going through a proxy, you can use the `ChanReader`

api.go

@@ -443,7 +443,7 @@ func apiError(resp http.ResponseWriter, err error) bool {
 
 type proxyToxics struct {
 	*Proxy
-	Toxics []toxics.Toxic `json:"toxics"`
+	Toxics []*toxics.ToxicWrapper `json:"toxics"`
 }
 
 func proxyWithToxics(proxy *Proxy) (result proxyToxics) {

link.go

@@ -18,12 +18,13 @@ import (
 // Input > ToxicStub > ToxicStub > Output
 //
 type ToxicLink struct {
-	stubs     []*toxics.ToxicStub
-	proxy     *Proxy
-	toxics    *ToxicCollection
-	input     *stream.ChanWriter
-	output    *stream.ChanReader
-	direction stream.Direction
+	stubs      []*toxics.ToxicStub
+	proxy      *Proxy
+	toxics     *ToxicCollection
+	input      *stream.ChanWriter
+	output     *stream.ChanReader
+	direction  stream.Direction
+	pairedLink *ToxicLink
 }
 
 func NewToxicLink(proxy *Proxy, collection *ToxicCollection, direction stream.Direction) *ToxicLink {
@@ -66,9 +67,7 @@ func (link *ToxicLink) Start(name string, source io.Reader, dest io.WriteCloser)
 		link.input.Close()
 	}()
 	for i, toxic := range link.toxics.chain[link.direction] {
-		if stateful, ok := toxic.Toxic.(toxics.StatefulToxic); ok {
-			link.stubs[i].State = stateful.NewState()
-		}
+		link.InitPairState(toxic)
 
 		go link.stubs[i].Run(toxic)
 	}
@@ -87,6 +86,22 @@ func (link *ToxicLink) Start(name string, source io.Reader, dest io.WriteCloser)
 	}()
 }
 
+func (link *ToxicLink) InitPairState(toxic *toxics.ToxicWrapper) {
+	// If the toxic is stateful, create a state object or copy it from the paired link.
+	if stateful, ok := toxic.Toxic.(toxics.StatefulToxic); ok {
+		if toxic.PairedToxic == nil || link.pairedLink.stubs[toxic.PairedToxic.Index].State == nil {
+			link.stubs[toxic.Index].State = stateful.NewState()
+		} else {
+			link.stubs[toxic.Index].State = link.pairedLink.stubs[toxic.PairedToxic.Index].State
+		}
+	}
+
+	// If the toxic is paired, synchronize the toxicity so they are always in the same state.
+	if toxic.PairedToxic != nil {
+		link.stubs[toxic.Index].Toxicity = link.pairedLink.stubs[toxic.PairedToxic.Index].Toxicity
+	}
+}
+
 // Add a toxic to the end of the chain.
 func (link *ToxicLink) AddToxic(toxic *toxics.ToxicWrapper) {
 	i := len(link.stubs)
@@ -98,9 +113,7 @@ func (link *ToxicLink) AddToxic(toxic *toxics.ToxicWrapper) {
 	if link.stubs[i-1].InterruptToxic() {
 		link.stubs[i-1].Output = newin
 
-		if stateful, ok := toxic.Toxic.(toxics.StatefulToxic); ok {
-			link.stubs[i].State = stateful.NewState()
-		}
+		link.InitPairState(toxic)
 
 		go link.stubs[i].Run(toxic)
 		go link.stubs[i-1].Run(link.toxics.chain[link.direction][i-1])

proxy.go

@@ -2,13 +2,11 @@ package toxiproxy
 
 import (
 	"errors"
+	"net"
 	"sync"
 
-	"github.com/Shopify/toxiproxy/stream"
 	"github.com/sirupsen/logrus"
 	tomb "gopkg.in/tomb.v1"
-
-	"net"
 )
 
 // Proxy represents the proxy in its entirity with all its links. The main
@@ -176,8 +174,7 @@ func (proxy *Proxy) server() {
 		proxy.connections.list[name+"upstream"] = upstream
 		proxy.connections.list[name+"downstream"] = client
 		proxy.connections.Unlock()
-		proxy.Toxics.StartLink(name+"upstream", client, upstream, stream.Upstream)
-		proxy.Toxics.StartLink(name+"downstream", upstream, client, stream.Downstream)
+		proxy.Toxics.StartLinks(name, client, upstream)
 	}
 }
 

toxic_collection.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"net"
 	"strings"
 	"sync"
 
@@ -61,18 +62,17 @@ func (c *ToxicCollection) GetToxic(name string) *toxics.ToxicWrapper {
 	return c.findToxicByName(name)
 }
 
-func (c *ToxicCollection) GetToxicArray() []toxics.Toxic {
+func (c *ToxicCollection) GetToxicArray() []*toxics.ToxicWrapper {
 	c.Lock()
 	defer c.Unlock()
 
-	result := make([]toxics.Toxic, 0)
+	result := make([]*toxics.ToxicWrapper, 0)
 	for dir := range c.chain {
-		for i, toxic := range c.chain[dir] {
-			if i == 0 {
-				// Skip the first noop toxic, it should not be visible
-				continue
+		for _, toxic := range c.chain[dir] {
+			if len(toxic.Name) > 0 {
+				// Skip toxics with no name, they should be hidden
+				result = append(result, toxic)
 			}
-			result = append(result, toxic)
 		}
 	}
 	return result
@@ -96,20 +96,31 @@ func (c *ToxicCollection) AddToxicJson(data io.Reader) (*toxics.ToxicWrapper, er
 		return nil, joinError(err, ErrBadRequestBody)
 	}
 
-	switch strings.ToLower(wrapper.Stream) {
-	case "downstream":
-		wrapper.Direction = stream.Downstream
-	case "upstream":
-		wrapper.Direction = stream.Upstream
-	default:
-		return nil, ErrInvalidStream
+	if toxics.New(wrapper) == nil {
+		return nil, ErrInvalidToxicType
 	}
-	if wrapper.Name == "" {
-		wrapper.Name = fmt.Sprintf("%s_%s", wrapper.Type, wrapper.Stream)
+
+	if wrapper.PairedToxic == nil {
+		switch strings.ToLower(wrapper.Stream) {
+		case "downstream":
+			wrapper.Direction = stream.Downstream
+		case "upstream":
+			wrapper.Direction = stream.Upstream
+		default:
+			return nil, ErrInvalidStream
+		}
+	} else {
+		wrapper.Stream = "both"
+		wrapper.Direction = stream.Downstream
+		wrapper.PairedToxic.Direction = stream.Upstream
 	}
 
-	if toxics.New(wrapper) == nil {
-		return nil, ErrInvalidToxicType
+	if wrapper.Name == "" {
+		if wrapper.PairedToxic != nil {
+			wrapper.Name = wrapper.Type
+		} else {
+			wrapper.Name = fmt.Sprintf("%s_%s", wrapper.Type, wrapper.Stream)
+		}
 	}
 
 	found := c.findToxicByName(wrapper.Name)
@@ -129,6 +140,9 @@ func (c *ToxicCollection) AddToxicJson(data io.Reader) (*toxics.ToxicWrapper, er
 	}
 
 	c.chainAddToxic(wrapper)
+	if wrapper.PairedToxic != nil {
+		c.chainAddToxic(wrapper.PairedToxic)
+	}
 	return wrapper, nil
 }
 
@@ -151,6 +165,10 @@ func (c *ToxicCollection) UpdateToxicJson(name string, data io.Reader) (*toxics.
 		}
 		toxic.Toxicity = attrs.Toxicity
 
+		if toxic.PairedToxic != nil {
+			toxic.PairedToxic.Toxicity = attrs.Toxicity
+			c.chainUpdateToxic(toxic.PairedToxic)
+		}
 		c.chainUpdateToxic(toxic)
 		return toxic, nil
 	}
@@ -163,19 +181,30 @@ func (c *ToxicCollection) RemoveToxic(name string) error {
 
 	toxic := c.findToxicByName(name)
 	if toxic != nil {
+		if toxic.PairedToxic != nil {
+			c.chainRemoveToxic(toxic.PairedToxic)
+		}
 		c.chainRemoveToxic(toxic)
 		return nil
 	}
 	return ErrToxicNotFound
 }
 
-func (c *ToxicCollection) StartLink(name string, input io.Reader, output io.WriteCloser, direction stream.Direction) {
+func (c *ToxicCollection) StartLinks(name string, client, upstream net.Conn) {
 	c.Lock()
 	defer c.Unlock()
 
-	link := NewToxicLink(c.proxy, c, direction)
-	link.Start(name, input, output)
-	c.links[name] = link
+	linkUp := NewToxicLink(c.proxy, c, stream.Upstream)
+	linkDown := NewToxicLink(c.proxy, c, stream.Downstream)
+
+	linkUp.pairedLink = linkDown
+	linkDown.pairedLink = linkUp
+
+	linkUp.Start(name+"upstream", client, upstream)
+	linkDown.Start(name+"downstream", upstream, client)
+
+	c.links[name+"upstream"] = linkUp
+	c.links[name+"downstream"] = linkDown
 }
 
 func (c *ToxicCollection) RemoveLink(name string) {
@@ -187,12 +216,8 @@ func (c *ToxicCollection) RemoveLink(name string) {
 // All following functions assume the lock is already grabbed
 func (c *ToxicCollection) findToxicByName(name string) *toxics.ToxicWrapper {
 	for dir := range c.chain {
-		for i, toxic := range c.chain[dir] {
-			if i == 0 {
-				// Skip the first noop toxic, it has no name
-				continue
-			}
-			if toxic.Name == name {
+		for _, toxic := range c.chain[dir] {
+			if len(toxic.Name) > 0 && toxic.Name == name {
 				return toxic
 			}
 		}