Merge pull request #30 from netsampler/bug/enricher-decode

bugfix: enricher decoding certain protobuf

README.md

@@ -181,7 +181,7 @@ with a database for Autonomous System Number and Country.
 Similar output options as GoFlow are provided.
 
 ```bash
-$ ./goflow2 -format=pb | ./enricher -db.asn path-to/GeoLite2-ASN.mmdb -db.country path-to/GeoLite2-Country.mmdb
+$ ./goflow2 -transport.file.sep= -format=pb -format.protobuf.fixedlen=true | ./enricher -db.asn path-to/GeoLite2-ASN.mmdb -db.country path-to/GeoLite2-Country.mmdb
 ```
 
 For a more scalable production setting, Kafka and protobuf are recommended.

cmd/enricher/main.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"bytes"
 	"context"
+	"encoding/binary"
 	"flag"
 	"fmt"
 	"io"
@@ -146,13 +147,28 @@ func main() {
 	rdr := bufio.NewReader(os.Stdin)
 
 	msg := &flowmessage.FlowMessageExt{}
+	msgLen := make([]byte, binary.MaxVarintLen64)
+	lenBufSize := len(msgLen)
 	for {
-		line, err := rdr.ReadBytes('\n')
+		n, err := rdr.Read(msgLen)
 		if err != nil && err != io.EOF {
 			log.Error(err)
 			continue
 		}
-		if len(line) == 0 {
+
+		len, vn := proto.DecodeVarint(msgLen[0:n])
+		if len == 0 {
+			continue
+		}
+
+		line := make([]byte, len)
+		if vn < lenBufSize {
+			copy(line[0:lenBufSize-vn], msgLen[vn:lenBufSize])
+		}
+
+		n, err = io.ReadFull(rdr, line[lenBufSize-vn:])
+		if err != nil && err != io.EOF {
+			log.Error(err)
 			continue
 		}
 		line = bytes.TrimSuffix(line, []byte("\n"))

transport/file/transport.go

@@ -14,6 +14,7 @@ import (
 
 type FileDriver struct {
 	fileDestination string
+	lineSeparator   string
 	w               io.Writer
 	file            *os.File
 	lock            *sync.RWMutex
@@ -22,6 +23,7 @@ type FileDriver struct {
 
 func (d *FileDriver) Prepare() error {
 	flag.StringVar(&d.fileDestination, "transport.file", "", "File/console output (empty for stdout)")
+	flag.StringVar(&d.lineSeparator, "transport.file.sep", "\n", "Line separator")
 	// idea: add terminal coloring based on key partitioning (if any)
 	return nil
 }
@@ -76,7 +78,7 @@ func (d *FileDriver) Send(key, data []byte) error {
 	d.lock.RLock()
 	w := d.w
 	d.lock.RUnlock()
-	_, err := fmt.Fprintln(w, string(data))
+	_, err := fmt.Fprint(w, string(data)+d.lineSeparator)
 	return err
 }