LocalAI/core/services/studio/egressapi.go

package studio

import (
	"context"
	"errors"
	"slices"
	"time"

	"github.com/mudler/LocalAI/core/services/studio/egress"
	"github.com/mudler/LocalAI/core/services/studio/egress/audit"
	"github.com/mudler/LocalAI/core/services/studio/session"
	"github.com/mudler/LocalAI/pkg/addons/contract"
)

// The routes of the network of a chat's commands: the chat's own
// view and preset for its owner, and the administrator's view.

var (
	// ErrEgressPreset means the preset is not off, registries or all.
	ErrEgressPreset = errors.New("studio: the network preset must be off, registries or all")
	// ErrEgressAboveCeiling means the preset is wider than the administrator
	// allows.
	ErrEgressAboveCeiling = errors.New("studio: the network preset is above the limit an administrator set")
	// ErrEgressNoRule means the chat has no rule for that host and port.
	ErrEgressNoRule = errors.New("studio: the chat has no network rule for this host and port")
)

// The reasons the network settings do not apply to a chat.
const (
	EgressAppliesHost = "host" // commands run on this machine, with its network
	EgressAppliesNone = "none" // no add-on serves bash, or it does not report that it enforces the policy
)

// EgressView is the network of one chat, as its network panel shows it.
type EgressView struct {
	Preset, Ceiling, Effective egress.Preset
	// Applies is true when the provider that runs the commands reports that it
	// enforces the network policy. It is false when commands run on this
	// machine, no add-on serves bash, or the provider does not report
	// enforcement; AppliesReason then says which.
	Applies       bool
	AppliesReason string
	// Addon is the add-on that runs the chat's commands and carries the
	// network-policy host control, which draws the network row of the chat:
	// set only for a provider that carries it. Without it no add-on draws the controls, and
	// the policy still holds.
	Addon string
	// Network names the network of Addon in text and gives the titles it
	// declares for the presets. Set with Addon.
	Network    *AddonNetwork
	Registries []egress.Registry
	// AllowedHosts is the administrator's always-allowed list, canonical.
	AllowedHosts []string
	Rules        []EgressRuleView
	Cards        []egress.Card
	// CardAbove holds the IDs of the pending cards whose host the ceiling does
	// not admit: no answer can make them work.
	CardAbove map[string]bool
}

// EgressRuleView is a rule with whether the ceiling still admits it.
type EgressRuleView struct {
	egress.Rule
	AboveLimit bool
}

// registriesView is the registries list with the administrator's extra hosts
// as one more entry.
func registriesView(s egress.Settings) []egress.Registry {
	out := append([]egress.Registry(nil), egress.BuiltinRegistries().List...)
	if len(s.ExtraRegistries) > 0 {
		out = append(out, egress.Registry{ID: "extra", Title: "Added by an administrator", Hosts: append([]string(nil), s.ExtraRegistries...)})
	}
	return out
}

