banger/internal/daemon/guest_sessions.go

package daemon

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

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

type guestSSHClient interface {
	Close() error
	RunScript(context.Context, string, io.Writer) error
	RunScriptOutput(context.Context, string) ([]byte, error)
	UploadFile(context.Context, string, os.FileMode, []byte, io.Writer) error
	StreamTar(context.Context, string, string, io.Writer) error
	StreamTarEntries(context.Context, string, []string, string, io.Writer) error
}

func (d *Daemon) waitForGuestSSH(ctx context.Context, address string, interval time.Duration) error {
	if d != nil && d.guestWaitForSSH != nil {
		return d.guestWaitForSSH(ctx, address, d.config.SSHKeyPath, interval)
	}
	return guest.WaitForSSH(ctx, address, d.config.SSHKeyPath, interval)
}

func (d *Daemon) dialGuest(ctx context.Context, address string) (guestSSHClient, error) {
	if d != nil && d.guestDial != nil {
		return d.guestDial(ctx, address, d.config.SSHKeyPath)
	}
	return guest.Dial(ctx, address, d.config.SSHKeyPath)
}

func (d *Daemon) waitForGuestSessionReadyHook(ctx context.Context, vm model.VMRecord, s model.GuestSession) (model.GuestSession, error) {
	if d != nil && d.waitForGuestSessionReady != nil {
		return d.waitForGuestSessionReady(ctx, vm, s)
	}
	return d.waitForGuestSessionReadyDefault(ctx, vm, s)
}

func (d *Daemon) waitForGuestSessionReadyDefault(ctx context.Context, vm model.VMRecord, s model.GuestSession) (model.GuestSession, error) {
	for {
		updated, err := d.refreshGuestSession(ctx, vm, s)
		if err == nil {
			s = updated
			if s.GuestPID != 0 || s.ExitCode != nil || s.Status == model.GuestSessionStatusRunning || s.Status == model.GuestSessionStatusFailed || s.Status == model.GuestSessionStatusExited {
				return s, nil
			}
		}
		select {
		case <-ctx.Done():
			return s, ctx.Err()
		case <-time.After(100 * time.Millisecond):
		}
	}
}

func (d *Daemon) refreshGuestSession(ctx context.Context, vm model.VMRecord, s model.GuestSession) (model.GuestSession, error) {
	if s.Status != model.GuestSessionStatusStarting && s.Status != model.GuestSessionStatusRunning && s.Status != model.GuestSessionStatusStopping {
		return s, nil
	}
	snapshot, err := d.inspectGuestSessionState(ctx, vm, s)
	if err != nil {
		return s, err
	}
	original := s
	session.ApplyStateSnapshot(&s, snapshot, vm.State == model.VMStateRunning && system.ProcessRunning(vm.Runtime.PID, vm.Runtime.APISockPath))
	if session.StateChanged(original, s) {
		s.UpdatedAt = model.Now()
		if err := d.store.UpsertGuestSession(ctx, s); err != nil {
			return s, err
		}
	}
	return s, nil
}

func (d *Daemon) inspectGuestSessionState(ctx context.Context, vm model.VMRecord, s model.GuestSession) (session.StateSnapshot, error) {
	if vm.State == model.VMStateRunning && system.ProcessRunning(vm.Runtime.PID, vm.Runtime.APISockPath) {
		client, err := guest.Dial(ctx, net.JoinHostPort(vm.Runtime.GuestIP, "22"), d.config.SSHKeyPath)
		if err != nil {
			return session.StateSnapshot{}, err
		}
		defer client.Close()
		var output bytes.Buffer
		if err := client.RunScript(ctx, session.InspectScript(s.ID), &output); err != nil {
			return session.StateSnapshot{}, session.FormatStepError("inspect guest session state", err, output.String())
		}
		return session.ParseState(output.String())
	}
	return d.inspectGuestSessionStateFromWorkDisk(ctx, vm, s.ID)
}

func (d *Daemon) inspectGuestSessionStateFromWorkDisk(ctx context.Context, vm model.VMRecord, sessionID string) (session.StateSnapshot, error) {
	runner := d.runner
	if runner == nil {
		runner = system.NewRunner()
	}
	workMount, cleanup, err := system.MountTempDir(ctx, runner, vm.Runtime.WorkDiskPath, false)
	if err != nil {
		return session.StateSnapshot{}, err
	}
	defer cleanup()
	stateDir := filepath.Join(workMount, session.RelativeStateDir(sessionID))
	return session.InspectStateFromDir(stateDir)
}

func (d *Daemon) findGuestSession(ctx context.Context, vmID, idOrName string) (model.GuestSession, error) {
	if strings.TrimSpace(idOrName) == "" {
		return model.GuestSession{}, errors.New("session id or name is required")
	}
	if s, err := d.store.GetGuestSession(ctx, vmID, idOrName); err == nil {
		return s, nil
	}
	sessions, err := d.store.ListGuestSessionsByVM(ctx, vmID)
	if err != nil {
		return model.GuestSession{}, err
	}
	matches := make([]model.GuestSession, 0, 1)
	for _, s := range sessions {
		if strings.HasPrefix(s.ID, idOrName) || strings.HasPrefix(s.Name, idOrName) {
			matches = append(matches, s)
		}
	}
	switch len(matches) {
	case 0:
		return model.GuestSession{}, fmt.Errorf("session %q not found", idOrName)
	case 1:
		return matches[0], nil
	default:
		return model.GuestSession{}, fmt.Errorf("multiple sessions match %q", idOrName)
	}
}