netsampler-goflow2/producer/proto/reflect.go

package protoproducer

import (
	"fmt"
	"reflect"

	"google.golang.org/protobuf/encoding/protowire"

	"github.com/netsampler/goflow2/v2/decoders/netflow"
)

type EndianType string
type ProtoType string

var (
	BigEndian    EndianType = "big"
	LittleEndian EndianType = "little"

	ProtoString ProtoType = "string"
	ProtoVarint ProtoType = "varint"

	ProtoTypeMap = map[string]ProtoType{
		string(ProtoString): ProtoString,
		string(ProtoVarint): ProtoVarint,
		"bytes":             ProtoString,
	}
)

func GetBytes(d []byte, offset int, length int) []byte {
	if length == 0 || offset < 0 {
		return nil
	}
	leftBytes := offset / 8
	rightBytes := (offset + length) / 8
	if (offset+length)%8 != 0 {
		rightBytes += 1
	}
	if leftBytes >= len(d) {
		return nil
	}
	if rightBytes > len(d) {
		rightBytes = len(d)
	}
	chunk := make([]byte, rightBytes-leftBytes)

	offsetMod8 := (offset % 8)
	shiftAnd := byte(0xff >> (8 - offsetMod8))

	var shifted byte
	for i := range chunk {
		j := len(chunk) - 1 - i
		cur := d[j+leftBytes]
		chunk[j] = (cur << offsetMod8) | shifted
		shifted = shiftAnd & cur
	}
	last := len(chunk) - 1
	shiftAndLast := byte(0xff << ((8 - ((offset + length) % 8)) % 8))
	chunk[last] = chunk[last] & shiftAndLast
	return chunk
}

func IsUInt(k reflect.Kind) bool {
	return k == reflect.Uint8 || k == reflect.Uint16 || k == reflect.Uint32 || k == reflect.Uint64
}

func IsInt(k reflect.Kind) bool {
	return k == reflect.Int8 || k == reflect.Int16 || k == reflect.Int32 || k == reflect.Int64
}

// Structure to help the MapCustom functions
// populate the protobuf data
type MapConfigBase struct {
	// Used if the field inside the protobuf exists
	// also serves as the field when rendering with text
	Destination string
	Endianness  EndianType

	// The following fields are used for mapping
	// when the destination field does not exist
	// inside the protobuf
	ProtoIndex int32
	ProtoType  ProtoType
	ProtoArray bool
}

func MapCustomNetFlow(flowMessage *ProtoProducerMessage, df netflow.DataField, mapper *NetFlowMapper) error {
	if mapper == nil {
		return nil
	}
	mapped, ok := mapper.Map(df)
	if ok {
		v := df.Value.([]byte)
		if err := MapCustom(flowMessage, v, mapped.MapConfigBase); err != nil {
			return err
		}
	}
	return nil
}

func MapCustom(flowMessage *ProtoProducerMessage, v []byte, cfg MapConfigBase) error {
	vfm := reflect.ValueOf(flowMessage)
	vfm = reflect.Indirect(vfm)

	fieldValue := vfm.FieldByName(cfg.Destination)

	if fieldValue.IsValid() {
		typeDest := fieldValue.Type()
		fieldValueAddr := fieldValue.Addr()

		if typeDest.Kind() == reflect.Slice {

			if typeDest.Elem().Kind() == reflect.Uint8 {
				fieldValue.SetBytes(v)
			} else {
				item := reflect.New(typeDest.Elem())

				if IsUInt(typeDest.Elem().Kind()) {
					if cfg.Endianness == LittleEndian {
						if err := DecodeUNumberLE(v, item.Interface()); err != nil {
							return err
						}
					} else {
						if err := DecodeUNumber(v, item.Interface()); err != nil {
							return err
						}
					}
				} else if IsInt(typeDest.Elem().Kind()) {
					if cfg.Endianness == LittleEndian {
						if err := DecodeNumberLE(v, item.Interface()); err != nil {
							return err
						}
					} else {
						if err := DecodeNumber(v, item.Interface()); err != nil {
							return err
						}
					}
				}

				itemi := reflect.Indirect(item)
				tmpFieldValue := reflect.Append(fieldValue, itemi)
				fieldValue.Set(tmpFieldValue)
			}

		} else if fieldValueAddr.IsValid() {
			if IsUInt(typeDest.Kind()) {
				if cfg.Endianness == LittleEndian {
					if err := DecodeUNumberLE(v, fieldValueAddr.Interface()); err != nil {
						return err
					}
				} else {
					if err := DecodeUNumber(v, fieldValueAddr.Interface()); err != nil {
						return err
					}
				}
			} else if IsInt(typeDest.Kind()) {
				if cfg.Endianness == LittleEndian {
					if err := DecodeNumberLE(v, fieldValueAddr.Interface()); err != nil {
						return err
					}
				} else {
					if err := DecodeNumber(v, fieldValueAddr.Interface()); err != nil {
						return err
					}
				}
			}

		}
	} else if cfg.ProtoIndex > 0 {

		fmr := flowMessage.ProtoReflect()
		unk := fmr.GetUnknown()

		if !cfg.ProtoArray {
			var offset int
			for offset < len(unk) {
				num, _, length := protowire.ConsumeField(unk[offset:])
				offset += length
				if int32(num) == cfg.ProtoIndex {
					// only one allowed
					break
				}
			}
		}

		var dstVar uint64
		if cfg.ProtoType == ProtoVarint {
			if cfg.Endianness == LittleEndian {
				if err := DecodeUNumberLE(v, &dstVar); err != nil {
					return err
				}
			} else {
				if err := DecodeUNumber(v, &dstVar); err != nil {
					return err
				}
			}
			// support signed int?
			unk = protowire.AppendTag(unk, protowire.Number(cfg.ProtoIndex), protowire.VarintType)
			unk = protowire.AppendVarint(unk, dstVar)
		} else if cfg.ProtoType == ProtoString {
			unk = protowire.AppendTag(unk, protowire.Number(cfg.ProtoIndex), protowire.BytesType)
			unk = protowire.AppendString(unk, string(v))
		} else {
			return fmt.Errorf("could not insert into protobuf unknown")
		}
		fmr.SetUnknown(unk)
	}
	return nil
}