// Egress returns the network of a chat for its owner. The chat of another
// user is session.ErrNotFound.
func (s *Service) Egress(ctx context.Context, user, conv string) (EgressView, error) {
	ctx, closeScope := s.executionScope(ctx, principal(user), conv)
	defer closeScope()
	p := principal(user)
	rec, err := s.d.Store.GetChat(ctx, p, conv)
	if err != nil {
		return EgressView{}, err
	}
	set, err := s.egressSettings(ctx)
	if err != nil {
		return EgressView{}, err
	}
	pol, err := s.egressPolicyOf(ctx, rec, "")
	if err != nil {
		return EgressView{}, err
	}
	v := EgressView{
		Preset: pol.chat.Preset, Ceiling: egress.Min(set.Ceiling, egress.PresetAll),
		Effective:    egress.Min(pol.chat.Preset, set.Ceiling),
		Registries:   registriesView(set),
		AllowedHosts: set.AllowedHostList(),
		Rules:        []EgressRuleView{},
		Cards:        append([]egress.Card{}, pol.rec.Egress.Cards...),
	}
	v.CardAbove = map[string]bool{}
	for _, c := range v.Cards {
		if c.State == egress.CardPending || c.State == egress.CardRefused {
			req := egress.Request{Host: c.Host, Port: c.Port, Kind: c.Kind, TLS: c.Port == 443 && c.Kind == egress.KindConnect, Private: c.Private}
			if !egress.Admits(set, req, egress.BuiltinRegistries()) {
				v.CardAbove[c.ID] = true
			}
		}
	}
	for _, r := range pol.rec.Egress.Rules {
		above := r.Allow && !egress.Admits(set, egress.Request{Host: r.Host, Port: r.Port, TLS: r.Port == 443}, egress.BuiltinRegistries())
		v.Rules = append(v.Rules, EgressRuleView{Rule: r, AboveLimit: above})
	}
	if v.Preset == "" {
		v.Preset = egress.PresetOff
	}
	cmds, err := s.CommandRunner(ctx, p, conv)
	switch {
	case err != nil:
		v.AppliesReason = EgressAppliesNone
	case cmds.State == CommandsAvailable && (cmds.Provider == BuiltinProvider || cmds.Isolation == "host"):
		v.AppliesReason = EgressAppliesHost
	case cmds.State == CommandsAvailable:
		// Isolation alone does not establish network-policy enforcement: the
		// settings apply only when the provider reports that it enforces them
		// in a status that is available and isolated.
		v.Applies = cmds.ExecutionStatus != nil && cmds.ExecutionStatus.NetworkPolicyEnforced
		if !v.Applies {
			v.AppliesReason = EgressAppliesNone
		}
		if c, ok := s.d.AddonTools.(AddonControls); ok && slices.Contains(c.Controls(ctx, p, cmds.Provider), contract.ControlNetworkPolicy) {
			v.Addon = cmds.Provider
			n := s.networkOf(ctx, p, cmds)
			v.Network = &n
		}
	default:
		v.AppliesReason = EgressAppliesNone
	}
	return v, nil
}

// SetEgressPreset sets the preset of a chat. A preset above the ceiling is
// refused; a sub-agent follows its parent and has none of its own.
func (s *Service) SetEgressPreset(ctx context.Context, user, conv string, preset egress.Preset) error {
	p := principal(user)
	if !preset.Valid() {
		return ErrEgressPreset
	}
	rec, err := s.d.Store.GetChat(ctx, p, conv)
	if err != nil {
		return err
	}
	if rec.Kind == session.KindSubAgent {
		return ErrSubAgentReadOnly
	}
	set, err := s.egressSettings(ctx)
	if err != nil {
		return err
	}
	if preset.Rank() > egress.Min(set.Ceiling, egress.PresetAll).Rank() {
		return ErrEgressAboveCeiling
	}
	if err := s.d.Store.UpdateEgress(ctx, p, conv, func(st *egress.ChatState) error {
		st.Preset = preset
		return nil
	}); err != nil {
		return err
	}
	s.publishEgress(p, conv, "")
	return nil
}

// DeleteEgressRule removes the allow or deny rule of a host and port.
func (s *Service) DeleteEgressRule(ctx context.Context, user, conv, host string, port int) error {
	p := principal(user)
	host = canonicalHost(host)
	err := s.d.Store.UpdateEgress(ctx, p, conv, func(st *egress.ChatState) error {
		if !st.DeleteRule(host, port) {
			return ErrEgressNoRule
		}
		return nil
	})
	if err != nil {
		return err
	}
	s.publishEgress(p, conv, "")
	return nil
}

// publishEgress tells every open tab of the chat that its network changed:
// a card (id) or the rules and the preset (empty id).
func (s *Service) publishEgress(p, conv, card string) {
	ev := session.Event{Kind: session.KindEgressCard}
	if card != "" {
		ev.IDs = []string{card}
	}
	s.bus.Publish(key(p, conv), ev)
}

