Introducing Jim, the Chaos Monkey

JIM.md

@@ -0,0 +1,53 @@
+Introduction to Jim
+===================
+
+Jim is the MailHog Chaos Monkey, inspired by Netflix.
+
+You can invite Jim to the party using the `invite-jim` flag:
+
+    MailHog -invite-jim
+
+With Jim around, things aren't going to work how you expect.
+
+### What can Jim do?
+
+* Reject connections
+* Rate limit connections
+* Reject authentication
+* Reject senders
+* Reject recipients
+
+It does this randomly, but within defined parameters.
+
+You can control these using the following command line flags:
+
+| Flag                  | Default | Description
+| --------------------- | ------- | ----
+| -invite-jim           | false   | Set to true to invite Jim
+| -jim-disconnect       | 0.005   | Chance of randomly disconnecting a session
+| -jim-accept           | 0.99    | Chance of accepting an incoming connection
+| -jim-linkspeed-affect | 0.1     | Chance of applying a rate limit
+| -jim-linkspeed-min    | 1024    | Minimum link speed (in bytes per second)
+| -jim-linkspeed-max    | 10240   | Maximum link speed (in bytes per second)
+| -jim-reject-sender    | 0.05    | Chance of rejecting a MAIL FROM command
+| -jim-reject-recipient | 0.05    | Chance of rejecting a RCPT TO command
+| -jim-reject-auth      | 0.05    | Chance of rejecting an AUTH command
+
+If you enable Jim, you enable all parts. To disable individual parts, set the chance
+of it happening to 0, e.g. to disable connection rate limiting:
+
+    MailHog -invite-jim -jim-linkspeed-affect=0
+
+### Examples
+
+Always rate limit to 1 byte per second:
+
+    MailHog -invite-jim -jim-linkspeed-affect=1 -jim-linkspeed-max=1 -jim-linkspeed-min=1
+
+Disconnect clients after approximately 5 commands:
+
+    MailHog -invite-jim -jim-disconnect=0.2
+
+Simulate a mobile connection (at 10-100kbps) for 10% of clients:
+
+    MailHog -invite-jim -jim-linkspeed-affect=0.1 -jim-linkspeed-min=1250 -jim-linkspeed-max=12500

MailHog-Server/config/config.go

