goreplay/pkg/raw/input_raw.go

package raw

import (
	"context"
	"fmt"
	"net"
	"strconv"
	"strings"
	"sync"

	"github.com/buger/goreplay/internal/capture"
	"github.com/buger/goreplay/internal/tcp"
	"github.com/buger/goreplay/pkg/http_proto"
	"github.com/buger/goreplay/pkg/plugin"
	"github.com/buger/goreplay/pkg/proto"

	"github.com/rs/zerolog/log"
)

// RAWInputConfig represents configuration that can be applied on raw input
type RAWInputConfig = capture.PcapOptions

// RAWInput used for intercepting traffic for given address
type RAWInput struct {
	sync.Mutex
	config         RAWInputConfig
	messageStats   []tcp.Stats
	listener       *capture.Listener
	messageParser  *tcp.MessageParser
	cancelListener context.CancelFunc
	closed         bool

	quit  chan bool // Channel used only to indicate goroutine should shutdown
	host  string
	ports []uint16
}

// NewRAWInput constructor for RAWInput. Accepts raw input config as arguments.
func NewRAWInput(address string, config RAWInputConfig) (i *RAWInput) {
	i = new(RAWInput)
	i.config = config
	i.quit = make(chan bool)

	host, _ports, err := net.SplitHostPort(address)
	if err != nil {
		// If we are reading pcap file, no port needed
		if strings.HasSuffix(address, "pcap") {
			host = address
			_ports = "0"
			err = nil
		} else if strings.HasPrefix(address, "k8s://") {
			portIndex := strings.LastIndex(address, ":")
			host = address[:portIndex]
			_ports = address[portIndex+1:]
		} else {
			log.Fatal().Err(err).Msg("input-raw: error while parsing address")
		}
	}

	if strings.HasSuffix(host, "pcap") {
		i.config.Engine = capture.EnginePcapFile
	}

	var ports []uint16
	if _ports != "" {
		portsStr := strings.Split(_ports, ",")

		for _, portStr := range portsStr {
			port, err := strconv.Atoi(strings.TrimSpace(portStr))
			if err != nil {
				log.Fatal().Err(err).Msg("input-raw: error while parsing port")
			}
			ports = append(ports, uint16(port))

		}
	}

	i.host = host
	i.ports = ports

	i.listen(address)

	return
}

// PluginRead reads meassage from this plugin
func (i *RAWInput) PluginRead() (*plugin.Message, error) {
	var msgTCP *tcp.Message
	var msg plugin.Message
	select {
	case <-i.quit:
		return nil, plugin.ErrorStopped
	case msgTCP = <-i.listener.Messages():
		msg.Data = msgTCP.Data()
	}

	var msgType byte = proto.ResponsePayload
	if msgTCP.Direction == tcp.DirIncoming {
		msgType = proto.RequestPayload
		if i.config.RealIPHeader != "" {
			msg.Data = http_proto.SetHeader(msg.Data, []byte(i.config.RealIPHeader), []byte(msgTCP.SrcAddr))
		}
	}
	msg.Meta = proto.PayloadHeader(msgType, msgTCP.UUID(), msgTCP.Start.UnixNano(), msgTCP.End.UnixNano()-msgTCP.Start.UnixNano())

	// to be removed....
	if msgTCP.Truncated {
		log.Debug().Msg("message truncated, increase copy-buffer-size")
	}
	// to be removed...
	if msgTCP.TimedOut {
		log.Debug().Msg("message timeout reached, increase input-raw-expire")
	}
	if i.config.Stats {
		stat := msgTCP.Stats
		go i.addStats(stat)
	}
	msgTCP = nil
	return &msg, nil
}

func (i *RAWInput) listen(address string) {
	var err error
	i.listener, err = capture.NewListener(i.host, i.ports, i.config)
	if err != nil {
		log.Fatal().Err(err).Msg("error while creating listener")
	}

	err = i.listener.Activate()
	if err != nil {
		log.Fatal().Err(err).Msg("error while activating listener")
	}

	var ctx context.Context
	ctx, i.cancelListener = context.WithCancel(context.Background())
	errCh := i.listener.ListenBackground(ctx)
	<-i.listener.Reading

	log.Info().Msg(i.String())

	go func() {
		<-errCh // the listener closed voluntarily
		i.Close()
	}()
}

func (i *RAWInput) String() string {
	return fmt.Sprintf("Intercepting traffic from: %s:%s", i.host, strings.Join(strings.Fields(fmt.Sprint(i.ports)), ","))
}

// GetStats returns the stats so far and reset the stats
func (i *RAWInput) GetStats() []tcp.Stats {
	i.Lock()
	defer func() {
		i.messageStats = []tcp.Stats{}
		i.Unlock()
	}()
	return i.messageStats
}

// Close closes the input raw listener
func (i *RAWInput) Close() error {
	i.Lock()
	defer i.Unlock()
	if i.closed {
		return nil
	}
	i.cancelListener()
	close(i.quit)
	i.closed = true
	return nil
}

func (i *RAWInput) addStats(mStats tcp.Stats) {
	i.Lock()
	if len(i.messageStats) >= 10000 {
		i.messageStats = []tcp.Stats{}
	}
	i.messageStats = append(i.messageStats, mStats)
	i.Unlock()
}