goreplay/middleware.go

package goreplay

import (
	"bufio"
	"context"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
	"sync"
	"syscall"
)

// Middleware represents a middleware object
type Middleware struct {
	command       string
	data          chan *Message
	Stdin         io.Writer
	Stdout        io.Reader
	commandCancel context.CancelFunc
	stop          chan bool // Channel used only to indicate goroutine should shutdown
	closed        bool
	mu            sync.RWMutex
}

// NewMiddleware returns new middleware
func NewMiddleware(command string) *Middleware {
	m := new(Middleware)
	m.command = command
	m.data = make(chan *Message, 1000)
	m.stop = make(chan bool)

	commands := strings.Split(command, " ")
	ctx, cancl := context.WithCancel(context.Background())
	m.commandCancel = cancl
	cmd := exec.CommandContext(ctx, commands[0], commands[1:]...)

	m.Stdout, _ = cmd.StdoutPipe()
	m.Stdin, _ = cmd.StdinPipe()

	cmd.Stderr = os.Stderr

	go m.read(m.Stdout)

	go func() {
		defer m.Close()
		var err error
		if err = cmd.Start(); err == nil {
			err = cmd.Wait()
		}
		if err != nil {
			if e, ok := err.(*exec.ExitError); ok {
				status := e.Sys().(syscall.WaitStatus)
				if status.Signal() == syscall.SIGKILL /*killed or context canceld */ {
					return
				}
			}
			Debug(0, fmt.Sprintf("[MIDDLEWARE] command[%q] error: %q", command, err.Error()))
		}
	}()

	return m
}

// ReadFrom start a worker to read from this plugin
func (m *Middleware) ReadFrom(plugin PluginReader) {
	Debug(2, fmt.Sprintf("[MIDDLEWARE] command[%q] Starting reading from %q", m.command, plugin))
	go m.copy(m.Stdin, plugin)
}

func (m *Middleware) copy(to io.Writer, from PluginReader) {
	var buf, dst []byte

	for {
		msg, err := from.PluginRead()
		if err != nil {
			return
		}
		if msg == nil || len(msg.Data) == 0 {
			continue
		}
		buf = msg.Data
		if Settings.PrettifyHTTP {
			buf = prettifyHTTP(msg.Data)
		}
		dstLen := (len(buf)+len(msg.Meta))*2 + 1
		// if enough space was previously allocated use it instead
		if dstLen > len(dst) {
			dst = make([]byte, dstLen)
		}
		n := hex.Encode(dst, msg.Meta)
		n += hex.Encode(dst[n:], buf)
		dst[n] = '\n'

		n, err = to.Write(dst[:n+1])
		if err == nil {
			continue
		}
		if m.isClosed() {
			return
		}
	}
}

func (m *Middleware) read(from io.Reader) {
	reader := bufio.NewReader(from)
	var line []byte
	var e error
	for {
		if line, e = reader.ReadBytes('\n'); e != nil {
			if m.isClosed() {
				return
			}
			continue
		}
		buf := make([]byte, (len(line)-1)/2)
		if _, err := hex.Decode(buf, line[:len(line)-1]); err != nil {
			Debug(0, fmt.Sprintf("[MIDDLEWARE] command[%q] failed to decode err: %q", m.command, err))
			continue
		}
		var msg Message
		msg.Meta, msg.Data = payloadMetaWithBody(buf)
		select {
		case <-m.stop:
			return
		case m.data <- &msg:
		}
	}

}

// PluginRead reads message from this plugin
func (m *Middleware) PluginRead() (msg *Message, err error) {
	select {
	case <-m.stop:
		return nil, ErrorStopped
	case msg = <-m.data:
	}

	return
}

func (m *Middleware) String() string {
	return fmt.Sprintf("Modifying traffic using %q command", m.command)
}

func (m *Middleware) isClosed() bool {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.closed
}

// Close closes this plugin
func (m *Middleware) Close() error {
	if m.isClosed() {
		return nil
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	m.commandCancel()
	close(m.stop)
	m.closed = true
	return nil
}