@@ -4,6 +4,7 @@ import (
 	"flag"
 	"log"
 
+	"github.com/ian-kent/Go-MailHog/MailHog-Server/monkey"
 	"github.com/ian-kent/Go-MailHog/data"
 	"github.com/ian-kent/Go-MailHog/storage"
 	"github.com/ian-kent/envconf"
@@ -30,12 +31,15 @@ type Config struct {
 	MongoDb      string
 	MongoColl    string
 	StorageType  string
+	InviteJim    bool
 	Storage      storage.Storage
 	MessageChan  chan *data.Message
 	Assets       func(asset string) ([]byte, error)
+	Monkey       monkey.ChaosMonkey
 }
 
 var cfg = DefaultConfig()
+var jim = &monkey.Jim{}
 
 func Configure() *Config {
 	switch cfg.StorageType {
@@ -56,6 +60,13 @@ func Configure() *Config {
 		log.Fatalf("Invalid storage type %s", cfg.StorageType)
 	}
 
+	if cfg.InviteJim {
+		jim.Configure(func(message string, args ...interface{}) {
+			log.Printf(message, args...)
+		})
+		cfg.Monkey = jim
+	}
+
 	return cfg
 }
 
@@ -67,4 +78,6 @@ func RegisterFlags() {
 	flag.StringVar(&cfg.MongoUri, "mongouri", envconf.FromEnvP("MH_MONGO_URI", "127.0.0.1:27017").(string), "MongoDB URI, e.g. 127.0.0.1:27017")
 	flag.StringVar(&cfg.MongoDb, "mongodb", envconf.FromEnvP("MH_MONGO_DB", "mailhog").(string), "MongoDB database, e.g. mailhog")
 	flag.StringVar(&cfg.MongoColl, "mongocoll", envconf.FromEnvP("MH_MONGO_COLLECTION", "messages").(string), "MongoDB collection, e.g. messages")
+	flag.BoolVar(&cfg.InviteJim, "invite-jim", envconf.FromEnvP("MH_INVITE_JIM", false).(bool), "Decide whether to invite Jim (beware, he causes trouble)")
+	jim.RegisterFlags()
 }

MailHog-Server/monkey/jim.go

@@ -0,0 +1,105 @@
+package monkey
+
+import (
+	"flag"
+	"math/rand"
+	"net"
+	"time"
+
+	"github.com/ian-kent/linkio"
+)
+
+// Jim is a chaos monkey
+type Jim struct {
+	disconnectChance      float64
+	acceptChance          float64
+	linkSpeedAffect       float64
+	linkSpeedMin          float64
+	linkSpeedMax          float64
+	rejectSenderChance    float64
+	rejectRecipientChance float64
+	rejectAuthChance      float64
+	logf                  func(message string, args ...interface{})
+}
+
+// RegisterFlags implements ChaosMonkey.RegisterFlags
+func (j *Jim) RegisterFlags() {
+	flag.Float64Var(&j.disconnectChance, "jim-disconnect", 0.005, "Chance of disconnect")
+	flag.Float64Var(&j.acceptChance, "jim-accept", 0.99, "Chance of accept")
+	flag.Float64Var(&j.linkSpeedAffect, "jim-linkspeed-affect", 0.1, "Chance of affecting link speed")
+	flag.Float64Var(&j.linkSpeedMin, "jim-linkspeed-min", 1024, "Minimum link speed (in bytes per second)")
+	flag.Float64Var(&j.linkSpeedMax, "jim-linkspeed-max", 10240, "Maximum link speed (in bytes per second)")
+	flag.Float64Var(&j.rejectSenderChance, "jim-reject-sender", 0.05, "Chance of rejecting a sender (MAIL FROM)")
+	flag.Float64Var(&j.rejectRecipientChance, "jim-reject-recipient", 0.05, "Chance of rejecting a recipient (RCPT TO)")
+	flag.Float64Var(&j.rejectAuthChance, "jim-reject-auth", 0.05, "Chance of rejecting authentication (AUTH)")
+}
+
+// Configure implements ChaosMonkey.Configure
+func (j *Jim) Configure(logf func(string, ...interface{})) {
+	j.logf = logf
+	rand.Seed(time.Now().Unix())
+}
+
+// Accept implements ChaosMonkey.Accept
+func (j *Jim) Accept(conn net.Conn) bool {
+	if rand.Float64() > j.acceptChance {
+		j.logf("Jim: Rejecting connection\n")
+		return false
+	}
+	j.logf("Jim: Allowing connection\n")
+	return true
+}
+
+// LinkSpeed implements ChaosMonkey.LinkSpeed
+func (j *Jim) LinkSpeed() *linkio.Throughput {
+	rand.Seed(time.Now().Unix())
+	if rand.Float64() < j.linkSpeedAffect {
+		lsDiff := j.linkSpeedMax - j.linkSpeedMin
+		lsAffect := j.linkSpeedMin + (lsDiff * rand.Float64())
+		f := linkio.Throughput(lsAffect) * linkio.BytePerSecond
+		j.logf("Jim: Restricting throughput to %s\n", f)
+		return &f
+	}
+	j.logf("Jim: Allowing unrestricted throughput")
+	return nil
+}
+
+// ValidRCPT implements ChaosMonkey.ValidRCPT
+func (j *Jim) ValidRCPT(rcpt string) bool {
+	if rand.Float64() < j.rejectRecipientChance {
+		j.logf("Jim: Rejecting recipient %s\n", rcpt)
+		return false
+	}
+	j.logf("Jim: Allowing recipient%s\n", rcpt)
+	return true
+}
+
+// ValidMAIL implements ChaosMonkey.ValidMAIL
+func (j *Jim) ValidMAIL(mail string) bool {
+	if rand.Float64() < j.rejectSenderChance {
+		j.logf("Jim: Rejecting sender %s\n", mail)
+		return false
+	}
+	j.logf("Jim: Allowing sender %s\n", mail)
+	return true
+}
+
+// ValidAUTH implements ChaosMonkey.ValidAUTH
+func (j *Jim) ValidAUTH(mechanism string, args ...string) bool {
+	if rand.Float64() < j.rejectAuthChance {
+		j.logf("Jim: Rejecting authentication %s: %s\n", mechanism, args)
+		return false
+	}
+	j.logf("Jim: Allowing authentication %s: %s\n", mechanism, args)
+	return true
+}
+
+// Disconnect implements ChaosMonkey.Disconnect
+func (j *Jim) Disconnect() bool {
+	if rand.Float64() < j.disconnectChance {
+		j.logf("Jim: Being nasty, kicking them off\n")
+		return true
+	}
+	j.logf("Jim: Being nice, letting them stay\n")
+	return false
+}

MailHog-Server/monkey/monkey.go

@@ -0,0 +1,28 @@
+package monkey
+
+import (
+	"net"
+
+	"github.com/ian-kent/linkio"
+)
+
+// ChaosMonkey should be implemented by chaos monkeys!
+type ChaosMonkey interface {
+	RegisterFlags()
+	Configure(func(string, ...interface{}))
+
+	// Accept is called for each incoming connection. Returning false closes the connection.
+	Accept(conn net.Conn) bool
+	// LinkSpeed sets the maximum connection throughput (in one direction)
+	LinkSpeed() *linkio.Throughput
+
+	// ValidRCPT is called for the RCPT command. Returning false signals an invalid recipient.
+	ValidRCPT(rcpt string) bool
+	// ValidMAIL is called for the MAIL command. Returning false signals an invalid sender.
+	ValidMAIL(mail string) bool
+	// ValidAUTH is called after authentication. Returning false signals invalid authentication.
+	ValidAUTH(mechanism string, args ...string) bool
+
+	// Disconnect is called after every read. Returning true will close the connection.
+	Disconnect() bool
+}

MailHog-Server/smtp/session.go

@@ -7,9 +7,11 @@ import (
 	"log"
 	"strings"
 
+	"github.com/ian-kent/Go-MailHog/MailHog-Server/monkey"
 	"github.com/ian-kent/Go-MailHog/data"
 	"github.com/ian-kent/Go-MailHog/smtp/protocol"
 	"github.com/ian-kent/Go-MailHog/storage"
+	"github.com/ian-kent/linkio"
 )
 
 // Session represents a SMTP session using net.TCPConn
@@ -21,13 +23,32 @@ type Session struct {
 	remoteAddress string
 	isTLS         bool
 	line          string
+	link          *linkio.Link
+
+	reader io.Reader
+	writer io.Writer
+	monkey monkey.ChaosMonkey
 }
 
 // Accept starts a new SMTP session using io.ReadWriteCloser
-func Accept(remoteAddress string, conn io.ReadWriteCloser, storage storage.Storage, messageChan chan *data.Message, hostname string) {
+func Accept(remoteAddress string, conn io.ReadWriteCloser, storage storage.Storage, messageChan chan *data.Message, hostname string, monkey monkey.ChaosMonkey) {
+	defer conn.Close()
+
 	proto := protocol.NewProtocol()
 	proto.Hostname = hostname
-	session := &Session{conn, proto, storage, messageChan, remoteAddress, false, ""}
+	var link *linkio.Link
+	reader := io.Reader(conn)
+	writer := io.Writer(conn)
+	if monkey != nil {
+		linkSpeed := monkey.LinkSpeed()
+		if linkSpeed != nil {
+			link = linkio.NewLink(*linkSpeed * linkio.BytePerSecond)
+			reader = link.NewLinkReader(io.Reader(conn))
+			writer = link.NewLinkWriter(io.Writer(conn))
+		}
+	}
+
+	session := &Session{conn, proto, storage, messageChan, remoteAddress, false, "", link, reader, writer, monkey}
 	proto.LogHandler = session.logf
 	proto.MessageReceivedHandler = session.acceptMessage
 	proto.ValidateSenderHandler = session.validateSender
@@ -37,19 +58,42 @@ func Accept(remoteAddress string, conn io.ReadWriteCloser, storage storage.Stora
 	session.logf("Starting session")
 	session.Write(proto.Start())
 	for session.Read() == true {
+		if monkey != nil && monkey.Disconnect != nil && monkey.Disconnect() {
+			session.conn.Close()
+			break
+		}
 	}
 	session.logf("Session ended")
 }
 
 func (c *Session) validateAuthentication(mechanism string, args ...string) (errorReply *protocol.Reply, ok bool) {
+	if c.monkey != nil {
+		ok := c.monkey.ValidAUTH(mechanism, args...)
+		if !ok {
+			// FIXME better error?
+			return protocol.ReplyUnrecognisedCommand(), false
+		}
+	}
 	return nil, true
 }
 
 func (c *Session) validateRecipient(to string) bool {
+	if c.monkey != nil {
+		ok := c.monkey.ValidRCPT(to)
+		if !ok {
+			return false
+		}
+	}
 	return true
 }
 
 func (c *Session) validateSender(from string) bool {
+	if c.monkey != nil {
+		ok := c.monkey.ValidMAIL(from)
+		if !ok {
+			return false
+		}
+	}
 	return true
 }
 
@@ -69,7 +113,7 @@ func (c *Session) logf(message string, args ...interface{}) {
 // Read reads from the underlying net.TCPConn
 func (c *Session) Read() bool {
 	buf := make([]byte, 1024)
-	n, err := io.Reader(c.conn).Read(buf)
+	n, err := c.reader.Read(buf)
 
 	if n == 0 {
 		c.logf("Connection closed by remote host\n")
@@ -112,6 +156,6 @@ func (c *Session) Write(reply *protocol.Reply) {
 		logText := strings.Replace(l, "\n", "\\n", -1)
 		logText = strings.Replace(logText, "\r", "\\r", -1)
 		c.logf("Sent %d bytes: '%s'", len(l), logText)
-		io.Writer(c.conn).Write([]byte(l))
+		c.writer.Write([]byte(l))
 	}
 }

MailHog-Server/smtp/session_test.go

@@ -40,7 +40,7 @@ func TestAccept(t *testing.T) {
 	Convey("Accept should handle a connection", t, func() {
 		frw := &fakeRw{}
 		mChan := make(chan *data.Message)
-		Accept("1.1.1.1:11111", frw, storage.CreateInMemory(), mChan, "localhost")
+		Accept("1.1.1.1:11111", frw, storage.CreateInMemory(), mChan, "localhost", nil)
 	})
 }
 
@@ -52,7 +52,7 @@ func TestSocketError(t *testing.T) {
 			},
 		}
 		mChan := make(chan *data.Message)
-		Accept("1.1.1.1:11111", frw, storage.CreateInMemory(), mChan, "localhost")
+		Accept("1.1.1.1:11111", frw, storage.CreateInMemory(), mChan, "localhost", nil)
 	})
 }
 
@@ -98,7 +98,7 @@ func TestAcceptMessage(t *testing.T) {
 			//So(m, ShouldNotBeNil)
 			wg.Done()
 		}()
-		Accept("1.1.1.1:11111", frw, storage.CreateInMemory(), mChan, "localhost")
+		Accept("1.1.1.1:11111", frw, storage.CreateInMemory(), mChan, "localhost", nil)
 		wg.Wait()
 		So(handlerCalled, ShouldBeTrue)
 	})

MailHog-Server/smtp/smtp.go

@@ -22,14 +22,22 @@ func Listen(cfg *config.Config, exitCh chan int) *net.TCPListener {
 			log.Printf("[SMTP] Error accepting connection: %s\n", err)
 			continue
 		}
-		defer conn.Close()
+
+		if cfg.Monkey != nil {
+			ok := cfg.Monkey.Accept(conn)
+			if !ok {
+				conn.Close()
+				continue
+			}
+		}
 
 		go Accept(
 			conn.(*net.TCPConn).RemoteAddr().String(),
 			io.ReadWriteCloser(conn),
 			cfg.Storage,
 			cfg.MessageChan,
 			cfg.Hostname,
+			cfg.Monkey,
 		)
 	}
 }

Makefile

@@ -25,6 +25,7 @@ deps:
 	go get github.com/ian-kent/go-log/log
 	go get github.com/ian-kent/envconf
 	go get github.com/ian-kent/goose
+	go get github.com/ian-kent/linkio
 	go get github.com/jteeuwen/go-bindata/...
 	go get labix.org/v2/mgo
 	# added to fix travis issues