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

Merge branch 'master' into pro-merge

Browse Source
This commit is contained in:
Urban Ishimwe
2020-06-13 11:39:51 +02:00
parent 540e42b9d4 a63ce611df
commit a7c5788052
13 changed files with 132 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Makefile
+2 -2
View File
@@ -64,8 +64,8 @@ bench:
$(RUN) go test $(LDFLAGS) -v -run NOT_EXISTING -bench $(BENCHMARK) -benchtime 5s
profile_test:
$(RUN) go test $(LDFLAGS) -run $(TEST) ./raw_socket_listener/. $(ARGS) -memprofile mem.mprof -cpuprofile cpu.out
$(RUN) go test $(LDFLAGS) -run $(TEST) ./raw_socket_listener/. $(ARGS) -c
$(RUN) go test $(LDFLAGS) -run $(TEST) ./capture/. $(ARGS) -memprofile mem.mprof -cpuprofile cpu.out
$(RUN) go test $(LDFLAGS) -run $(TEST) ./capture/. $(ARGS) -c
# Used mainly for debugging, because docker container do not have access to parent machine ports
run:
raw_socket_listener/listener.go → capture/listener.go
+47 -53
View File
@@ -1,21 +1,15 @@
/*
Package rawSocket provides traffic sniffier using RAW sockets.
Capture traffic from socket using RAW_SOCKET's
http://en.wikipedia.org/wiki/Raw_socket
RAW_SOCKET allow you listen for traffic on any port (e.g. sniffing) because they operate on IP level.
Package capture provides traffic sniffier using RAW sockets.
Capture traffic from socket using RAW_SOCKET's http://en.wikipedia.org/wiki/Raw_socket
RAW_SOCKET allows you to listen for traffic from any port (e.g. sniffing) because they operate on IP level.
Ports is TCP feature, same as flow control, reliable transmission and etc.
This package implements own TCP layer: TCP packets is parsed using tcp_packet.go, and flow control is managed by tcp_message.go
*/
package rawSocket
package capture
import (
"bytes"
"encoding/binary"
"fmt"
"io"
"log"
"net"
@@ -33,8 +27,6 @@ import (
"github.com/google/gopacket/pcap"
)
var _ = fmt.Println
type packet struct {
srcIP []byte
data []byte
@@ -43,7 +35,7 @@ type packet struct {
// Listener handle traffic capture
type Listener struct {
mu sync.Mutex
sync.Mutex
// buffer of TCPMessages waiting to be send
// ID -> TCPMessage
messages map[tcpID]*TCPMessage
@@ -82,8 +74,8 @@ type Listener struct {
conn net.PacketConn
pcapHandles []*pcap.Handle
quit chan bool
readyCh chan bool
quit chan bool
ready bool
protocol TCPProtocol
}
@@ -108,7 +100,6 @@ func NewListener(addr string, port string, engine int, trackResponse bool, expir
l.packetsChan = make(chan *packet, 10000)
l.messagesChan = make(chan *TCPMessage, 10000)
l.quit = make(chan bool)
l.readyCh = make(chan bool, 1)
l.messages = make(map[tcpID]*TCPMessage)
l.ackAliases = make(map[uint32]uint32)
@@ -142,6 +133,8 @@ func NewListener(addr string, port string, engine int, trackResponse bool, expir
go l.readPcap()
case EnginePcapFile:
go l.readPcapFile()
case EngineRawSocket:
go l.readRAWSocket()
default:
log.Fatal("Unknown traffic interception engine:", engine)
}
@@ -343,6 +336,7 @@ func (t *Listener) readPcap() {
go func(device pcap.Interface) {
inactive, err := pcap.NewInactiveHandle(device.Name)
if err != nil {
inactive.CleanUp()
log.Println("Pcap Error while opening device", device.Name, err)
wg.Done()
return
@@ -362,7 +356,6 @@ func (t *Listener) readPcap() {
} else {
inactive.SetSnapLen(65536)
}
inactive.SetTimeout(t.messageExpire)
inactive.SetPromisc(true)
inactive.SetImmediateMode(t.immediateMode)
@@ -382,7 +375,7 @@ func (t *Listener) readPcap() {
defer handle.Close()
t.mu.Lock()
t.Lock()
t.pcapHandles = append(t.pcapHandles, handle)
var bpfDstHost, bpfSrcHost string
@@ -428,7 +421,7 @@ func (t *Listener) readPcap() {
return
}
}
t.mu.Unlock()
t.Unlock()
var decoder gopacket.Decoder
@@ -529,7 +522,7 @@ func (t *Listener) readPcap() {
continue
}
dataOffset := (data[12] & 0xF0) >> 4
dataOffset := data[12] >> 4
isFIN := data[13]&0x01 != 0
// We need only packets with data inside
@@ -589,7 +582,9 @@ func (t *Listener) readPcap() {
}
wg.Wait()
t.readyCh <- true
t.Lock()
t.ready = true
t.Unlock()
}
func (t *Listener) readPcapFile() {
@@ -603,7 +598,9 @@ func (t *Listener) readPcapFile() {
}
}
t.readyCh <- true
t.Lock()
t.ready = true
t.Unlock()
packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
for {
@@ -643,7 +640,7 @@ func (t *Listener) readPcapFile() {
continue
}
dataOffset := (data[12] & 0xF0) >> 4
dataOffset := data[12] >> 4
isFIN := data[13]&0x01 != 0
// We need only packets with data inside
@@ -660,32 +657,41 @@ func (t *Listener) readPcapFile() {
func (t *Listener) readRAWSocket() {
conn, e := net.ListenPacket("ip:tcp", t.addr)
t.conn = conn
if e != nil {
log.Fatal(e)
}
defer t.conn.Close()
buf := make([]byte, 64*1024) // 64kb
t.readyCh <- true
type RSPacket struct {
buf []byte
addr net.Addr
err error
n int
}
var bufChan = make(chan *RSPacket, 1000)
t.Lock()
t.ready = true
t.Unlock()
go func() {
for {
// Re-allocate data object to avoid data collision
var buf [64 * 104 * 1024]byte
// Note: ReadFrom receive messages without IP header
n, addr, err := t.conn.ReadFrom(buf[:])
bufChan <- &RSPacket{buf[:], addr, err, n}
}
}()
for {
// Note: ReadFrom receive messages without IP header
n, addr, err := t.conn.ReadFrom(buf)
if err != nil {
if strings.HasSuffix(err.Error(), "closed network connection") {
packet := <-bufChan
if packet.err != nil {
if strings.HasSuffix(packet.err.Error(), "closed network connection") {
return
} else {
continue
}
continue
}
if n > 0 {
if t.isValidPacket(buf[:n]) {
t.packetsChan <- t.buildPacket([]byte(addr.(*net.IPAddr).IP), buf[:n], time.Now())
if packet.n > 0 {
if t.isValidPacket(packet.buf[:packet.n]) {
t.packetsChan <- t.buildPacket([]byte(packet.addr.(*net.IPAddr).IP), packet.buf[:packet.n], time.Now())
}
}
}
@@ -704,16 +710,13 @@ func (t *Listener) isValidPacket(buf []byte) bool {
// http://en.wikipedia.org/wiki/Transmission_Control_Protocol
destPort := binary.BigEndian.Uint16(buf[2:4])
srcPort := binary.BigEndian.Uint16(buf[0:2])
// Because RAW_SOCKET can't be bound to port, we have to control it by ourself
if destPort == t.port || (t.trackResponse && srcPort == t.port) {
// Get the 'data offset' (size of the TCP header in 32-bit words)
dataOffset := (buf[12] & 0xF0) >> 4
dataOffset := buf[12] >> 4
// We need only packets with data inside
// Check that the buffer is larger than the size of the TCP header
if len(buf) > int(dataOffset*4) {
// We should create new buffer because go slices is pointers. So buffer data shoud be immutable.
return true
}
}
@@ -884,15 +887,6 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
}
}
func (t *Listener) IsReady() bool {
select {
case <-t.readyCh:
return true
case <-time.After(5 * time.Second):
return false
}
}
// Receiver TCP messages from the listener channel
func (t *Listener) Receiver() chan *TCPMessage {
return t.messagesChan
raw_socket_listener/listener_test.go → capture/listener_test.go
+39 -45
View File
@@ -1,4 +1,4 @@
package rawSocket
package capture
import (
"bytes"
@@ -11,8 +11,7 @@ import (
func TestRawListenerInput(t *testing.T) {
var req, resp *TCPMessage
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false)
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false, false)
defer listener.Close()
reqPacket := buildPacket(true, 1, 1, []byte("GET / HTTP/1.1\r\n\r\n"), time.Now())
@@ -77,7 +76,7 @@ func responsePacket(prev *TCPPacket, payload []byte) *TCPPacket {
}
func TestHEADRequestNoBody(t *testing.T) {
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false)
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false, false)
defer listener.Close()
reqPacket := firstPacket([]byte("HEAD / HTTP/1.1\r\nContent-Length: 0\r\n\r\n"))
@@ -111,7 +110,7 @@ func TestHEADRequestNoBody(t *testing.T) {
}
func TestSingleAck100Continue(t *testing.T) {
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false)
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false, false)
defer listener.Close()
reqPacket1 := firstPacket([]byte("POST / HTTP/1.1\r\nExpect: 100-continue\r\nContent-Length: 4\r\n\r\n"))
@@ -130,7 +129,7 @@ func TestSingleAck100Continue(t *testing.T) {
}
func Test100ContinueWithoutWaiting(t *testing.T) {
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false)
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false, false)
defer listener.Close()
req1 := firstPacket([]byte("POST / HTTP/1.1\r\nExpect: 100-continue\r\nContent-Length: 4\r\n\r\n"))
@@ -146,7 +145,7 @@ func Test100ContinueWithoutWaiting(t *testing.T) {
// Client first sends data without waiting 100-continue, but once response received, generate packets based on Ack payload
func Test100ContinueMixed(t *testing.T) {
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false)
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false, false)
defer listener.Close()
req1 := firstPacket([]byte("POST / HTTP/1.1\r\nExpect: 100-continue\r\nContent-Length: 12\r\n\r\n"))
@@ -164,7 +163,7 @@ func Test100ContinueMixed(t *testing.T) {
}
func TestDoubleAck100Continue(t *testing.T) {
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false)
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false, false)
defer listener.Close()
reqPacket1 := firstPacket([]byte("POST / HTTP/1.1\r\nExpect: 100-continue\r\nContent-Length: 4\r\n\r\n"))
@@ -187,7 +186,7 @@ func TestDoubleAck100Continue(t *testing.T) {
func TestRawListenerInputResponseByClose(t *testing.T) {
var req, resp *TCPMessage
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false)
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false, false)
defer listener.Close()
reqPacket := buildPacket(true, 1, 1, []byte("GET / HTTP/1.1\r\n\r\n"), time.Now())
@@ -227,7 +226,7 @@ func TestRawListenerInputResponseByClose(t *testing.T) {
func TestRawListenerInputWithoutResponse(t *testing.T) {
var req *TCPMessage
listener := NewListener("", "0", EnginePcap, false, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false)
listener := NewListener("", "0", EnginePcap, false, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false, false)
defer listener.Close()
reqPacket := buildPacket(true, 1, 1, []byte("GET / HTTP/1.1\r\n\r\n"), time.Now())
@@ -249,7 +248,7 @@ func TestRawListenerInputWithoutResponse(t *testing.T) {
func TestRawListenerResponse(t *testing.T) {
var req, resp *TCPMessage
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false)
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false, false)
defer listener.Close()
reqPacket := firstPacket([]byte("GET / HTTP/1.1\r\n\r\n"))
@@ -297,7 +296,7 @@ func get100ContinuePackets() (req []*TCPPacket, resp []*TCPPacket) {
}
func TestShort100Continue(t *testing.T) {
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false)
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false, false)
defer listener.Close()
req, resp := get100ContinuePackets()
@@ -309,7 +308,7 @@ func TestShort100Continue(t *testing.T) {
// Response comes before Request
func Test100ContinueWrongOrder(t *testing.T) {
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false)
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false, false)
defer listener.Close()
req, resp := get100ContinuePackets()
@@ -420,14 +419,6 @@ func testChunkedSequence(t *testing.T, listener *Listener, packets ...*TCPPacket
t.Fatal("packetsChan non empty:", listener.packetsChan)
}
if len(listener.messagesChan) != 0 {
t.Fatal("messagesChan non empty:", <-listener.messagesChan)
}
if len(listener.messages) != 0 {
t.Fatal("Messages non empty:", listener.messages)
}
if len(listener.ackAliases) != 0 {
t.Fatal("ackAliases non empty:", listener.ackAliases)
}
@@ -439,30 +430,34 @@ func testChunkedSequence(t *testing.T, listener *Listener, packets ...*TCPPacket
if len(listener.respAliases) != 0 {
t.Fatal("respAliases non empty:", listener.respAliases)
}
if len(listener.respWithoutReq) != 0 {
t.Fatal("respWithoutReq non empty:", listener.respWithoutReq)
}
}
func permutation(n int, list []*TCPPacket) []*TCPPacket {
if len(list) == 1 {
return list
// permutation using heap algorithm https://en.wikipedia.org/wiki/Heap%27s_algorithm
func permutation(a []*TCPPacket, f func([]*TCPPacket)) {
n := len(a)
c := make([]int, n)
f(a)
i := 0
for i < n {
if c[i] < i {
if i&1 != 1 {
a[0], a[i] = a[i], a[0]
} else {
a[c[i]], a[i] = a[i], a[c[i]]
}
f(a)
c[i]++
i = 0
} else {
c[i] = 0
i++
}
}
k := n % len(list)
first := []*TCPPacket{list[k]}
next := make([]*TCPPacket, len(list)-1)
copy(next, append(list[:k], list[k+1:]...))
return append(first, permutation(n/len(list), next)...)
}
// Response comes before Request
func TestRawListenerChunkedWrongOrder(t *testing.T) {
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false)
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false, false)
defer listener.Close()
reqPacket1 := firstPacket([]byte("POST / HTTP/1.1\r\nTransfer-Encoding: chunked\r\nExpect: 100-continue\r\n\r\n"))
@@ -474,12 +469,11 @@ func TestRawListenerChunkedWrongOrder(t *testing.T) {
respPacket2 := responsePacket(reqPacket4, []byte("HTTP/1.1 200 OK\r\n\r\n"))
// Should re-construct message from all possible combinations
for i := 0; i < 6*5*4*3*2*1; i++ {
packets := permutation(i, []*TCPPacket{reqPacket1, reqPacket2, reqPacket3, reqPacket4, respPacket1, respPacket2})
testChunkedSequence(t, listener, packets...)
f := func(p []*TCPPacket) {
testChunkedSequence(t, listener, p...)
}
// Should re-construct message from all possible combinations
permutation([]*TCPPacket{reqPacket1, reqPacket2, reqPacket3, reqPacket4, respPacket1, respPacket2}, f)
}
func chunkedPostMessage() []*TCPPacket {
@@ -532,7 +526,7 @@ func getMessage() []*TCPPacket {
// Response comes before Request
func TestRawListenerBench(t *testing.T) {
l := NewListener("", "0", EnginePcap, true, 200*time.Millisecond, ProtocolHTTP, "", "", 0, false)
l := NewListener("", "0", EnginePcap, true, 200*time.Millisecond, ProtocolHTTP, "", "", 0, false, false)
defer l.Close()
// Should re-construct message from all possible combinations
@@ -583,7 +577,7 @@ func TestRawListenerBench(t *testing.T) {
func TestResponseZeroContentLength(t *testing.T) {
var req, resp *TCPMessage
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false)
listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP, "", "", 0, false, false)
defer listener.Close()
reqPacket := firstPacket([]byte("POST /api/setup/install HTTP/1.1\r\nHost: localhost:22936\r\nUser-Agent: curl/7.57.0\r\nAccept: */*\r\nContent-Length: 0\r\nContent-Type: application/x-www-form-urlencoded\r\n\r\n"))
raw_socket_listener/tcp_message.go → capture/tcp_message.go
+5 -2
View File
@@ -1,4 +1,4 @@
package rawSocket
package capture
import (
"bytes"
@@ -16,10 +16,13 @@ import (
var _ = log.Println
// TCPProtocol is a number to indicate type of protocol
type TCPProtocol uint8
const (
ProtocolHTTP TCPProtocol = 0
// ProtocolHTTP ...
ProtocolHTTP TCPProtocol = 0
// ProtocolBinary ...
ProtocolBinary TCPProtocol = 1
)
raw_socket_listener/tcp_message_test.go → capture/tcp_message_test.go
+2 -3
View File
@@ -1,4 +1,4 @@
package rawSocket
package capture
import (
"bytes"
@@ -36,8 +36,7 @@ func buildMessage(p *TCPPacket) *TCPMessage {
if p.SrcPort == 1 {
isIncoming = true
}
m := NewTCPMessage(p.Seq, p.Ack, isIncoming, ProtocolHTTP)
m := NewTCPMessage(p.Seq, p.Ack, isIncoming, ProtocolHTTP, "")
m.AddPacket(p)
return m
raw_socket_listener/tcp_packet.go → capture/tcp_packet.go
+1 -5
View File
@@ -1,15 +1,12 @@
package rawSocket
package capture
import (
"encoding/binary"
"log"
"strconv"
"strings"
"time"
)
var _ = log.Println
// TCP Flags
const (
fFIN = 1 << iota
@@ -102,7 +99,6 @@ func (t *TCPPacket) dump() *packet {
}
copy(packetData[16:], t.Data)
return &packet{
srcIP: packetSrcIP,
data: packetData,
emitter_test.go
+1 -1
View File
@@ -127,7 +127,7 @@ func TestEmitterSplitRoundRobin(t *testing.T) {
close(quit)
if counter1 == 0 || counter2 == 0 {
if counter1 == 0 || counter2 == 0 || counter1 != counter2 {
t.Errorf("Round robin should split traffic equally: %d vs %d", counter1, counter2)
}
input_raw.go
+2 -9
View File
@@ -5,8 +5,8 @@ import (
"net"
"time"
raw "github.com/buger/goreplay/capture"
"github.com/buger/goreplay/proto"
raw "github.com/buger/goreplay/raw_socket_listener"
)
// RAWInput used for intercepting traffic for given address
@@ -59,7 +59,6 @@ func NewRAWInput(address string, engine int, trackResponse bool, expire time.Dur
}
i.listen(address)
i.listener.IsReady()
return
}
@@ -89,7 +88,6 @@ func (i *RAWInput) listen(address string) {
Debug("Listening for traffic on: " + address)
host, port, err := net.SplitHostPort(address)
if err != nil {
log.Fatalf("input-raw: error while parsing address: %s", err)
}
@@ -103,13 +101,8 @@ func (i *RAWInput) listen(address string) {
select {
case <-i.quit:
return
default:
case i.data <- <-ch: // Receiving TCPMessage object
}
// Receiving TCPMessage object
m := <-ch
i.data <- m
}
}()
}
output_s3.go
+14 -12
View File
@@ -3,16 +3,17 @@ package main
import (
_ "bufio"
"fmt"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/session"
"github.com/aws/aws-sdk-go/service/s3"
_ "github.com/aws/aws-sdk-go/service/s3/s3manager"
_ "io"
"log"
"math/rand"
"os"
"path/filepath"
"strings"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/session"
"github.com/aws/aws-sdk-go/service/s3"
_ "github.com/aws/aws-sdk-go/service/s3/s3manager"
)
// S3Output output plugin
@@ -22,7 +23,7 @@ type S3Output struct {
buffer *FileOutput
session *session.Session
config *FileOutputConfig
closeC chan struct{}
closeCh chan struct{}
}
// NewS3Output constructor for FileOutput, accepts path
@@ -42,18 +43,18 @@ func NewS3Output(pathTemplate string, config *FileOutputConfig) *S3Output {
}
rnd := rand.Int63()
buffer_name := fmt.Sprintf("gor_output_s3_%d_buf_", rnd)
bufferName := fmt.Sprintf("gor_output_s3_%d_buf_", rnd)
pathParts := strings.Split(pathTemplate, "/")
buffer_name += pathParts[len(pathParts)-1]
bufferName += pathParts[len(pathParts)-1]
if strings.HasSuffix(o.pathTemplate, ".gz") {
buffer_name += ".gz"
bufferName += ".gz"
}
buffer_path := filepath.Join(config.bufferPath, buffer_name)
bufferPath := filepath.Join(config.bufferPath, bufferName)
o.buffer = NewFileOutput(buffer_path, config)
o.buffer = NewFileOutput(bufferPath, config)
o.connect()
return o
@@ -74,6 +75,7 @@ func (o *S3Output) String() string {
return "S3 output: " + o.pathTemplate
}
// Close close the buffer of the S3 connection
func (o *S3Output) Close() error {
return o.buffer.Close()
}
@@ -121,7 +123,7 @@ func (o *S3Output) onBufferUpdate(path string) {
os.Remove(path)
if o.closeC != nil {
o.closeC <- struct{}{}
if o.closeCh != nil {
o.closeCh <- struct{}{}
}
}
pro.go
+2
View File
@@ -1,3 +1,5 @@
package main
// PRO this value indicates if goreplay is running in PRO mode..
// it must not be modified explicitly in production
var PRO = false
s3_reader.go
+4
View File
@@ -12,6 +12,7 @@ import (
"github.com/aws/aws-sdk-go/service/s3"
)
// S3ReadCloser ...
type S3ReadCloser struct {
bucket string
key string
@@ -49,6 +50,7 @@ func awsConfig() *aws.Config {
return config
}
// NewS3ReadCloser returns new instance of S3 read closer
func NewS3ReadCloser(path string) *S3ReadCloser {
if !PRO {
log.Fatal("Using S3 input and output require PRO license")
@@ -68,6 +70,7 @@ func NewS3ReadCloser(path string) *S3ReadCloser {
}
}
// Read reads buffer from s3 session
func (s *S3ReadCloser) Read(b []byte) (n int, e error) {
if s.readBytes == 0 || s.readBytes+len(b) > s.offset {
svc := s3.New(s.sess)
@@ -96,6 +99,7 @@ func (s *S3ReadCloser) Read(b []byte) (n int, e error) {
return s.buf.Read(b)
}
// Close is here to make S3ReadCloser satisfy ReadCloser interface
func (s *S3ReadCloser) Close() error {
return nil
}
s3_test.go
+7 -6
View File
@@ -4,13 +4,14 @@ package main
import (
"fmt"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/s3"
"math/rand"
"os"
"path/filepath"
"testing"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/s3"
)
func TestS3Output(t *testing.T) {
@@ -63,7 +64,7 @@ func TestS3OutputQueueLimit(t *testing.T) {
path := fmt.Sprintf("s3://test-gor/%d/requests.gz", rnd)
output := NewS3Output(path, &FileOutputConfig{queueLimit: 100})
output.closeC = make(chan struct{}, 3)
output.closeCh = make(chan struct{}, 3)
svc := s3.New(output.session)
@@ -77,7 +78,7 @@ func TestS3OutputQueueLimit(t *testing.T) {
output.Write([]byte("1 1 1\ntest"))
for i := 0; i < 3; i++ {
<-output.closeC
<-output.closeCh
}
params := &s3.ListObjectsInput{
@@ -110,7 +111,7 @@ func TestInputFileFromS3(t *testing.T) {
path := fmt.Sprintf("s3://test-gor-eu/%d/requests.gz", rnd)
output := NewS3Output(path, &FileOutputConfig{queueLimit: 5000})
output.closeC = make(chan struct{}, 10)
output.closeCh = make(chan struct{}, 10)
for i := 0; i <= 20000; i++ {
output.Write([]byte("1 1 1\ntest"))
@@ -123,7 +124,7 @@ func TestInputFileFromS3(t *testing.T) {
output.Write([]byte("1 1 1\ntest"))
for i := 0; i < 2; i++ {
<-output.closeC
<-output.closeCh
}
input := NewFileInput(fmt.Sprintf("s3://test-gor-eu/%d", rnd), false)
tcp_client.go
+6
View File
@@ -10,6 +10,7 @@ import (
"time"
)
// TCPClientConfig client configuration
type TCPClientConfig struct {
Debug bool
ConnectionTimeout time.Duration
@@ -18,6 +19,7 @@ type TCPClientConfig struct {
Secure bool
}
// TCPClient client connection properties
type TCPClient struct {
baseURL string
addr string
@@ -27,6 +29,7 @@ type TCPClient struct {
redirectsCount int
}
// NewTCPClient returns new TCPClient
func NewTCPClient(addr string, config *TCPClientConfig) *TCPClient {
if config.Timeout.Nanoseconds() == 0 {
config.Timeout = 5 * time.Second
@@ -44,6 +47,7 @@ func NewTCPClient(addr string, config *TCPClientConfig) *TCPClient {
return client
}
// Connect creates a tcp connection of the client
func (c *TCPClient) Connect() (err error) {
c.Disconnect()
@@ -62,6 +66,7 @@ func (c *TCPClient) Connect() (err error) {
return
}
// Disconnect closes the client connection
func (c *TCPClient) Disconnect() {
if c.conn != nil {
c.conn.Close()
@@ -92,6 +97,7 @@ func (c *TCPClient) isAlive() bool {
return true
}
// Send sends data over created tcp connection
func (c *TCPClient) Send(data []byte) (response []byte, err error) {
// Don't exit on panic
defer func() {