feat: Support connection reset fault injection

internal/middleware/fault.go

@@ -17,15 +17,23 @@ type FaultSession struct {
 	Settings  FaultSettings
 }
 
+// Describes the fault injection settings for a session. The fault chain is
+// a series of fault injectors that are applied to the request in order. The
+// order of faults is:
+// - Delay
+// - Reset
+// - Reject
+// - Error
 type FaultSettings struct {
-	// NOTE: The way these fields are ordered represents their precedence in the
-	// fault chain.
-
 	// DelayMS is the number of milliseconds to delay the request.
 	DelayMS int64 `json:"delay_ms"`
+
 	// DelayCount is the number of times to delay the request.
 	DelayCount int `json:"delay_count"`
 
+	// ResetCount is the number of times to reset the connection.
+	ResetCount int `json:"reset_count"`
+
 	// RejectCount is the number of times to reject the request without a response.
 	// A value greater than 0 enables this fault injector.
 	RejectCount int `json:"reject_count"`
@@ -36,6 +44,7 @@ type FaultSettings struct {
 	// NOTE: Error injection only takes effect after all rejections have
 	// resolved if both of these injectors are enabled.
 	ErrorCount int `json:"error_count"`
+
 	// ErrorCode is the status code to return when the error injector is enabled.
 	ErrorCode int `json:"error_code"`
 }
@@ -81,6 +90,10 @@ func Fault(h http.Handler) http.Handler {
 
 		var faults []fault.Injector
 
+		// Since multiple injectors can be enabled, need to count the number of
+		// requests based on prior injector counts
+		countOffset := 0
+
 		if settings.DelayMS > 0 && reqCount < settings.DelayCount {
 			inj, err := fault.NewSlowInjector(time.Millisecond * time.Duration(settings.DelayMS))
 			if err != nil {
@@ -91,7 +104,14 @@ func Fault(h http.Handler) http.Handler {
 			faults = append(faults, inj)
 		}
 
-		countOffset := 0
+		// Delay injector does not increase the count offset.
+
+		if settings.ResetCount > 0 && reqCount < settings.ResetCount+countOffset {
+			faults = append(faults, &ConnectionResetInjector{})
+		}
+
+		countOffset += settings.ResetCount
+
 		if settings.RejectCount > 0 && reqCount < settings.RejectCount+countOffset {
 			inj, err := fault.NewRejectInjector()
 			if err != nil {
@@ -103,6 +123,7 @@ func Fault(h http.Handler) http.Handler {
 		}
 
 		countOffset += settings.RejectCount
+
 		if settings.ErrorCode > 0 && reqCount < (settings.ErrorCount+countOffset) {
 			inj, err := fault.NewErrorInjector(settings.ErrorCode, fault.WithStatusText("Injected error"))
 			if err != nil {

internal/middleware/reset_injector.go

@@ -0,0 +1,38 @@
+package middleware
+
+import (
+	"net"
+	"net/http"
+
+	"github.com/lingrino/go-fault"
+)
+
+var _ fault.Injector = (*ConnectionResetInjector)(nil)
+
+// Injects a connection error by closing the connection immediately while also
+// simulating a TCP RST via setting SO_LINGER to 0.
+type ConnectionResetInjector struct{}
+
+func (i *ConnectionResetInjector) Handler(_ http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		hijacker, ok := w.(http.Hijacker)
+
+		if !ok {
+			http.Error(w, "connection hijacking not supported", http.StatusInternalServerError)
+			return
+		}
+
+		conn, _, err := hijacker.Hijack()
+
+		if err != nil {
+			http.Error(w, "failed to hijack connection", http.StatusInternalServerError)
+			return
+		}
+
+		if tcpConn, ok := conn.(*net.TCPConn); ok {
+			_ = tcpConn.SetLinger(0) // Best effort RST on close
+		}
+
+		conn.Close()
+	})
+}