banger/internal/daemon/session_controller.go

package daemon

import (
	"errors"
	"io"
	"net"
	"sync"

	"banger/internal/guest"
	"banger/internal/sessionstream"
)

type guestSessionController struct {
	stream    *guest.StreamSession
	streams   []*guest.StreamSession
	stdin     io.WriteCloser
	attachMu  sync.Mutex
	attach    net.Conn
	writeMu   sync.Mutex
	closeOnce sync.Once
}

func (c *guestSessionController) setAttach(conn net.Conn) error {
	c.attachMu.Lock()
	defer c.attachMu.Unlock()
	if c.attach != nil {
		return errors.New("session already has an active attach")
	}
	c.attach = conn
	return nil
}

func (c *guestSessionController) clearAttach(conn net.Conn) {
	c.attachMu.Lock()
	defer c.attachMu.Unlock()
	if c.attach == conn {
		c.attach = nil
	}
}

func (c *guestSessionController) writeFrame(channel byte, payload []byte) {
	c.attachMu.Lock()
	conn := c.attach
	c.attachMu.Unlock()
	if conn == nil {
		return
	}
	c.writeMu.Lock()
	err := sessionstream.WriteFrame(conn, channel, payload)
	c.writeMu.Unlock()
	if err != nil {
		_ = conn.Close()
		c.clearAttach(conn)
	}
}

func (c *guestSessionController) writeControl(message sessionstream.ControlMessage) {
	c.attachMu.Lock()
	conn := c.attach
	c.attachMu.Unlock()
	if conn == nil {
		return
	}
	c.writeMu.Lock()
	err := sessionstream.WriteControl(conn, message)
	c.writeMu.Unlock()
	if err != nil {
		_ = conn.Close()
		c.clearAttach(conn)
	}
}

func (c *guestSessionController) close() error {
	if c == nil {
		return nil
	}
	var err error
	c.closeOnce.Do(func() {
		c.attachMu.Lock()
		conn := c.attach
		c.attach = nil
		c.attachMu.Unlock()
		if conn != nil {
			err = errors.Join(err, conn.Close())
		}
		if c.stdin != nil {
			err = errors.Join(err, c.stdin.Close())
		}
		if c.stream != nil {
			err = errors.Join(err, c.stream.Close())
		}
		for _, stream := range c.streams {
			if stream != nil {
				err = errors.Join(err, stream.Close())
			}
		}
	})
	return err
}

type guestSessionStateSnapshot struct {
	Status    string
	GuestPID  int
	ExitCode  *int
	Alive     bool
	LastError string
}

func (d *Daemon) setGuestSessionController(id string, controller *guestSessionController) {
	d.mu.Lock()
	defer d.mu.Unlock()
	d.sessionControllers[id] = controller
}

func (d *Daemon) claimGuestSessionController(id string, controller *guestSessionController) bool {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.sessionControllers[id] != nil {
		return false
	}
	d.sessionControllers[id] = controller
	return true
}

func (d *Daemon) getGuestSessionController(id string) *guestSessionController {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.sessionControllers[id]
}

func (d *Daemon) clearGuestSessionController(id string) *guestSessionController {
	d.mu.Lock()
	defer d.mu.Unlock()
	controller := d.sessionControllers[id]
	delete(d.sessionControllers, id)
	return controller
}

func (d *Daemon) closeGuestSessionControllers() error {
	d.mu.Lock()
	controllers := make([]*guestSessionController, 0, len(d.sessionControllers))
	for _, controller := range d.sessionControllers {
		controllers = append(controllers, controller)
	}
	d.sessionControllers = nil
	d.mu.Unlock()
	var err error
	for _, controller := range controllers {
		err = errors.Join(err, controller.close())
	}
	return err
}