github.com/castai/kvisor@v1.7.1-0.20240516114728-b3572a2607b5/tools/codegen/generator.go

package main

import (
	"errors"
	"fmt"
	"strings"
	"unicode"

	"github.com/iancoleman/strcase"
)

func generateTypes(events []eventDefinition) (string, error) {
	sink := &strings.Builder{}

	_, err := sink.WriteString(generateTypesFileHeader())
	if err != nil {
		return "", err
	}

	_, err = sink.WriteString("\n")
	if err != nil {
		return "", err
	}

	for i, definition := range events {
		err = generateStruct(sink, definition)
		if err != nil {
			return "", err
		}

		if i < len(events)-1 {
			_, err = sink.WriteString("\n")
			if err != nil {
				return "", err
			}
		}
	}

	return sink.String(), nil
}

func generateParsers(targetPackage string, events []eventDefinition) (string, error) {
	sink := &strings.Builder{}

	_, err := sink.WriteString(generateParsersFileHeader(targetPackage))
	if err != nil {
		return "", err
	}

	_, err = sink.WriteString("\n")
	if err != nil {
		return "", err
	}

	for _, definition := range events {
		err = generatePerEventParserFunction(sink, definition)
		if err != nil {
			return "", err
		}
		_, err = sink.WriteString("\n")
		if err != nil {
			return "", err
		}
	}

	err = generateParserFunction(sink, events)
	if err != nil {
		return "", err
	}

	return sink.String(), nil
}

func generateTypesFileHeader() string {
	return fmt.Sprintf(`// Code generated by tools/codegen; DO NOT EDIT.

package types

type Args interface {
  args()
}

// internalArgs is a marker type to distinguish Args interface from basically any
type internalArgs struct{}

func (c internalArgs) args() {}
`)
}

func generateParsersFileHeader(targetPackage string) string {
	return fmt.Sprintf(`// Code generated by tools/codegen; DO NOT EDIT.

package %s

import (
  "errors"

  "github.com/castai/kvisor/pkg/ebpftracer/events"
  "github.com/castai/kvisor/pkg/ebpftracer/types"
)

var (
  ErrUnknownArgsType error = errors.New("unknown args type")
)

// eventMaxByteSliceBufferSize is used to determine the max slice size allowed for different
// event types. For example, most events have a max size of 4096, but for network related events
// there is no max size (this is represented as -1).
func eventMaxByteSliceBufferSize(id events.ID) int {
  // For non network event, we have a max byte slice size of 4096
  if id < events.NetPacketBase || id > events.MaxNetID {
    return 4096
  }

  // Network events do not have a max buffer size.
  return -1
}
`, targetPackage)
}

func generateStruct(sink *strings.Builder, definition eventDefinition) error {
	sink.WriteString(fmt.Sprintf(`type %s struct {
  internalArgs
`, generateArgName(definition)))

	if len(definition.params) > 0 {
		_, err := sink.WriteRune('\n')
		if err != nil {
			return err
		}
	}

	for _, p := range definition.params {
		paramLine, err := generateParamStructField(p)
		if err != nil {
			return err
		}

		_, err = sink.WriteString(paramLine)
		if err != nil {
			return err
		}

		_, err = sink.WriteString("\n")
		if err != nil {
			return err
		}
	}
	_, err := sink.WriteString("}\n")
	if err != nil {
		return err
	}

	return nil
}

func generateParamStructField(p param) (string, error) {
	goType, err := toGolangType(p.paramType)
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("  %s %s", generateParamName(p), goType), nil
}

func capitalize(str string) string {
	runes := []rune(str)
	runes[0] = unicode.ToUpper(runes[0])
	return string(runes)
}

func generateParamName(p param) string {
	// This should transform parameters called e.g. `src_old` to `SrcOld`
	return strcase.ToCamel(p.name)
}

func generateArgName(definition eventDefinition) string {
	return fmt.Sprintf("%sArgs", capitalize(definition.event))
}

func toGolangType(t ArgType) (string, error) {
	switch t {
	case noneT:
		return "", errors.New("cannot handle erorr type none!")
	case u8T:
		return "uint8", nil
	case u16T:
		return "uint16", nil
	case intT:
		return "int32", nil
	case uintT, devT, modeT:
		return "uint32", nil
	case longT:
		return "int64", nil
	case ulongT, offT, sizeT:
		return "uint64", nil
	case boolT:
		return "bool", nil
	case pointerT:
		return "uintptr", nil
	case sockAddrT:
		return "Sockaddr", nil
	case credT:
		return "SlimCred", nil
	case strT:
		return "string", nil
	case strArrT, argsArrT:
		return "[]string", nil
	case bytesT:
		return "[]byte", nil
	case intArr2T:
		return "[2]int32", nil
	case uint64ArrT:
		return "[]uint64", nil
	case timespecT:
		return "float64", nil
	case tupleT:
		return "AddrTuple", nil
	case protoDNST:
		return "*ProtoDNS", nil
	}

	return "", fmt.Errorf("unknown event type: %d", t)
}

func indent(str string, indentSize int) string {
	if len(str) == 0 {
		return str
	}

	indent := strings.Repeat(" ", indentSize)

	var result strings.Builder

	for _, line := range strings.Split(str, "\n") {
		if len(strings.TrimSpace(line)) > 0 {
			_, err := result.WriteString(indent)
			if err != nil {
				panic(err)
			}

			_, err = result.WriteString(line)
			if err != nil {
				panic(err)
			}
		}

		_, err := result.WriteRune('\n')
		if err != nil {
			panic(err)
		}
	}

	return result.String()
}

