func (tr TemplateRecord) register(s session.Session) {
if s == nil {
return
}
if debug {
debugLog.Println("register template:", tr)
}
s.Lock()
defer s.Unlock()
s.AddTemplate(tr)
}
func (p *Packet) UnmarshalFlowSets(r io.Reader, s session.Session, t *Translate) error {
if debug {
debugLog.Printf("decoding %d flow sets\n", p.Header.Count)
}
for i := uint16(0); i < p.Header.Count; i++ {
// Read the next set header
header := FlowSetHeader{}
if err := header.Unmarshal(r); err != nil {
if debug {
debugLog.Println("failed to read flow set header:", err)
}
return err
}
switch header.ID {
case 0: // Template FlowSet
tfs := TemplateFlowSet{}
tfs.Header = header
readSize := int(tfs.Header.Length) - tfs.Header.Len()
if readSize < 4 {
if debug {
debugLog.Printf("short read size of %d\n", readSize)
}
return io.ErrShortBuffer
}
data := make([]byte, readSize)
if _, err := r.Read(data); err != nil {
if debug {
debugLog.Printf("failed to read %d bytes: %v\n", readSize, err)
}
return err
}
if err := tfs.UnmarshalRecords(bytes.NewBuffer(data)); err != nil {
return err
}
if debug {
debugLog.Printf("unmarshaled %d records: %v\n", len(tfs.Records), tfs)
}
for _, tr := range tfs.Records {
tr.register(s)
}
p.TemplateFlowSets = append(p.TemplateFlowSets, tfs)
return nil
case 1: // Options Template FlowSet
ofs := OptionsTemplateFlowSet{}
ofs.Header = header
readSize := int(ofs.Header.Length) - ofs.Header.Len()
if readSize < 4 {
return io.ErrShortBuffer
}
data := make([]byte, readSize)
if _, err := r.Read(data); err != nil {
return err
}
p.OptionsTemplateFlowSets = append(p.OptionsTemplateFlowSets, ofs)
default:
dfs := DataFlowSet{}
dfs.Header = header
if dfs.Header.Length < 4 {
return io.ErrShortBuffer
}
data := make([]byte, int(dfs.Header.Length)-dfs.Header.Len())
if _, err := r.Read(data); err != nil {
return err
}
var (
tm session.Template
tr TemplateRecord
ok bool
)
// If we don't have a session, or no template to resolve the Data
// Set contained Data Records, we'll store the raw bytes in stead.
if s == nil {
if debug {
debugLog.Printf("no session, storing %d raw bytes in data set\n", len(data))
}
dfs.Bytes = data
continue
}
if tm, ok = s.GetTemplate(header.ID); !ok {
if debug {
debugLog.Printf("no template for id=%d, storing %d raw bytes in data set\n", header.ID, len(data))
}
dfs.Bytes = data
continue
}
if tr, ok = tm.(TemplateRecord); !ok {
if debug {
debugLog.Printf("no template record, got %T, storing %d raw bytes in data set\n", tm, len(data))
}
dfs.Bytes = data
continue
}
if err := dfs.Unmarshal(bytes.NewBuffer(data), tr, t); err != nil {
return err
}
p.DataFlowSets = append(p.DataFlowSets, dfs)
return nil
}
}
return nil
}
// UnmarshalSets will, based on the Message length, unmarshal all sets in the
// message.
func (m *Message) UnmarshalSets(r io.Reader, s session.Session, t *Translate) error {
// Read the rest of the message, containing the sets.
data := make([]byte, int(m.Header.Length)-m.Header.Len())
if _, err := r.Read(data); err != nil {
return err
}
buffer := bytes.NewBuffer(data)
for buffer.Len() > 0 {
// Read the next set header
header := SetHeader{}
if err := header.Unmarshal(buffer); err != nil {
return err
}
if debug {
debugLog.Println("set header:", header)
}
if int(header.Length) < header.Len() {
return io.ErrUnexpectedEOF
}
data := make([]byte, int(header.Length)-header.Len())
if _, err := buffer.Read(data); err != nil {
return err
}
if debug {
debugLog.Printf("received set of %d bytes\n", len(data))
}
switch {
case header.ID < 2: // Legacy header ID, should not happen.
if debug {
debugLog.Println("received legacy set id", header.ID)
}
return errProtocol("received invalid set id")
case header.ID == 2: // Template set
ts := TemplateSet{}
ts.Header = header
if err := ts.UnmarshalRecords(bytes.NewBuffer(data)); err != nil {
return err
}
m.TemplateSets = append(m.TemplateSets, ts)
for _, tr := range ts.Records {
tr.register(s)
}
case header.ID == 3: // Options Template set
ots := OptionsTemplateSet{}
ots.Header = header
if err := ots.UnmarshalRecords(bytes.NewBuffer(data)); err != nil {
return err
}
m.OptionsTemplateSets = append(m.OptionsTemplateSets, ots)
case header.ID >= 4 && header.ID <= 255:
if debug {
debugLog.Println("received reserved set id", header.ID)
}
// Silently dropped
default:
ds := DataSet{}
ds.Header = header
var (
tm session.Template
tr TemplateRecord
ok bool
)
// If we don't have a session, or no template to resolve the Data
// Set contained Data Records, we'll store the raw bytes in stead.
if s == nil {
if debug {
debugLog.Printf("no session, storing %d raw bytes in data set\n", len(data))
}
ds.Bytes = data
continue
}
if tm, ok = s.GetTemplate(header.ID); !ok {
if debug {
debugLog.Printf("no template for id=%d, storing %d raw bytes in data set\n", header.ID, len(data))
}
ds.Bytes = data
continue
}
if tr, ok = tm.(TemplateRecord); !ok {
if debug {
debugLog.Printf("no template record, got %T, storing %d raw bytes in data set\n", tm, len(data))
}
ds.Bytes = data
continue
}
if err := ds.Unmarshal(bytes.NewBuffer(data), tr, t); err != nil {
return err
}
m.DataSets = append(m.DataSets, ds)
}
}
return nil
}