banger/internal/daemon/session_stream.go

package daemon

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"net"
	"path/filepath"
	"strings"

	"banger/internal/api"
	sess "banger/internal/daemon/session"
	"banger/internal/guest"
	"banger/internal/model"
	"banger/internal/system"
)

func (d *Daemon) GuestSessionLogs(ctx context.Context, params api.GuestSessionLogsParams) (api.GuestSessionLogsResult, error) {
	vm, err := d.FindVM(ctx, params.VMIDOrName)
	if err != nil {
		return api.GuestSessionLogsResult{}, err
	}
	session, err := d.findGuestSession(ctx, vm.ID, params.SessionIDOrName)
	if err != nil {
		return api.GuestSessionLogsResult{}, err
	}
	streamName := strings.TrimSpace(params.Stream)
	if streamName == "" {
		streamName = "stdout"
	}
	tailLines := params.TailLines
	if tailLines <= 0 {
		tailLines = sess.LogTailLineDefault
	}
	path := session.StdoutLogPath
	if streamName == "stderr" {
		path = session.StderrLogPath
	}
	content, err := d.readGuestSessionLog(ctx, vm, session, streamName, tailLines)
	if err != nil {
		return api.GuestSessionLogsResult{}, err
	}
	return api.GuestSessionLogsResult{Session: session, Stream: streamName, Path: path, Content: content}, nil
}

func (d *Daemon) SendToGuestSession(ctx context.Context, params api.GuestSessionSendParams) (api.GuestSessionSendResult, error) {
	vm, err := d.FindVM(ctx, params.VMIDOrName)
	if err != nil {
		return api.GuestSessionSendResult{}, err
	}
	session, err := d.findGuestSession(ctx, vm.ID, params.SessionIDOrName)
	if err != nil {
		return api.GuestSessionSendResult{}, err
	}
	if session.StdinMode != model.GuestSessionStdinPipe {
		return api.GuestSessionSendResult{}, errors.New("session does not have a stdin pipe")
	}
	if session.Status != model.GuestSessionStatusRunning {
		return api.GuestSessionSendResult{}, fmt.Errorf("session is not running (status=%s)", session.Status)
	}
	if !d.vmAlive(vm) {
		return api.GuestSessionSendResult{}, fmt.Errorf("vm %q is not running", vm.Name)
	}
	if len(params.Payload) == 0 {
		return api.GuestSessionSendResult{Session: session}, nil
	}
	client, err := d.dialGuest(ctx, net.JoinHostPort(vm.Runtime.GuestIP, "22"))
	if err != nil {
		return api.GuestSessionSendResult{}, fmt.Errorf("dial guest: %w", err)
	}
	defer client.Close()
	tmpPath := fmt.Sprintf("/tmp/banger-send-%s.bin", session.ID[:8])
	var uploadLog bytes.Buffer
	if err := client.UploadFile(ctx, tmpPath, 0o600, params.Payload, &uploadLog); err != nil {
		return api.GuestSessionSendResult{}, fmt.Errorf("upload payload: %w", err)
	}
	sendScript := fmt.Sprintf(
		"set -euo pipefail\ncat %s >> %s\nrm -f %s\n",
		sess.ShellQuote(tmpPath),
		sess.ShellQuote(sess.StdinPipePath(session.ID)),
		sess.ShellQuote(tmpPath),
	)
	var sendLog bytes.Buffer
	if err := client.RunScript(ctx, sendScript, &sendLog); err != nil {
		return api.GuestSessionSendResult{}, fmt.Errorf("send to session: %w: %s", err, strings.TrimSpace(sendLog.String()))
	}
	return api.GuestSessionSendResult{Session: session, BytesWritten: len(params.Payload)}, nil
}

func (d *Daemon) readGuestSessionLog(ctx context.Context, vm model.VMRecord, session model.GuestSession, stream string, tailLines int) (string, error) {
	if d.vmAlive(vm) {
		client, err := guest.Dial(ctx, net.JoinHostPort(vm.Runtime.GuestIP, "22"), d.config.SSHKeyPath)
		if err != nil {
			return "", err
		}
		defer client.Close()
		path := session.StdoutLogPath
		if stream == "stderr" {
			path = session.StderrLogPath
		}
		var output bytes.Buffer
		script := fmt.Sprintf("set -euo pipefail\nif [ -f %s ]; then tail -n %d %s; fi\n", sess.ShellQuote(path), tailLines, sess.ShellQuote(path))
		if err := client.RunScript(ctx, script, &output); err != nil {
			return "", sess.FormatStepError("read guest session log", err, output.String())
		}
		return output.String(), nil
	}
	runner := d.runner
	if runner == nil {
		runner = system.NewRunner()
	}
	workMount, cleanup, err := system.MountTempDir(ctx, runner, vm.Runtime.WorkDiskPath, false)
	if err != nil {
		return "", err
	}
	defer cleanup()
	logPath := filepath.Join(workMount, sess.RelativeStateDir(session.ID), stream+".log")
	return sess.TailFileContent(logPath, tailLines)
}