DDI/agent/role/dhcp/control.go

// SPDX-License-Identifier: Apache-2.0

package dhcp

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"time"
)

// Kea control channel result codes (Kea ARM 2.6 and 3.x, "Management API",
// "Data Syntax").
const (
	ResultSuccess     = 0
	ResultError       = 1
	ResultUnsupported = 2
	ResultEmpty       = 3
	ResultConflict    = 4
)

// maxResponse bounds a control channel response (lease pages are the
// largest ones).
const maxResponse = 256 << 20

// Response is the answer to a control command.
type Response struct {
	Result    int             `json:"result"`
	Text      string          `json:"text"`
	Arguments json.RawMessage `json:"arguments"`
}

// Err returns nil for result 0 and an error otherwise.
func (r *Response) Err(command string) error {
	if r.Result == ResultSuccess {
		return nil
	}
	return &CommandError{Command: command, Result: r.Result, Text: r.Text}
}

// CommandError is a non-zero result of a control command.
type CommandError struct {
	Command string
	Result  int
	Text    string
}

func (e *CommandError) Error() string {
	return fmt.Sprintf("kea %s: result %d: %s", e.Command, e.Result, e.Text)
}

// Command sends one command over the UNIX control socket of kea-dhcp4 and
// returns the response. It is used by the role and by tests.
//
// Protocol (Kea ARM 2.6 and 3.x, "Management API"): the client writes a JSON
// object {"command": ..., "arguments": ...}; the server answers with one
// JSON object {"result", "text", "arguments"} and closes the connection.
// Answers relayed by the (deprecated) Control Agent are JSON lists of such
// objects, so a list is accepted as well. The response is read as one JSON
// value, which does not depend on the server closing the connection. Kea
// itself aborts a command that is not received within 10 seconds.
func (r *Role) Command(ctx context.Context, command string, args any) (*Response, error) {
	ctx, cancel := context.WithTimeout(ctx, r.cfg.CommandTimeout)
	defer cancel()

	req := struct {
		Command   string `json:"command"`
		Arguments any    `json:"arguments,omitempty"`
	}{command, args}
	body, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("kea %s: %w", command, err)
	}

	var d net.Dialer
	conn, err := d.DialContext(ctx, "unix", r.cfg.ControlSocket)
	if err != nil {
		return nil, fmt.Errorf("kea %s: connect to control socket: %w", command, err)
	}
	defer func() { _ = conn.Close() }()
	if dl, ok := ctx.Deadline(); ok {
		_ = conn.SetDeadline(dl)
	}
	// Close the connection when ctx ends, so a cancellation interrupts reads.
	stop := context.AfterFunc(ctx, func() { _ = conn.SetDeadline(time.Unix(1, 0)) })
	defer stop()

	if _, err := conn.Write(body); err != nil {
		return nil, fmt.Errorf("kea %s: send: %w", command, err)
	}
	var raw json.RawMessage
	dec := json.NewDecoder(io.LimitReader(conn, maxResponse))
	if err := dec.Decode(&raw); err != nil {
		if ctx.Err() != nil {
			err = errors.Join(ctx.Err(), err)
		}
		return nil, fmt.Errorf("kea %s: read response: %w", command, err)
	}
	return decodeResponse(command, raw)
}

func decodeResponse(command string, raw json.RawMessage) (*Response, error) {
	raw = bytes.TrimSpace(raw)
	if len(raw) > 0 && raw[0] == '[' {
		var list []json.RawMessage
		if err := json.Unmarshal(raw, &list); err != nil {
			return nil, fmt.Errorf("kea %s: invalid response: %w", command, err)
		}
		if len(list) == 0 {
			return nil, fmt.Errorf("kea %s: empty response list", command)
		}
		raw = list[0]
	}
	var resp Response
	// Decode result separately so that a missing result is an error rather
	// than "success".
	var probe struct {
		Result *int `json:"result"`
	}
	if err := json.Unmarshal(raw, &probe); err != nil {
		return nil, fmt.Errorf("kea %s: invalid response: %w", command, err)
	}
	if probe.Result == nil {
		return nil, fmt.Errorf("kea %s: response without result", command)
	}
	if err := json.Unmarshal(raw, &resp); err != nil {
		return nil, fmt.Errorf("kea %s: invalid response: %w", command, err)
	}
	return &resp, nil
}