diff --git a/parser.go b/parser.go index 546710d..028959d 100644 --- a/parser.go +++ b/parser.go @@ -24,6 +24,10 @@ var ErrRead = errors.New("short read - malformed packet?") // error are encountered. var ErrProtocol = errors.New("protocol error") +// ErrUnknownTemplate is returned when a message's data set refers to a template +// which has not yet been seen by the parser. +var ErrUnknownTemplate = errors.New("unknown template") + // A Message is the top level construct representing an IPFIX message. A well // formed message contains one or more sets of data or template information. type Message struct { @@ -491,6 +495,28 @@ func calcMinRecLen(tpl []TemplateFieldSpecifier) uint16 { return minLen } +// LookupTemplateRecords returns the list of template records referred to by the data sets +// in the given message. It also includes any templates already extant in the message. +func (s *Session) LookupTemplateRecords(m Message) ([]TemplateRecord, error) { + tr := m.TemplateRecords +dataLoop: + for _, dr := range m.DataRecords { + // First walk and make sure this template isn't already in the return list + for _, t := range tr { + if t.TemplateID == dr.TemplateID { + continue dataLoop + } + } + // If we got this far, we haven't seen the template ID yet + tfs := s.lookupTemplateFieldSpecifiers(dr.TemplateID) + if len(tfs) == 0 { + return nil, ErrUnknownTemplate + } + tr = append(tr, TemplateRecord{TemplateID: dr.TemplateID, FieldSpecifiers: tfs}) + } + return tr, nil +} + func (s *Session) lookupTemplateFieldSpecifiers(tid uint16) []TemplateFieldSpecifier { var tpl []TemplateFieldSpecifier @@ -586,3 +612,165 @@ func (s *Session) LoadTemplateRecords(trecs []TemplateRecord) { s.registerTemplateRecord(&tr) } } + +// Returns overall length of the message, the length of the template set, and +// the length of the data set(s). +func (s *Session) calculateMarshalledLength(m Message) (int, int, int, error) { + var length int + var tmplLen, dataLen int + length += msgHeaderLength // there will always be a header + if len(m.TemplateRecords) > 0 { + // We will be creating a template set + tmplLen += setHeaderLength + for _, rec := range m.TemplateRecords { + // Each template record implies a record header + tmplLen += 4 + for _, field := range rec.FieldSpecifiers { + if field.EnterpriseID == 0 { + tmplLen += 4 + } else { + tmplLen += 8 + } + } + } + } + if len(m.DataRecords) > 0 { + currentTemplate := m.DataRecords[0].TemplateID + tpl := s.lookupTemplateFieldSpecifiers(currentTemplate) + dataLen += setHeaderLength + for _, dr := range m.DataRecords { + if dr.TemplateID != currentTemplate { + dataLen += setHeaderLength + tpl = s.lookupTemplateFieldSpecifiers(dr.TemplateID) + currentTemplate = dr.TemplateID + } + for i, field := range dr.Fields { + if i > len(tpl) { + return 0, 0, 0, errors.New("too many fields for template") + } + // Handle variable-length fields + if tpl[i].Length == 0xffff { + if len(field) < 0xff { + dataLen += 1 // 1 byte for the length + dataLen += len(field) // and then the field itself + } else { + dataLen += 3 // 3 bytes for the length + dataLen += len(field) // and then the field itself + } + } else { + dataLen += int(tpl[i].Length) + } + } + } + } + length += tmplLen + length += dataLen + return length, tmplLen, dataLen, nil +} + +// Marshall a Message struct back into a raw IPFIX buffer +func (s *Session) Marshal(m Message) ([]byte, error) { + // First we'll calculate how big the message will be + length, tmplLen, _, err := s.calculateMarshalledLength(m) + if err != nil { + return []byte{}, nil + } + + // Now make the empty buffer (brings us down to 1 allocation per call to Marshal) + message := make([]byte, length) + + // Populate the header + binary.BigEndian.PutUint16(message[:2], m.Header.Version) + binary.BigEndian.PutUint32(message[4:8], m.Header.ExportTime) + binary.BigEndian.PutUint32(message[8:12], m.Header.SequenceNumber) + binary.BigEndian.PutUint32(message[12:16], m.Header.DomainID) + binary.BigEndian.PutUint16(message[2:4], uint16(length)) + offset := msgHeaderLength + + // Build template records set + if len(m.TemplateRecords) > 0 { + // construct template set header + binary.BigEndian.PutUint16(message[offset:offset+2], 2) // type is always 2 for templates + binary.BigEndian.PutUint16(message[offset+2:offset+4], uint16(tmplLen)) // we calculated earlier + offset += 4 + + for _, rec := range m.TemplateRecords { + // Build the record header (template ID + number of field specifiers) + binary.BigEndian.PutUint16(message[offset:offset+2], rec.TemplateID) + binary.BigEndian.PutUint16(message[offset+2:offset+4], uint16(len(rec.FieldSpecifiers))) + offset += 4 + // Now build out the fields + for _, field := range rec.FieldSpecifiers { + if field.EnterpriseID == 0 { + // No enterprise needed + binary.BigEndian.PutUint16(message[offset:offset+2], field.FieldID) + binary.BigEndian.PutUint16(message[offset+2:offset+4], field.Length) + offset += 4 + } else { + binary.BigEndian.PutUint16(message[offset:offset+2], field.FieldID+0x8000) + binary.BigEndian.PutUint16(message[offset+2:offset+4], field.Length) + binary.BigEndian.PutUint32(message[offset+4:offset+8], field.EnterpriseID) + offset += 8 + } + } + } + } + + // Build data record section + // It's possible that there were multiple sets with alternating templates, + // e.g. set 0 used template 256, set 1 used template 300, set 2 used 256 again. + if len(m.DataRecords) > 0 { + currentTemplate := m.DataRecords[0].TemplateID + tpl := s.lookupTemplateFieldSpecifiers(currentTemplate) + setStart := offset // where the current set started + // initialize the first set header + // We only write the template ID now, since we won't be sure of length until later + binary.BigEndian.PutUint16(message[offset:offset+2], currentTemplate) + offset += 4 + for _, dr := range m.DataRecords { + if dr.TemplateID != currentTemplate { + // We've transitioned templates, make a new set + if offset-setStart > setHeaderLength { + binary.BigEndian.PutUint16(message[setStart+2:setStart+4], uint16(offset-setStart)) + } + // now set up for the next set + setStart = offset + currentTemplate = dr.TemplateID + tpl = s.lookupTemplateFieldSpecifiers(dr.TemplateID) + // We only write the template ID now, since we won't be sure of length until later + binary.BigEndian.PutUint16(message[offset:offset+2], dr.TemplateID) + offset += 4 + } + for i, field := range dr.Fields { + if i > len(tpl) { + return message, errors.New("too many fields for template") + } + // Handle variable-length fields + if tpl[i].Length == 0xffff { + if len(field) < 0xff { + message[offset] = uint8(len(field)) + offset++ + copy(message[offset:offset+len(field)], field) + offset += len(field) + } else { + message[offset] = 0xff + binary.BigEndian.PutUint16(message[offset+1:offset+3], uint16(len(field))) + offset += 3 + copy(message[offset:offset+len(field)], field) + offset += len(field) + } + } else { + copy(message[offset:offset+len(field)], field) + offset += len(field) + } + } + } + // Emit the final set length + if offset-setStart > setHeaderLength { + binary.BigEndian.PutUint16(message[setStart+2:setStart+4], uint16(offset-setStart)) + } + + } + + return message, nil +} diff --git a/parser_test.go b/parser_test.go index d199700..b547556 100644 --- a/parser_test.go +++ b/parser_test.go @@ -45,6 +45,48 @@ func TestParseTemplateSet(t *testing.T) { } } +func TestMarshalTemplateSet(t *testing.T) { + packet, _ := hex.DecodeString("000a008c51ec4264000000000b20bdbe0002007c283b0008001c0010800c000400003c258003000800003c258004000800003c258012ffff00003c258001ffff00003c25801cffff00003c25001b0010c2ac0008000c0004800c000400003c258003000800003c258004000800003c258012ffff00003c258001ffff00003c25801cffff00003c2500080004") + p := ipfix.NewSession() + + r := bytes.NewBuffer(packet) + msg, err := p.ParseReader(r) + if err != nil { + t.Fatal("ParseReader failed", err) + } + + marshalled, err := p.Marshal(msg) + if err != nil { + t.Fatal("Marshal failed", err) + } + + if !bytes.Equal(marshalled, packet) { + t.Fatal("Expected didn't match marshalled bytes") + } +} + +func TestMarshalDataSet(t *testing.T) { + p0, _ := hex.DecodeString("000a008c51ec4264000000000b20bdbe0002007c283b0008001c0010800c000400003c258003000800003c258004000800003c258012ffff00003c258001ffff00003c25801cffff00003c25001b0010c2ac0008000c0004800c000400003c258003000800003c258004000800003c258012ffff00003c258001ffff00003c25801cffff00003c2500080004") + p1, _ := hex.DecodeString("000a05b051ec4270000000000b20bdbec2ac05a0ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043000116fcb8ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b525043005e489f46ac10200300000026000000000000019f0000000000000160000e4265696e6720616e616c797a656400c27ef905ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043007aa7519c0808080800000000000000000000008d00000000000000550003444e5300ac102082ac10200f0000000000000000000000940000000000000147000f426974546f7272656e74204b52504300b228265c1859c1570000000000000000000000000000000000000064000f426974546f7272656e74204b52504300ac10200fac10200f0000000000000000000000920000000000000145000f426974546f7272656e74204b525043007b75a68ad92bb37f00000000000000000000006e0000000000000064000f426974546f7272656e74204b52504300ac10200fac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043004f972c247449d8f200000000000000000000006e0000000000000064000f426974546f7272656e74204b52504300ac10200fac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b5250430048b682a4ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b52504300595cc40dac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b5250430057451cc1ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b525043005465e5a8ac1020ff00000000000000000000000000000000000000af001a44726f70626f78204c414e2073796e6320646973636f766572790764726f70626f78ac102013ac10200f00000000000000000000008f000000000000014b000f426974546f7272656e74204b5250430001ab3c06ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b52504300befcacc8ffffffff00000000000000000000000000000000000000af001a44726f70626f78204c414e2073796e6320646973636f766572790764726f70626f78ac102013ac10200300000025000000000000019e0000000000000167000e4265696e6720616e616c797a656400c27ef905ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043006ca28bcdac10200f000000000000000000000091000000000000011c000f426974546f7272656e74204b52504300b13531caac10200f000000000000000000000068000000000000005f000f426974546f7272656e74204b5250430053df9212ac10200f0000000000000000000000940000000000000159000f426974546f7272656e74204b525043005f43f0b2ac10200f0000000000000000000001220000000000000252000f426974546f7272656e74204b52504300567ce6fbac10200100000000000000000000005a000000000000005a00034e545000ac102080ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b5250430055550ef7ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b52504300ba9322a2ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043004579e7114b01bf5300000000000000000000006e0000000000000064000f426974546f7272656e74204b52504300ac10200fac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043005cf46adf") + p := ipfix.NewSession() + + p.ParseBuffer(p0) + + msg, err := p.ParseBuffer(p1) + if err != nil { + t.Fatal("ParseReader failed", err) + } + + marshalled, err := p.Marshal(msg) + if err != nil { + t.Fatal("Marshal failed", err) + } + + if !bytes.Equal(marshalled, p1) { + t.Fatalf("Expected (%v bytes) didn't match marshalled (%v bytes)", len(p1), len(marshalled)) + } +} + func TestParseTemplateIDAliasing(t *testing.T) { packet, _ := hex.DecodeString("000a017c51ec4264000000000b20bdbe0002016c283b0008001c0010800c000400003c258003000800003c258004000800003c258012ffff00003c258001ffff00003c25801cffff00003c25001b0010c2ac0008000c0004800c000400003c258003000800003c258004000800003c258012ffff00003c258001ffff00003c25801cffff00003c250008000412340008001c0010800c000400003c258003000800003c258004000800003c258012ffff00003c258001ffff00003c25801cffff00003c25001b0010abcd0008000c0004800c000400003c258003000800003c258004000800003c258012ffff00003c258001ffff00003c25801cffff00003c250008000412340008001c0010800c000400003c258003000800003c258004000800003c258012ffff00003c258001ffff00003c25801cffff00003c25001b0010abcd0008000c0004800c000400003c258003000800003c258004000800003c258012ffff00003c258001ffff00003c25801cffff00003c2500080004") p := ipfix.NewSession(ipfix.WithIDAliasing(true)) @@ -334,6 +376,29 @@ func BenchmarkParseBuffer(b *testing.B) { } } +func BenchmarkMarshal(b *testing.B) { + p0, _ := hex.DecodeString("000a008c51ec4264000000000b20bdbe0002007c283b0008001c0010800c000400003c258003000800003c258004000800003c258012ffff00003c258001ffff00003c25801cffff00003c25001b0010c2ac0008000c0004800c000400003c258003000800003c258004000800003c258012ffff00003c258001ffff00003c25801cffff00003c2500080004") + p1, _ := hex.DecodeString("000a05b051ec4270000000000b20bdbec2ac05a0ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043000116fcb8ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b525043005e489f46ac10200300000026000000000000019f0000000000000160000e4265696e6720616e616c797a656400c27ef905ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043007aa7519c0808080800000000000000000000008d00000000000000550003444e5300ac102082ac10200f0000000000000000000000940000000000000147000f426974546f7272656e74204b52504300b228265c1859c1570000000000000000000000000000000000000064000f426974546f7272656e74204b52504300ac10200fac10200f0000000000000000000000920000000000000145000f426974546f7272656e74204b525043007b75a68ad92bb37f00000000000000000000006e0000000000000064000f426974546f7272656e74204b52504300ac10200fac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043004f972c247449d8f200000000000000000000006e0000000000000064000f426974546f7272656e74204b52504300ac10200fac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b5250430048b682a4ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b52504300595cc40dac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b5250430057451cc1ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b525043005465e5a8ac1020ff00000000000000000000000000000000000000af001a44726f70626f78204c414e2073796e6320646973636f766572790764726f70626f78ac102013ac10200f00000000000000000000008f000000000000014b000f426974546f7272656e74204b5250430001ab3c06ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b52504300befcacc8ffffffff00000000000000000000000000000000000000af001a44726f70626f78204c414e2073796e6320646973636f766572790764726f70626f78ac102013ac10200300000025000000000000019e0000000000000167000e4265696e6720616e616c797a656400c27ef905ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043006ca28bcdac10200f000000000000000000000091000000000000011c000f426974546f7272656e74204b52504300b13531caac10200f000000000000000000000068000000000000005f000f426974546f7272656e74204b5250430053df9212ac10200f0000000000000000000000940000000000000159000f426974546f7272656e74204b525043005f43f0b2ac10200f0000000000000000000001220000000000000252000f426974546f7272656e74204b52504300567ce6fbac10200100000000000000000000005a000000000000005a00034e545000ac102080ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b5250430055550ef7ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b52504300ba9322a2ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043004579e7114b01bf5300000000000000000000006e0000000000000064000f426974546f7272656e74204b52504300ac10200fac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043005cf46adf") + + p := ipfix.NewSession() + _, err := p.ParseBuffer(p0) + if err != nil { + b.Fatal("ParseReader failed", err) + } + msg, err := p.ParseBuffer(p1) + + b.ResetTimer() + b.ReportAllocs() + b.SetBytes(1) + + for i := 0; i < b.N; i++ { + _, err = p.Marshal(msg) + if err != nil { + b.Error("Marshal failed", err) + } + } +} + func TestParsingTemplateAndDataRecords(t *testing.T) { packet, _ := hex.DecodeString("000a00405685b3700000000000bc614e000200140100000300080004000c0004000200040100001cc0a800c9c0a80001000000ebc0a800cac0a800010000002a") p := ipfix.NewSession() @@ -373,3 +438,24 @@ func TestParsingBufferAll(t *testing.T) { } } } + +func TestLookupTemplateRecords(t *testing.T) { + p0, _ := hex.DecodeString("000a008c51ec4264000000000b20bdbe0002007c283b0008001c0010800c000400003c258003000800003c258004000800003c258012ffff00003c258001ffff00003c25801cffff00003c25001b0010c2ac0008000c0004800c000400003c258003000800003c258004000800003c258012ffff00003c258001ffff00003c25801cffff00003c2500080004") + p1, _ := hex.DecodeString("000a05b051ec4270000000000b20bdbec2ac05a0ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043000116fcb8ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b525043005e489f46ac10200300000026000000000000019f0000000000000160000e4265696e6720616e616c797a656400c27ef905ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043007aa7519c0808080800000000000000000000008d00000000000000550003444e5300ac102082ac10200f0000000000000000000000940000000000000147000f426974546f7272656e74204b52504300b228265c1859c1570000000000000000000000000000000000000064000f426974546f7272656e74204b52504300ac10200fac10200f0000000000000000000000920000000000000145000f426974546f7272656e74204b525043007b75a68ad92bb37f00000000000000000000006e0000000000000064000f426974546f7272656e74204b52504300ac10200fac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043004f972c247449d8f200000000000000000000006e0000000000000064000f426974546f7272656e74204b52504300ac10200fac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b5250430048b682a4ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b52504300595cc40dac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b5250430057451cc1ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b525043005465e5a8ac1020ff00000000000000000000000000000000000000af001a44726f70626f78204c414e2073796e6320646973636f766572790764726f70626f78ac102013ac10200f00000000000000000000008f000000000000014b000f426974546f7272656e74204b5250430001ab3c06ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b52504300befcacc8ffffffff00000000000000000000000000000000000000af001a44726f70626f78204c414e2073796e6320646973636f766572790764726f70626f78ac102013ac10200300000025000000000000019e0000000000000167000e4265696e6720616e616c797a656400c27ef905ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043006ca28bcdac10200f000000000000000000000091000000000000011c000f426974546f7272656e74204b52504300b13531caac10200f000000000000000000000068000000000000005f000f426974546f7272656e74204b5250430053df9212ac10200f0000000000000000000000940000000000000159000f426974546f7272656e74204b525043005f43f0b2ac10200f0000000000000000000001220000000000000252000f426974546f7272656e74204b52504300567ce6fbac10200100000000000000000000005a000000000000005a00034e545000ac102080ac10200f00000000000000000000008c000000000000013a000f426974546f7272656e74204b5250430055550ef7ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b52504300ba9322a2ac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043004579e7114b01bf5300000000000000000000006e0000000000000064000f426974546f7272656e74204b52504300ac10200fac10200f0000000000000000000000910000000000000136000f426974546f7272656e74204b525043005cf46adf") + + p := ipfix.NewSession() + _, err := p.ParseBuffer(p0) + if err != nil { + t.Fatal("ParseReader failed", err) + } + msg, err := p.ParseBuffer(p1) + + trs, err := p.LookupTemplateRecords(msg) + if err != nil { + t.Fatal("LookupTemplateRecords failed", err) + } + + if len(trs) != 1 { + t.Fatalf("LookupTemplateRecords returned %d records instead of expected 1", len(trs)) + } +}