mirror of
https://github.com/netsampler/goflow2.git
synced 2024-05-06 15:54:52 +00:00
e52a053c9c
* Improve format registration * less interleaving (json does not require importing protobuf) * generic text renderer * escape for strings in text/json formatter
47 lines
991 B
Go
47 lines
991 B
Go
package protobuf
|
|
|
|
import (
|
|
"context"
|
|
"flag"
|
|
"fmt"
|
|
"github.com/golang/protobuf/proto"
|
|
"github.com/netsampler/goflow2/format"
|
|
"github.com/netsampler/goflow2/format/common"
|
|
)
|
|
|
|
type ProtobufDriver struct {
|
|
fixedLen bool
|
|
}
|
|
|
|
func (d *ProtobufDriver) Prepare() error {
|
|
common.HashFlag()
|
|
flag.BoolVar(&d.fixedLen, "format.protobuf.fixedlen", false, "Prefix the protobuf with message length")
|
|
return nil
|
|
}
|
|
|
|
func (d *ProtobufDriver) Init(context.Context) error {
|
|
return common.ManualHashInit()
|
|
}
|
|
|
|
func (d *ProtobufDriver) Format(data interface{}) ([]byte, []byte, error) {
|
|
msg, ok := data.(proto.Message)
|
|
if !ok {
|
|
return nil, nil, fmt.Errorf("message is not protobuf")
|
|
}
|
|
key := common.HashProtoLocal(msg)
|
|
|
|
if !d.fixedLen {
|
|
b, err := proto.Marshal(msg)
|
|
return []byte(key), b, err
|
|
} else {
|
|
buf := proto.NewBuffer([]byte{})
|
|
err := buf.EncodeMessage(msg)
|
|
return []byte(key), buf.Bytes(), err
|
|
}
|
|
}
|
|
|
|
func init() {
|
|
d := &ProtobufDriver{}
|
|
format.RegisterFormatDriver("pb", d)
|
|
}
|