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
}

// sessionRegistry owns the live guest-session controller map. Its lock is
// independent of Daemon.mu so guest-session lookups do not contend with
// unrelated daemon state.
type sessionRegistry struct {
	mu     sync.Mutex
	byID   map[string]*guestSessionController
	closed bool
}

func newSessionRegistry() sessionRegistry {
	return sessionRegistry{byID: make(map[string]*guestSessionController)}
}

func (r *sessionRegistry) set(id string, controller *guestSessionController) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.closed {
		return
	}
	r.byID[id] = controller
}

func (r *sessionRegistry) claim(id string, controller *guestSessionController) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.closed {
		return false
	}
	if r.byID[id] != nil {
		return false
	}
	r.byID[id] = controller
	return true
}

func (r *sessionRegistry) get(id string) *guestSessionController {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.byID[id]
}

func (r *sessionRegistry) clear(id string) *guestSessionController {
	r.mu.Lock()
	defer r.mu.Unlock()
	controller := r.byID[id]
	delete(r.byID, id)
	return controller
}

func (r *sessionRegistry) closeAll() error {
	r.mu.Lock()
	controllers := make([]*guestSessionController, 0, len(r.byID))
	for _, controller := range r.byID {
		controllers = append(controllers, controller)
	}
	r.byID = nil
	r.closed = true
	r.mu.Unlock()
	var err error
	for _, controller := range controllers {
		err = errors.Join(err, controller.close())
	}
	return err
}

func (d *Daemon) setGuestSessionController(id string, controller *guestSessionController) {
	d.sessions.set(id, controller)
}

func (d *Daemon) claimGuestSessionController(id string, controller *guestSessionController) bool {
	return d.sessions.claim(id, controller)
}

func (d *Daemon) getGuestSessionController(id string) *guestSessionController {
	return d.sessions.get(id)
}

func (d *Daemon) clearGuestSessionController(id string) *guestSessionController {
	return d.sessions.clear(id)
}

func (d *Daemon) closeGuestSessionControllers() error {
	return d.sessions.closeAll()
}