client_linux.go

//go:build linux
// +build linux

package taskstats

import (
	"fmt"
	"os"
	"unsafe"

	"github.com/mdlayher/genetlink"
	"github.com/mdlayher/netlink"
	"github.com/mdlayher/netlink/nlenc"
	"golang.org/x/sys/unix"
)

// Fixed structure sizes.
const sizeofCGroupStats = int(unsafe.Sizeof(unix.CGroupStats{}))

var _ osClient = &client{}

// A client is a Linux-specific taskstats client.
type client struct {
	c      *genetlink.Conn
	family genetlink.Family
}

// newClient opens a connection to the taskstats family using
// generic netlink.
func newClient() (*client, error) {
	c, err := genetlink.Dial(nil)
	if err != nil {
		return nil, err
	}

	// Best effort.
	_ = c.SetOption(netlink.ExtendedAcknowledge, true)

	return initClient(c)
}

// initClient is the internal client constructor used in some tests.
func initClient(c *genetlink.Conn) (*client, error) {
	f, err := c.GetFamily(unix.TASKSTATS_GENL_NAME)
	if err != nil {
		_ = c.Close()
		return nil, err
	}

	return &client{
		c:      c,
		family: f,
	}, nil
}

// Close implements osClient.
func (c *client) Close() error {
	return c.c.Close()
}

// PID implements osClient.
func (c *client) PID(pid int) (*Stats, error) {
	return c.getStats(pid, unix.TASKSTATS_CMD_ATTR_PID, unix.TASKSTATS_TYPE_AGGR_PID)
}

// TGID implements osClient.
func (c *client) TGID(tgid int) (*Stats, error) {
	return c.getStats(tgid, unix.TASKSTATS_CMD_ATTR_TGID, unix.TASKSTATS_TYPE_AGGR_TGID)
}

func (c *client) getStats(id int, cmdAttr, typeAggr uint16) (*Stats, error) {
	// Query taskstats for information using a specific ID.
	attrs := []netlink.Attribute{{
		Type: cmdAttr,
		Data: nlenc.Uint32Bytes(uint32(id)),
	}}

	msg, err := c.execute(unix.TASKSTATS_CMD_GET, attrs)
	if err != nil {
		return nil, err
	}

	return parseMessage(*msg, typeAggr)
}

// CGroupStats implements osClient.
func (c *client) CGroupStats(path string) (*CGroupStats, error) {
	// Open cgroup path so its file descriptor can be passed to taskstats.
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	// Query taskstats for cgroup information using the file descriptor.
	attrs := []netlink.Attribute{{
		Type: unix.CGROUPSTATS_CMD_ATTR_FD,
		Data: nlenc.Uint32Bytes(uint32(f.Fd())),
	}}

	msg, err := c.execute(unix.CGROUPSTATS_CMD_GET, attrs)
	if err != nil {
		return nil, err
	}

	return parseCGroupMessage(*msg)
}

// execute executes a single generic netlink command and returns its response.
func (c *client) execute(cmd uint8, attrs []netlink.Attribute) (*genetlink.Message, error) {
	b, err := netlink.MarshalAttributes(attrs)
	if err != nil {
		return nil, err
	}

	msg := genetlink.Message{
		Header: genetlink.Header{
			Command: cmd,
			Version: unix.TASKSTATS_VERSION,
		},
		Data: b,
	}

	msgs, err := c.c.Execute(msg, c.family.ID, netlink.Request)
	if err != nil {
		// We don't want to expose netlink errors directly to callers, so unpack
		// the error for use with os.IsPermission and similar.
		oerr, ok := err.(*netlink.OpError)
		if !ok {
			// Expect all errors to conform to netlink.OpError.
			return nil, fmt.Errorf("taskstats: netlink operation returned non-netlink error (please file a bug: https://github.com/mdlayher/taskstats): %v", err)
		}

		return nil, oerr.Err
	}

	if l := len(msgs); l != 1 {
		return nil, fmt.Errorf("taskstats: unexpected number of response messages: %d", l)
	}

	return &msgs[0], nil
}

// parseCGroupMessage attempts to parse a CGroupStats structure from a generic netlink message.
func parseCGroupMessage(m genetlink.Message) (*CGroupStats, error) {
	attrs, err := netlink.UnmarshalAttributes(m.Data)
	if err != nil {
		return nil, err
	}

	for _, a := range attrs {
		// Only parse cgroupstats structure.
		if a.Type != unix.CGROUPSTATS_TYPE_CGROUP_STATS {
			continue
		}

		// Verify that the byte slice containing a unix.CGroupStats is the
		// size expected by this package, so we don't blindly cast the
		// byte slice into a structure of the wrong size.
		if want, got := sizeofCGroupStats, len(a.Data); want != got {
			return nil, fmt.Errorf("unexpected cgroupstats structure size, want %d, got %d", want, got)
		}

		cs := *(*unix.CGroupStats)(unsafe.Pointer(&a.Data[0]))
		return parseCGroupStats(cs)
	}

	// No taskstats response found.
	return nil, os.ErrNotExist
}

// parseMessage attempts to parse a Stats structure from a generic netlink message.
func parseMessage(m genetlink.Message, typeAggr uint16) (*Stats, error) {
	attrs, err := netlink.UnmarshalAttributes(m.Data)
	if err != nil {
		return nil, err
	}

	for _, a := range attrs {
		// Only parse ID+stats structure.
		if a.Type != typeAggr {
			continue
		}

		nattrs, err := netlink.UnmarshalAttributes(a.Data)
		if err != nil {
			return nil, err
		}

		for _, na := range nattrs {
			// Only parse Stats element since caller would already have ID.
			if na.Type != unix.TASKSTATS_TYPE_STATS {
				continue
			}

			// Assume the kernel returns a structure compatible with the
			// taskstats struct and cast directly.
			return parseStats(*(*unix.Taskstats)(unsafe.Pointer(&na.Data[0])))
		}
	}

	// No taskstats response found.
	return nil, os.ErrNotExist
}