// networkOf is the network of the add-on that runs the commands in cmds: what
// the add-on declares, else the title of its commands tool.
func (s *Service) networkOf(ctx context.Context, user string, cmds Commands) AddonNetwork {
	n := AddonNetwork{Title: cmds.Title}
	if b, ok := s.d.AddonTools.(AddonNetworks); ok {
		if d, ok := b.Network(ctx, user, cmds.Provider); ok {
			n.Presets = d.Presets
			if d.Title != "" {
				n.Title = d.Title
			}
		}
	}
	return n
}

// EgressAdmin is the administrator's view of the network of commands.
type EgressAdmin struct {
	Settings   egress.Settings
	AuditDays  int
	Registries []egress.Registry
	// Runner is the add-on that runs commands for the administrator, and
	// Isolation how ("isolated", "host" or "" when unknown). Empty Runner:
	// none serves, and the settings apply to nothing yet.
	Runner, Isolation string
	// Network names the network of Runner in text, when one serves.
	Network *AddonNetwork
	Last24h audit.Counts
}

// AdminEgress returns the settings in force, the registries list with the
// extra hosts, whether an add-on serves and the counts of the last 24 hours.
func (s *Service) AdminEgress(ctx context.Context, user string) (EgressAdmin, error) {
	set, err := s.egressSettings(ctx)
	if err != nil {
		return EgressAdmin{}, err
	}
	set.Ceiling = egress.Min(set.Ceiling, egress.PresetAll)
	set.Default = set.DefaultPreset()
	if len(set.Ports) == 0 {
		set.Ports = []int{80, 443}
	}
	out := EgressAdmin{Settings: set, Registries: registriesView(set), AuditDays: DefaultEgressAuditDays}
	if s.d.EgressAuditDays != nil {
		if d := s.d.EgressAuditDays(); d >= 1 && d <= 365 {
			out.AuditDays = d
		}
	}
	out.Runner, err = s.servingRunner(ctx, principal(user))
	if err != nil {
		return EgressAdmin{}, err
	}
	if out.Runner != "" {
		if iso, ok := s.d.AddonTools.(CommandIsolation); ok {
			out.Isolation = iso.Isolation(ctx, principal(user), out.Runner)
		}
		n := s.networkOf(ctx, principal(user), Commands{Provider: out.Runner, Title: s.commandTitle(ctx, principal(user), out.Runner)})
		out.Network = &n
	}
	c, err := s.egressAudit().Count(ctx, s.now().Add(-24*time.Hour))
	if err != nil {
		return EgressAdmin{}, err
	}
	out.Last24h = c
	return out, nil
}

// EgressEvents reads the network audit, newest first.
func (s *Service) EgressEvents(ctx context.Context, q audit.Query) ([]audit.Event, error) {
	return s.egressAudit().Events(ctx, q)
}

// servingRunner is the add-on that runs the commands of the user's chats now:
// of the running add-ons that replace bash, the pinned one or the oldest
// install, or "" when none runs or the pinned one is stopped (the built-in
// serves). It asks the broker, so every frontend of a cluster answers the
// same, also right after a restart, and an add-on that was removed is gone.
func (s *Service) servingRunner(ctx context.Context, user string) (string, error) {
	pins, err := s.toolPins(ctx)
	if err != nil {
		return "", err
	}
	if s.d.AddonTools == nil {
		return "", nil
	}
	defs, err := s.d.AddonTools.Tools(ctx, user)
	if err != nil {
		return "", nil
	}
	var running []string
	installed := map[string]string{}
	for _, d := range defs {
		if d.Shadows == commandTool && !slices.Contains(running, d.AddonID) {
			running = append(running, d.AddonID)
			installed[d.AddonID] = d.InstalledAt
		}
	}
	if len(running) == 0 {
		return "", nil
	}
	pin := pins[commandTool]
	if s.shadowersFor(ctx, user).pinnedButIdle(commandTool, pin, running) {
		return "", nil
	}
	return ChooseProvider(running, installed, pin), nil
}