func generatePerEventParserFunction(sink *strings.Builder, definition eventDefinition) error {
	eventName := generateArgName(definition)

	_, err := sink.WriteString(fmt.Sprintf(`func Parse%s(decoder *Decoder) (types.%s, error) {
`,
		eventName, eventName))
	if err != nil {
		return err
	}

	if len(definition.params) == 0 {
		_, err = sink.WriteString(fmt.Sprintf(`  return types.%s{}, nil
}
`, eventName))
		if err != nil {
			return err
		}
		return nil
	}

	_, err = sink.WriteString(fmt.Sprintf(`  var result types.%s
  var err error

`, eventName))
	if err != nil {
		return err
	}

	_, err = sink.WriteString(generateParseNumArgsCode(definition))
	if err != nil {
		return err
	}

	_, err = sink.WriteString("\n")
	if err != nil {
		return err
	}

	_, err = sink.WriteString(`  for arg := 0; arg < int(numArgs); arg++ {
`)
	if err != nil {
		return err
	}

	_, err = sink.WriteString(indent(generateCurrentArgCode(definition), 2))
	if err != nil {
		return err
	}

	_, err = sink.WriteString(indent(`  switch currArg {`, 2))
	if err != nil {
		return err
	}

	for i, p := range definition.params {
		_, err = sink.WriteString(indent(fmt.Sprintf(`  case %d:`, i), 2))
		if err != nil {
			return err
		}

		line, err := getDecoderCode(definition, p)
		if err != nil {
			return err
		}

		_, err = sink.WriteString(indent(line, 4))
		if err != nil {
			return err
		}
	}

	_, err = sink.WriteString(`    }
  }
`)
	if err != nil {
		return err
	}

	_, err = sink.WriteString(`  return result, nil
}
`)
	if err != nil {
		return err
	}

	return nil
}
func generateParseNumArgsCode(definition eventDefinition) string {
	return fmt.Sprintf(`  var numArgs uint8
  err = decoder.DecodeUint8(&numArgs)
%s
`, generateDecoderErrorCheck(definition))
}

func generateCurrentArgCode(definition eventDefinition) string {
	return fmt.Sprintf(`  var currArg uint8
  err = decoder.DecodeUint8(&currArg)
%s
`, generateDecoderErrorCheck(definition))
}

func getDecoderCode(definition eventDefinition, p param) (string, error) {
	paramName := generateParamName(p)
	switch p.paramType {
	case noneT:
		return "", errors.New("cannot handle erorr type none!")
	case u8T:
		return fmt.Sprintf(`  err = decoder.DecodeUint8(&result.%s)
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case u16T:
		return fmt.Sprintf(`  err = decoder.DecodeUint16(&result.%s)
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case intT:
		return fmt.Sprintf(`  err = decoder.DecodeInt32(&result.%s)
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case uintT, devT, modeT:
		return fmt.Sprintf(`  err = decoder.DecodeUint32(&result.%s)
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case longT:
		return fmt.Sprintf(`  err = decoder.DecodeInt64(&result.%s)
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case ulongT, offT, sizeT:
		return fmt.Sprintf(`  err = decoder.DecodeUint64(&result.%s)
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case boolT:
		return fmt.Sprintf(`  err = decoder.DecodeBool(&result.%s)
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case pointerT:
		return fmt.Sprintf(`  var data%s uint64
  err = decoder.DecodeUint64(&data%s)
%s
  result.%s = uintptr(data%s)`, paramName, paramName, generateDecoderErrorCheck(definition), paramName, paramName), nil
	case sockAddrT:
		return fmt.Sprintf(`  result.%s, err = decoder.ReadSockaddrFromBuff()
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case credT:
		return fmt.Sprintf(`  err = decoder.DecodeSlimCred(&result.%s)
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case strT:
		return fmt.Sprintf(`  result.%s, err = decoder.ReadStringFromBuff()
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case strArrT:
		return fmt.Sprintf(`  result.%s, err = decoder.ReadStringArrayFromBuff()
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case argsArrT:
		return fmt.Sprintf(`  result.%s, err = decoder.ReadArgsArrayFromBuff()
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case bytesT:
		return fmt.Sprintf(`  result.%s, err = decoder.ReadMaxByteSliceFromBuff(eventMaxByteSliceBufferSize(events.%s))
%s`, paramName, definition.event, generateDecoderErrorCheck(definition)), nil
	case intArr2T:
		return fmt.Sprintf(`  err = decoder.DecodeIntArray(result.%s[:], 2)
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case uint64ArrT:
		return fmt.Sprintf(`  err = decoder.DecodeUint64Array(&result.%s)
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case timespecT:
		return fmt.Sprintf(`  result.%s, err = decoder.ReadTimespec()
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case tupleT:
		return fmt.Sprintf(`  result.%s, err = decoder.ReadAddrTuple()
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	case protoDNST:
		return fmt.Sprintf(`  result.%s, err = decoder.ReadProtoDNS()
%s`, paramName, generateDecoderErrorCheck(definition)), nil
	}

	return "", fmt.Errorf("unknown event type: %d", p.paramType)
}

func generateDecoderErrorCheck(definition eventDefinition) string {
	return fmt.Sprintf(`  if err != nil {
    return types.%s{}, err
  }`, generateArgName(definition))
}

func generateParserFunction(sink *strings.Builder, definitions []eventDefinition) error {
	_, err := sink.WriteString(`func ParseArgs(decoder *Decoder, event events.ID) (types.Args, error) {
  switch event {
`)
	if err != nil {
		return err
	}

	for _, definition := range definitions {
		sink.WriteString(fmt.Sprintf(`  case events.%s:
    return Parse%s(decoder)
`, definition.event, generateArgName(definition)))
	}

	_, err = sink.WriteString(`  }

  return nil, ErrUnknownArgsType
}
`)
	if err != nil {
		return err
	}
	return nil
}