blob: 114d250018b433aa65cb40d64da049a2cd773ef7 [file] [log] [blame]
package btf
import (
"bufio"
"bytes"
"errors"
"fmt"
"io"
"strings"
"golang.org/x/exp/maps"
"golang.org/x/exp/slices"
)
type stringTable struct {
base *stringTable
offsets []uint32
prevIdx int
strings []string
}
// sizedReader is implemented by bytes.Reader, io.SectionReader, strings.Reader, etc.
type sizedReader interface {
io.Reader
Size() int64
}
func readStringTable(r sizedReader, base *stringTable) (*stringTable, error) {
// When parsing split BTF's string table, the first entry offset is derived
// from the last entry offset of the base BTF.
firstStringOffset := uint32(0)
if base != nil {
idx := len(base.offsets) - 1
firstStringOffset = base.offsets[idx] + uint32(len(base.strings[idx])) + 1
}
// Derived from vmlinux BTF.
const averageStringLength = 16
n := int(r.Size() / averageStringLength)
offsets := make([]uint32, 0, n)
strings := make([]string, 0, n)
offset := firstStringOffset
scanner := bufio.NewScanner(r)
scanner.Split(splitNull)
for scanner.Scan() {
str := scanner.Text()
offsets = append(offsets, offset)
strings = append(strings, str)
offset += uint32(len(str)) + 1
}
if err := scanner.Err(); err != nil {
return nil, err
}
if len(strings) == 0 {
return nil, errors.New("string table is empty")
}
if firstStringOffset == 0 && strings[0] != "" {
return nil, errors.New("first item in string table is non-empty")
}
return &stringTable{base, offsets, 0, strings}, nil
}
func splitNull(data []byte, atEOF bool) (advance int, token []byte, err error) {
i := bytes.IndexByte(data, 0)
if i == -1 {
if atEOF && len(data) > 0 {
return 0, nil, errors.New("string table isn't null terminated")
}
return 0, nil, nil
}
return i + 1, data[:i], nil
}
func (st *stringTable) Lookup(offset uint32) (string, error) {
if st.base != nil && offset <= st.base.offsets[len(st.base.offsets)-1] {
return st.base.lookup(offset)
}
return st.lookup(offset)
}
func (st *stringTable) lookup(offset uint32) (string, error) {
// Fast path: zero offset is the empty string, looked up frequently.
if offset == 0 && st.base == nil {
return "", nil
}
// Accesses tend to be globally increasing, so check if the next string is
// the one we want. This skips the binary search in about 50% of cases.
if st.prevIdx+1 < len(st.offsets) && st.offsets[st.prevIdx+1] == offset {
st.prevIdx++
return st.strings[st.prevIdx], nil
}
i, found := slices.BinarySearch(st.offsets, offset)
if !found {
return "", fmt.Errorf("offset %d isn't start of a string", offset)
}
// Set the new increment index, but only if its greater than the current.
if i > st.prevIdx+1 {
st.prevIdx = i
}
return st.strings[i], nil
}
// Num returns the number of strings in the table.
func (st *stringTable) Num() int {
return len(st.strings)
}
// stringTableBuilder builds BTF string tables.
type stringTableBuilder struct {
length uint32
strings map[string]uint32
}
// newStringTableBuilder creates a builder with the given capacity.
//
// capacity may be zero.
func newStringTableBuilder(capacity int) *stringTableBuilder {
var stb stringTableBuilder
if capacity == 0 {
// Use the runtime's small default size.
stb.strings = make(map[string]uint32)
} else {
stb.strings = make(map[string]uint32, capacity)
}
// Ensure that the empty string is at index 0.
stb.append("")
return &stb
}
// Add a string to the table.
//
// Adding the same string multiple times will only store it once.
func (stb *stringTableBuilder) Add(str string) (uint32, error) {
if strings.IndexByte(str, 0) != -1 {
return 0, fmt.Errorf("string contains null: %q", str)
}
offset, ok := stb.strings[str]
if ok {
return offset, nil
}
return stb.append(str), nil
}
func (stb *stringTableBuilder) append(str string) uint32 {
offset := stb.length
stb.length += uint32(len(str)) + 1
stb.strings[str] = offset
return offset
}
// Lookup finds the offset of a string in the table.
//
// Returns an error if str hasn't been added yet.
func (stb *stringTableBuilder) Lookup(str string) (uint32, error) {
offset, ok := stb.strings[str]
if !ok {
return 0, fmt.Errorf("string %q is not in table", str)
}
return offset, nil
}
// Length returns the length in bytes.
func (stb *stringTableBuilder) Length() int {
return int(stb.length)
}
// AppendEncoded appends the string table to the end of the provided buffer.
func (stb *stringTableBuilder) AppendEncoded(buf []byte) []byte {
n := len(buf)
buf = append(buf, make([]byte, stb.Length())...)
strings := buf[n:]
for str, offset := range stb.strings {
copy(strings[offset:], str)
}
return buf
}
// Copy the string table builder.
func (stb *stringTableBuilder) Copy() *stringTableBuilder {
return &stringTableBuilder{
stb.length,
maps.Clone(stb.strings),
}
}