lpra/internal/server/hub.go

package server

import (
	"sync"
	"time"

	"github.com/x-grk/longpolling-remote-access/internal/proto"
)

// Hub is the in-memory broker between agents (HTTP long-polling) and the panel.
// It carries three kinds of work to an agent through its single control poll:
// terminal sessions, port-forward streams, and one-shot tasks (file manager).
// Nothing here is persisted.
type Hub struct {
	mu sync.Mutex

	// jobs holds control messages queued for an agent but not yet delivered via
	// its control poll.
	jobs map[string][]proto.ControlMessage
	// ctrlWaiters wakes a blocked control poll the moment a job is enqueued.
	ctrlWaiters map[string]chan struct{}

	sessions map[string]*Session
	tasks    map[string]*taskEntry

	// Presence: control polls held open right now per agent, and when each
	// agent's last one ended. started is when this hub came up.
	polls    map[string]int
	lastPoll map[string]time.Time
	started  time.Time
}

// taskEntry is a pending one-shot task awaiting the agent's result.
type taskEntry struct {
	env    proto.TaskEnvelope
	result chan proto.TaskResult
}

// Session is one live byte stream between the panel and an agent (a PTY terminal
// or a raw TCP forward - the hub treats both identically).
type Session struct {
	ID      string
	AgentID string

	mu       sync.Mutex
	inBuf    []proto.InFrame
	inSignal chan struct{} // size-1; pulsed when inBuf gains data
	inSeq    uint64        // number of the last browser message taken

	out chan proto.OutMessage // agent -> panel

	closed    chan struct{}
	closeOnce sync.Once
}

func NewHub() *Hub {
	return &Hub{
		jobs:        map[string][]proto.ControlMessage{},
		ctrlWaiters: map[string]chan struct{}{},
		sessions:    map[string]*Session{},
		tasks:       map[string]*taskEntry{},
		polls:       map[string]int{},
		lastPoll:    map[string]time.Time{},
		started:     time.Now(),
	}
}

// enqueueJob appends a control message for an agent and wakes its poll.
func (h *Hub) enqueueJob(agentID string, msg proto.ControlMessage) {
	h.mu.Lock()
	h.jobs[agentID] = append(h.jobs[agentID], msg)
	if w := h.ctrlWaiters[agentID]; w != nil {
		select {
		case w <- struct{}{}:
		default:
		}
	}
	h.mu.Unlock()
}

func (h *Hub) newSession(agentID string) *Session {
	s := &Session{
		ID:       randToken()[:16],
		AgentID:  agentID,
		inSignal: make(chan struct{}, 1),
		out:      make(chan proto.OutMessage, 1024),
		closed:   make(chan struct{}),
	}
	h.mu.Lock()
	h.sessions[s.ID] = s
	h.mu.Unlock()
	return s
}

// CreateSession registers a terminal session and queues it for the agent. user
// is the OS user to run the shell as ("" = the agent's own user).
func (h *Hub) CreateSession(agentID, user string) *Session {
	s := h.newSession(agentID)
	h.enqueueJob(agentID, proto.ControlMessage{Type: proto.CtrlOpenTerminal, SessionID: s.ID, User: user})
	return s
}

// CreateForward registers a port-forward stream (agent dials 127.0.0.1:port).
func (h *Hub) CreateForward(agentID string, port int) *Session {
	s := h.newSession(agentID)
	h.enqueueJob(agentID, proto.ControlMessage{Type: proto.CtrlOpenForward, SessionID: s.ID, Port: port})
	return s
}

// pollGrace is how long an agent stays online between two control polls: it
// re-polls at once, so this only has to cover a slow round trip or one retry.
const pollGrace = 10 * time.Second

// Online reports whether the agent is connected: it holds a control poll open
// or has just finished one. ok is false while that cannot be told yet (right
// after a start, before agents had time to come back).
func (h *Hub) Online(agentID string) (online, ok bool) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.polls[agentID] > 0 {
		return true, true
	}
	if t, seen := h.lastPoll[agentID]; seen {
		return time.Since(t) < pollGrace, true
	}
	return false, time.Since(h.started) > pollHold+pollGrace
}

// Forget drops what the hub keeps about a deleted agent.
func (h *Hub) Forget(agentID string) {
	h.mu.Lock()
	delete(h.jobs, agentID)
	delete(h.ctrlWaiters, agentID)
	delete(h.lastPoll, agentID)
	h.mu.Unlock()
}

// WaitControl blocks until a job is queued for agentID, the timeout elapses
// or the agent hangs up (done). The agent counts as online meanwhile.
func (h *Hub) WaitControl(done <-chan struct{}, agentID string, timeout time.Duration) proto.ControlMessage {
	h.mu.Lock()
	h.polls[agentID]++
	h.mu.Unlock()
	defer func() {
		h.mu.Lock()
		if h.polls[agentID]--; h.polls[agentID] <= 0 {
			delete(h.polls, agentID)
		}
		// A poll cut off by the agent is not a sign of life.
		select {
		case <-done:
		default:
			h.lastPoll[agentID] = time.Now()
		}
		h.mu.Unlock()
	}()
	t := time.NewTimer(timeout)
	defer t.Stop()
	deadline := t.C
	for {
		h.mu.Lock()
		if q := h.jobs[agentID]; len(q) > 0 {
			msg := q[0]
			h.jobs[agentID] = q[1:]
			h.mu.Unlock()
			return msg
		}
		w := h.ctrlWaiters[agentID]
		if w == nil {
			w = make(chan struct{}, 1)
			h.ctrlWaiters[agentID] = w
		}
		h.mu.Unlock()

		select {
		case <-w:
		case <-deadline:
			return proto.ControlMessage{Type: proto.CtrlIdle}
		case <-done:
			// The job, if one comes, stays queued for the next poll.
			return proto.ControlMessage{Type: proto.CtrlIdle}
		}
	}
}

