feat: initial commit
This commit is contained in:
redtea committed
1 parent
353f62e257
commit
c6852c6f1b
62 files changed
+5911
-2
No files matched your search
@@ -0,0 +1,258 @@
|
||||
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
|
||||
}
|
||||
|
||||
// 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
|
||||
|
||||
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{},
|
||||
}
|
||||
}
|
||||
|
||||
// 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.
|
||||
func (h *Hub) CreateSession(agentID string) *Session {
|
||||
s := h.newSession(agentID)
|
||||
h.enqueueJob(agentID, proto.ControlMessage{Type: proto.CtrlOpenTerminal, SessionID: s.ID})
|
||||
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
|
||||
}
|
||||
|
||||
// WaitControl blocks until a job is queued for agentID or timeout elapses.
|
||||
func (h *Hub) WaitControl(agentID string, timeout time.Duration) proto.ControlMessage {
|
||||
deadline := time.After(timeout)
|
||||
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}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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 }
|
||||
|
||||
// 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:
|
||||
}
|
||||
}
|
||||
|
||||
// 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) {
|
||||
var batch []proto.OutMessage
|
||||
select {
|
||||
case m := <-s.out:
|
||||
batch = append(batch, m)
|
||||
case <-time.After(timeout):
|
||||
return nil, true
|
||||
case <-s.closed:
|
||||
return nil, false
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case m := <-s.out:
|
||||
batch = append(batch, m)
|
||||
default:
|
||||
return batch, true
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user