lpra/internal/server/handlers_agent.go

package server

import (
	"encoding/json"
	"net/http"
	"time"

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

// Long-poll hold window. Must stay below nginx proxy_read_timeout (120s).
const pollHold = 25 * time.Second

func writeJSON(w http.ResponseWriter, v any) {
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(v)
}

// POST /api/agent/enroll - exchange the shared enroll secret for a per-agent id+token.
func (s *Server) handleEnroll(w http.ResponseWriter, r *http.Request) {
	secret := r.Header.Get("X-Enroll-Secret")
	if !constantTimeEqual(secret, s.db.EnrollSecret()) {
		http.Error(w, "forbidden", http.StatusForbidden)
		return
	}
	var req proto.EnrollRequest
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<16)).Decode(&req); err != nil {
		http.Error(w, "bad request", http.StatusBadRequest)
		return
	}
	id, token, err := s.enroll(req.Name, req.Hostname, req.OS, req.Arch, req.Version, req.User)
	if err != nil {
		http.Error(w, "internal error", http.StatusInternalServerError)
		return
	}
	writeJSON(w, proto.EnrollResponse{AgentID: id, Token: token})
}

// GET /api/agent/poll - control channel. Held until a job is queued or timeout.
func (s *Server) handlePoll(w http.ResponseWriter, r *http.Request) {
	a := s.authenticateAgent(r)
	if a == nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	// The agent sends its version and OS user occasionally; record them when
	// present. Agents older than 0.0.2 send no user, so keep what we have.
	s.viaDirectA.Store(a.ID, s.viaDirect(r))
	v := r.Header.Get("X-Agent-Version")
	u := firstNonEmpty(r.Header.Get("X-Agent-User"), a.RunUser)
	if v != "" && (v != a.Version || u != a.RunUser) {
		_ = s.db.TouchAgentInfo(a.ID, v, u)
	} else {
		_ = s.db.TouchAgent(a.ID)
	}
	msg := s.hub.WaitControl(r.Context().Done(), a.ID, pollHold)
	// Set after the hold so a setting changed meanwhile reaches the agent now.
	s.setTunnelHeader(w)
	w.Header().Set(proto.TransportHeader, s.transportOf(a.ID))
	writeJSON(w, msg)
}

// tunnelAddrHeader carries the direct-tunnel address ("host:port", host empty
// = the panel's host) to agents, on every poll and on /api/agent/config.
const tunnelAddrHeader = "X-Tunnel-Addr"

func (s *Server) setTunnelHeader(w http.ResponseWriter) {
	if addr := s.tunnelAdvert(); addr != "" {
		w.Header().Set(tunnelAddrHeader, addr)
	}
	if d := s.db.DirectAddress(); d != "" {
		w.Header().Set(proto.FallbackHeader, d)
		w.Header().Set(proto.MainHeader, s.db.HostAddress())
	}
}

// GET /api/agent/config - settings an agent picks up when it connects, without
// waiting out a poll hold.
func (s *Server) handleAgentConfig(w http.ResponseWriter, r *http.Request) {
	a := s.authenticateAgent(r)
	if a == nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	s.setTunnelHeader(w)
	w.Header().Set(proto.TransportHeader, s.transportOf(a.ID))
	w.WriteHeader(http.StatusNoContent)
}

// GET /api/agent/task/{id} - the agent fetches the task to run.
func (s *Server) handleTaskFetch(w http.ResponseWriter, r *http.Request) {
	if s.authenticateAgent(r) == nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	env, ok := s.hub.GetTask(r.PathValue("id"))
	if !ok {
		http.Error(w, "no task", http.StatusGone)
		return
	}
	writeJSON(w, env)
}

// POST /api/agent/task/{id}/result - the agent returns the task result.
func (s *Server) handleTaskResult(w http.ResponseWriter, r *http.Request) {
	if s.authenticateAgent(r) == nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	var res proto.TaskResult
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 64<<20)).Decode(&res); err != nil {
		http.Error(w, "bad request", http.StatusBadRequest)
		return
	}
	s.hub.CompleteTask(r.PathValue("id"), res)
	w.WriteHeader(http.StatusNoContent)
}

// GET /api/agent/term/{sid}/in - downstream (panel -> agent) long-poll.
func (s *Server) handleTermIn(w http.ResponseWriter, r *http.Request) {
	if s.authenticateAgent(r) == nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	sess := s.hub.Session(r.PathValue("sid"))
	if sess == nil {
		http.Error(w, "no session", http.StatusGone)
		return
	}
	frames, alive := sess.PollIn(pollHold)
	if !alive {
		http.Error(w, "closed", http.StatusGone)
		return
	}
	writeJSON(w, proto.InBatch{Frames: frames})
}

// POST /api/agent/term/{sid}/out - upstream (agent PTY -> panel).
func (s *Server) handleTermOut(w http.ResponseWriter, r *http.Request) {
	if s.authenticateAgent(r) == nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	sess := s.hub.Session(r.PathValue("sid"))
	if sess == nil {
		http.Error(w, "no session", http.StatusGone)
		return
	}
	var msg proto.OutMessage
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&msg); err != nil {
		http.Error(w, "bad request", http.StatusBadRequest)
		return
	}
	if !sess.PushOut(msg) {
		http.Error(w, "closed", http.StatusGone)
		return
	}
	if msg.Exit {
		// End the stream now, but keep the session reachable for a moment so
		// the panel can still collect the last output (e.g. why the shell
		// failed to start).
		sess.close()
		time.AfterFunc(pollHold, func() { s.hub.CloseSession(sess.ID) })
	}
	w.WriteHeader(http.StatusNoContent)
}