func (h *Hub) Session(id string) *Session {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.sessions[id]
}

func (h *Hub) CloseSession(id string) {
	h.mu.Lock()
	s := h.sessions[id]
	delete(h.sessions, id)
	h.mu.Unlock()
	if s != nil {
		s.close()
	}
}

// --- one-shot tasks ---

// CreateTask queues a task and returns its id plus a channel for the result.
func (h *Hub) CreateTask(agentID, typ string, args []byte) (string, chan proto.TaskResult) {
	id := randToken()[:16]
	entry := &taskEntry{
		env:    proto.TaskEnvelope{Type: typ, Args: args},
		result: make(chan proto.TaskResult, 1),
	}
	h.mu.Lock()
	h.tasks[id] = entry
	h.mu.Unlock()
	h.enqueueJob(agentID, proto.ControlMessage{Type: proto.CtrlTask, TaskID: id})
	return id, entry.result
}

// GetTask returns the envelope the agent should execute.
func (h *Hub) GetTask(id string) (proto.TaskEnvelope, bool) {
	h.mu.Lock()
	defer h.mu.Unlock()
	e := h.tasks[id]
	if e == nil {
		return proto.TaskEnvelope{}, false
	}
	return e.env, true
}

// CompleteTask delivers a result to the waiting panel handler.
func (h *Hub) CompleteTask(id string, res proto.TaskResult) {
	h.mu.Lock()
	e := h.tasks[id]
	h.mu.Unlock()
	if e != nil {
		select {
		case e.result <- res:
		default:
		}
	}
}

// DiscardTask drops a task from the registry.
func (h *Hub) DiscardTask(id string) {
	h.mu.Lock()
	delete(h.tasks, id)
	h.mu.Unlock()
}

func (s *Session) close() { s.closeOnce.Do(func() { close(s.closed) }) }

// Done returns a channel closed when the session ends.
func (s *Session) Done() <-chan struct{} { return s.closed }

// TakeSeq reports whether the browser message numbered seq is new, and takes
// note of it. The browser sends them one at a time in rising order and sends
// one again when it got no answer, so a number not above the last is a repeat.
func (s *Session) TakeSeq(seq uint64) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	if seq <= s.inSeq {
		return false
	}
	s.inSeq = seq
	return true
}

// PushIn appends a downstream frame (panel -> agent) and wakes the agent poll.
func (s *Session) PushIn(f proto.InFrame) {
	s.mu.Lock()
	s.inBuf = append(s.inBuf, f)
	s.mu.Unlock()
	select {
	case s.inSignal <- struct{}{}:
	default:
	}
}

// Unpoll puts frames taken by PollIn back in front, for frames that could not
// be delivered.
func (s *Session) Unpoll(frames []proto.InFrame) {
	s.mu.Lock()
	s.inBuf = append(append([]proto.InFrame{}, frames...), s.inBuf...)
	s.mu.Unlock()
	select {
	case s.inSignal <- struct{}{}:
	default:
	}
}

// PollIn drains queued downstream frames, blocking up to timeout for the first.
func (s *Session) PollIn(timeout time.Duration) ([]proto.InFrame, bool) {
	deadline := time.After(timeout)
	for {
		s.mu.Lock()
		if len(s.inBuf) > 0 {
			b := s.inBuf
			s.inBuf = nil
			s.mu.Unlock()
			return b, true
		}
		s.mu.Unlock()

		select {
		case <-s.inSignal:
		case <-deadline:
			return nil, true
		case <-s.closed:
			return nil, false
		}
	}
}

// PushOut delivers agent output to the panel side. Returns false if closed.
func (s *Session) PushOut(m proto.OutMessage) bool {
	select {
	case s.out <- m:
		return true
	case <-s.closed:
		return false
	}
}

func (s *Session) Out() <-chan proto.OutMessage { return s.out }

// PollOut drains queued upstream messages, blocking up to timeout for the first.
func (s *Session) PollOut(timeout time.Duration) ([]proto.OutMessage, bool) {
	return s.PollOutUntil(nil, timeout)
}

// PollOutUntil is PollOut that also gives up when done closes, taking
// nothing: for a reader that has gone, so no output is lost on it.
func (s *Session) PollOutUntil(done <-chan struct{}, timeout time.Duration) ([]proto.OutMessage, bool) {
	var batch []proto.OutMessage
	select {
	case m := <-s.out:
		batch = append(batch, m)
	case <-done:
		return nil, true
	case <-time.After(timeout):
		return nil, true
	case <-s.closed:
		// Hand over output queued before the close (e.g. the reason a shell
		// failed to start) instead of dropping it.
		select {
		case m := <-s.out:
			batch = append(batch, m)
		default:
			return nil, false
		}
	}
	for {
		select {
		case m := <-s.out:
			batch = append(batch, m)
		default:
			return batch, true
		}
	}
}