github/goreplay
4
0
Fork 0
mirror of https://github.com/buger/goreplay.git synced 2024-04-21 12:32:02 +00:00
Code Issues Releases Wiki Activity

Ensure that raw input gets closed

Browse Source
This commit is contained in:
Leonid Bugaev
2015-08-20 08:56:10 +03:00
parent 33cdcd4d13
commit 6974a6e278
5 changed files with 75 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
emitter.go
+12 -1
View File
@@ -32,8 +32,19 @@ func Start(stop chan int) {
for {
select {
case <-stop:
for _, in := range Plugins.Inputs {
if c, ok := in.(io.Closer); ok {
c.Close()
}
}
for _, out := range Plugins.Outputs {
if c, ok := out.(io.Closer); ok {
c.Close()
}
}
return
case <-time.After(1 * time.Second):
case <-time.After(time.Second):
}
}
}
input_raw.go
+16 -2
View File
@@ -13,6 +13,8 @@ type RAWInput struct {
data chan *raw.TCPMessage
address string
expire time.Duration
quit chan bool
listener *raw.Listener
}
// NewRAWInput constructor for RAWInput. Accepts address with port as argument.
@@ -21,6 +23,7 @@ func NewRAWInput(address string, expire time.Duration) (i *RAWInput) {
i.data = make(chan *raw.TCPMessage)
i.address = address
i.expire = expire
i.quit = make(chan bool)
go i.listen(address)
@@ -56,11 +59,17 @@ func (i *RAWInput) listen(address string) {
log.Fatal("input-raw: error while parsing address", err)
}
listener := raw.NewListener(host, port, i.expire, true)
i.listener = raw.NewListener(host, port, i.expire, true)
for {
select {
case <-i.quit:
return
default:
}
// Receiving TCPMessage object
m := listener.Receive()
m := i.listener.Receive()
i.data <- m
}
@@ -69,3 +78,8 @@ func (i *RAWInput) listen(address string) {
func (i *RAWInput) String() string {
return "RAW Socket input: " + i.address
}
func (i *RAWInput) Close() {
i.listener.Close()
close(i.quit)
}
input_raw_test.go
+8 -3
View File
@@ -28,8 +28,13 @@ func TestRAWInput(t *testing.T) {
var respCounter, reqCounter int64
defer func(){
log.Println(reqCounter, respCounter)
}()
input := NewRAWInput(originAddr, testRawExpire)
defer input.Close()
output := NewTestOutput(func(data []byte) {
if data[0] == '1' {
atomic.AddInt64(&reqCounter, 1)
@@ -37,8 +42,6 @@ func TestRAWInput(t *testing.T) {
atomic.AddInt64(&respCounter, 1)
}
log.Println(reqCounter, respCounter)
wg.Done()
})
@@ -58,7 +61,6 @@ func TestRAWInput(t *testing.T) {
}
wg.Wait()
close(quit)
}
@@ -80,6 +82,7 @@ func TestInputRAW100Expect(t *testing.T) {
originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
input := NewRAWInput(originAddr, testRawExpire)
defer input.Close()
// We will use it to get content of raw HTTP request
testOutput := NewTestOutput(func(data []byte) {
@@ -142,6 +145,7 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
input := NewRAWInput(originAddr, testRawExpire)
defer input.Close()
replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
defer r.Body.Close()
@@ -200,6 +204,7 @@ func TestInputRAWLargePayload(t *testing.T) {
originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
input := NewRAWInput(originAddr, testRawExpire)
defer input.Close()
replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
req.Body = http.MaxBytesReader(w, req.Body, 1*1024*1024)
middleware_test.go
+2
View File
@@ -117,6 +117,7 @@ func TestEchoMiddleware(t *testing.T) {
// Catch traffic from one service
input := NewRAWInput(from.Listener.Addr().String(), testRawExpire)
defer input.Close()
// And redirect to another
output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: false})
@@ -174,6 +175,7 @@ func TestTokenMiddleware(t *testing.T) {
fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
// Catch traffic from one service
input := NewRAWInput(fromAddr, testRawExpire)
defer input.Close()
// And redirect to another
output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: true})
raw_socket_listener/listener.go
+37 -18
View File
@@ -20,6 +20,7 @@ import (
"runtime/debug"
"strconv"
"time"
"strings"
)
// Listener handle traffic capture
@@ -50,6 +51,9 @@ type Listener struct {
messageExpire time.Duration
captureResponse bool
conn net.PacketConn
quit chan bool
}
type request struct {
@@ -58,30 +62,31 @@ type request struct {
}
// NewListener creates and initializes new Listener object
func NewListener(addr string, port string, expire time.Duration, captureResponse bool) (rawListener *Listener) {
rawListener = &Listener{captureResponse: captureResponse}
func NewListener(addr string, port string, expire time.Duration, captureResponse bool) (l *Listener) {
l = &Listener{captureResponse: captureResponse}
rawListener.packetsChan = make(chan *TCPPacket, 10000)
rawListener.messagesChan = make(chan *TCPMessage, 10000)
rawListener.messageDelChan = make(chan *TCPMessage, 10000)
l.packetsChan = make(chan *TCPPacket, 10000)
l.messagesChan = make(chan *TCPMessage, 10000)
l.messageDelChan = make(chan *TCPMessage, 10000)
l.quit = make(chan bool)
rawListener.messages = make(map[string]*TCPMessage)
rawListener.ackAliases = make(map[uint32]uint32)
rawListener.seqWithData = make(map[uint32]uint32)
rawListener.respAliases = make(map[uint32]*request)
l.messages = make(map[string]*TCPMessage)
l.ackAliases = make(map[uint32]uint32)
l.seqWithData = make(map[uint32]uint32)
l.respAliases = make(map[uint32]*request)
rawListener.addr = addr
l.addr = addr
_port, _ := strconv.Atoi(port)
rawListener.port = uint16(_port)
l.port = uint16(_port)
if expire.Nanoseconds() == 0 {
expire = 2000 * time.Millisecond
}
rawListener.messageExpire = expire
l.messageExpire = expire
go rawListener.listen()
go rawListener.readRAWSocket()
go l.listen()
go l.readRAWSocket()
return
}
@@ -89,6 +94,9 @@ func NewListener(addr string, port string, expire time.Duration, captureResponse
func (t *Listener) listen() {
for {
select {
case <-t.quit:
t.conn.Close()
return
// If message ready for deletion it means that its also complete or expired by timeout
case message := <-t.messageDelChan:
delete(t.ackAliases, message.Ack)
@@ -109,21 +117,26 @@ func (t *Listener) listen() {
}
func (t *Listener) readRAWSocket() {
conn, e := net.ListenPacket("ip4:tcp", t.addr)
t.conn = conn
if e != nil {
log.Fatal(e)
}
defer conn.Close()
defer t.conn.Close()
for {
buf := make([]byte, 64*1024) // 64kb
// Note: ReadFrom receive messages without IP header
n, addr, err := conn.ReadFrom(buf)
n, addr, err := t.conn.ReadFrom(buf)
if err != nil {
log.Println("Error:", err)
continue
if strings.HasSuffix(err.Error(), "closed network connection") {
return
} else {
log.Println("Raw listener error:", err)
continue
}
}
if n > 0 {
@@ -240,3 +253,9 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
func (t *Listener) Receive() *TCPMessage {
return <-t.messagesChan
}
func (t *Listener) Close() {
close(t.quit)
t.conn.Close()
return
}