natsrpc

package
v0.2.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Jun 1, 2026 License: MIT Imports: 19 Imported by: 0

Documentation

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func ConsumeChunks 

func ConsumeChunks(
	ctx context.Context,
	js jetstream.JetStream,
	streamName, consumerName, chunkSubject string,
	chunkParallelism int,
	sem *resourceprofile.ResizableSemaphore,
	scaler ChunkScaler,
	processFn func(ctx context.Context, chunk *ChunkMessage) error,
) error

ConsumeChunks pulls chunks from a shared durable consumer scoped to chunkSubject and processes them concurrently. sem caps concurrency; if nil, a fixed semaphore of size chunkParallelism is created. scaler is optional. processFn errors trigger a nak; ConsumeChunks returns once NumPending == 0.

func StartHeartbeat 

func StartHeartbeat(ctx context.Context, msg jetstream.Msg, interval time.Duration) context.CancelFunc

StartHeartbeat calls msg.InProgress() every interval to prevent JetStream redelivery while the message is being processed. The returned cancel stops it.

Types

type AgentInfo 

type AgentInfo struct {
	ID            string  `json:"id"`
	Name          string  `json:"name"`
	Version       string  `json:"version"`
	Uptime        string  `json:"uptime"`
	UptimeSeconds float64 `json:"uptime_seconds"`
	TasksRunning  int     `json:"tasks_running"`
}

AgentInfo contains agent identity and status.

type ChunkMessage 

type ChunkMessage struct {
	ChunkID         string   `json:"chunk_id"`
	Targets         []string `json:"targets"`
	PublicTemplates []string `json:"public_templates,omitempty"`
	// PrivateTemplates maps a template file name to base64-encoded YAML.
	PrivateTemplates map[string]string `json:"private_templates,omitempty"`
	ScanConfig       string            `json:"scan_configuration,omitempty"`

	// Enrichment fields (enumeration chunks).
	EnrichmentID   string `json:"enrichment_id,omitempty"`
	EnrichmentType string `json:"enrichment_type,omitempty"`
	EnumConfig     string `json:"enumeration_configuration,omitempty"`
}

ChunkMessage is a single unit of work decoded from the group stream. Scan chunks decode from ZSTD-compressed protobuf (ScanRequest); enumeration chunks from plain protobuf (AssetEnrichmentRequest).

type ChunkScaler 

type ChunkScaler interface {
	RecordChunkDuration(d time.Duration)
}

ChunkScaler receives chunk duration reports for adaptive scaling.

type DebugData 

type DebugData struct {
	Agent       AgentInfo   `json:"agent"`
	System      SystemInfo  `json:"system"`
	Process     ProcessInfo `json:"process"`
	Runtime     RuntimeInfo `json:"runtime"`
	ActiveTasks []TaskInfo  `json:"active_tasks,omitempty"`
}

DebugData is returned by the "debug" direct handler.

type GroupMetrics 

type GroupMetrics struct {
	// ChunksPending is the count of chunks sitting in the stream that have not
	// yet been delivered to any agent. Sum of NumPending across all
	// "chunks-*" consumers. This is the primary scale-up signal: a rising
	// value means the group cannot keep up with the inbound work.
	ChunksPending int64 `json:"chunks_pending"`

	// ChunksInflight is the count of chunks currently being processed by some
	// agent (delivered, heartbeating, not yet acked). Sum of NumAckPending
	// across all "chunks-*" consumers. Useful to compute true backlog
	// (pending + inflight) versus just queue depth.
	ChunksInflight int64 `json:"chunks_inflight"`

	// ActiveScans is the number of scans the group is currently working on.
	// Counted as "chunks-*" consumers with pending+inflight > 0.
	ActiveScans int64 `json:"active_scans"`

	// WorkPending is the number of scan/enumeration work messages queued to
	// this agent's work consumer but not yet pulled. Because every agent in
	// the group has its own durable work consumer that replays the entire
	// work stream, this value is the same across all agents in the group by
	// construction — so it doubles as a "scans queued, not yet started"
	// signal at the group level.
	WorkPending int64 `json:"work_pending"`

	// OldestConsumerAgeSec is the age in seconds of the oldest active chunk
	// consumer (one with non-zero pending+inflight). 0 if no active scans.
	// A growing value with non-zero ChunksPending indicates the queue is
	// stalled — agents may be stuck or the work is too slow to drain.
	OldestConsumerAgeSec int64 `json:"oldest_consumer_age_seconds"`

	// CollectedAt is when this snapshot was taken (RFC3339 UTC).
	CollectedAt string `json:"collected_at"`

	// CollectionDurationMs is how long the last collection pass took. Useful
	// to spot consumer-list latency growing out of hand.
	CollectionDurationMs int64 `json:"collection_duration_ms"`

	// CollectionErrors is the cumulative count of consumer-info read failures
	// since process start. Treat as a counter for Prometheus, even though we
	// store it as int64 here.
	CollectionErrors int64 `json:"collection_errors_total"`
}

GroupMetrics aggregates JetStream backlog across the whole agent group.

Every agent in the same group computes the same numbers — the metrics are stream-scoped, not per-agent. A customer HPA can scrape any one pod and get the group view; the operator can poll any one agent over NATS RPC.

Field semantics:

func (*GroupMetrics) String 

func (g *GroupMetrics) String() string

String renders a single-line summary, useful for logging.

type GroupMetricsCollector 

type GroupMetricsCollector struct {
	// contains filtered or unexported fields
}

GroupMetricsCollector walks the JetStream stream once per cache window and produces a GroupMetrics snapshot. Safe for concurrent use.

func NewGroupMetricsCollector 

func NewGroupMetricsCollector(js jetstream.JetStream, streamName, workConsumerName string, cacheTTL time.Duration) *GroupMetricsCollector

NewGroupMetricsCollector constructs a collector. Pass the JetStream handle (typically WorkerPool.JS()), the group stream name, and the agent's local work consumer name. cacheTTL controls how often Get triggers a fresh walk.

func (*GroupMetricsCollector) Get 

func (c *GroupMetricsCollector) Get(ctx context.Context) *GroupMetrics

Get returns the most recent snapshot, refreshing if the cache is stale. Concurrent callers during a refresh see the previous snapshot until the refresh completes — never block.

func (*GroupMetricsCollector) Snapshot 

func (c *GroupMetricsCollector) Snapshot() *GroupMetrics

Snapshot returns the last cached snapshot without triggering a refresh. Returns a zero-valued snapshot if Get has never run yet — callers can rely on a non-nil pointer.

type HTTPXRequest 

type HTTPXRequest struct {
	Target string `json:"target"`
}

HTTPXRequest is the payload for the "httpx" RPC method.

type HandlerFunc 

type HandlerFunc func(ctx context.Context, method string, data []byte) (any, error)

HandlerFunc processes an RPC request. method is the extracted suffix from the NATS subject (e.g. "httpx", "nuclei-retest"). data is the raw JSON body of the NATS message. Return value is marshalled into Response.Data on success.

type HealthCheckData 

type HealthCheckData struct {
	AgentID      string   `json:"agent_id"`
	AgentName    string   `json:"agent_name"`
	Version      string   `json:"version"`
	Uptime       string   `json:"uptime"`
	NumCPU       int      `json:"num_cpu"`
	MemTotalMB   uint64   `json:"mem_total_mb"`
	TasksRunning int      `json:"tasks_running"`
	ActiveScans  []string `json:"active_scans,omitempty"`
	ActiveEnums  []string `json:"active_enums,omitempty"`
	Idle         bool     `json:"idle"`                 // true if idle > 1 min
	IdleSince    string   `json:"idle_since,omitempty"` // RFC3339 timestamp if idle > 1 min
}

HealthCheckData is returned by the "health-check" broadcast handler.

type LogsRequest 

type LogsRequest struct {
	Offset int    `json:"offset"`          // 0 = oldest available entry
	Limit  int    `json:"limit"`           // max entries to return (default 100, max 500)
	Since  string `json:"since,omitempty"` // RFC3339 UTC
	Until  string `json:"until,omitempty"` // RFC3339 UTC
}

LogsRequest is the payload for the "logs" RPC method.

type LogsResponse 

type LogsResponse struct {
	Lines  []string `json:"lines"`
	Total  int      `json:"total"`
	Offset int      `json:"offset"`
	Limit  int      `json:"limit"`
}

LogsResponse is returned by the "logs" direct handler.

type MetricPoint 

type MetricPoint struct {
	T             string  `json:"t"`   // RFC3339 UTC
	CPU           float64 `json:"cpu"` // cpu_percent
	RSSMB         uint64  `json:"rss_mb"`
	HeapMB        uint64  `json:"heap_mb"` // heap_alloc_mb
	FDUsed        int     `json:"fd_used"`
	FDLimit       int     `json:"fd_limit"`
	MemTotalMB    uint64  `json:"mem_total_mb"`
	MemAvailMB    uint64  `json:"mem_avail_mb"`
	Goroutines    int     `json:"goroutines"`
	ActiveWorkers int32   `json:"active_workers"` // chunks currently being processed
	Capacity      int     `json:"capacity"`       // max concurrent chunks (parallelism)
}

MetricPoint is a single data point for time-series graphing.

type MetricsRequest 

type MetricsRequest struct {
	Range string `json:"range"`           // "5m","15m","30m","1h","3h","6h","24h","custom"
	Start string `json:"start,omitempty"` // RFC3339 UTC, required when range="custom"
	End   string `json:"end,omitempty"`   // RFC3339 UTC, required when range="custom"
}

MetricsRequest is the payload for the "metrics" RPC method.

type MetricsResponse 

type MetricsResponse struct {
	Range        string        `json:"range"`
	Since        string        `json:"since"`         // actual start (RFC3339 UTC)
	Until        string        `json:"until"`         // actual end (RFC3339 UTC)
	TotalSamples int           `json:"total_samples"` // raw count in DB for this range
	Returned     int           `json:"returned"`      // points returned after downsampling
	Points       []MetricPoint `json:"points"`
}

MetricsResponse is returned by the "metrics" direct handler.

type NucleiRetestRequest 

type NucleiRetestRequest struct {
	Targets         []string `json:"targets"`
	TemplateID      string   `json:"template_id"`
	TemplateEncoded string   `json:"template_encoded,omitempty"`
	TemplateURL     string   `json:"template_url,omitempty"`
	VulnID          string   `json:"vuln_id,omitempty"`
}

NucleiRetestRequest is the payload for the "nuclei-retest" RPC method. Template resolution priority: template_encoded > template_url > template_id.

type PortProbeRequest 

type PortProbeRequest struct {
	Host string `json:"host"`
	Port int    `json:"port"`
}

PortProbeRequest is the payload for the "port-probe" RPC method.

type ProcessInfo 

type ProcessInfo struct {
	PID        int     `json:"pid"`
	MemAllocMB float64 `json:"mem_alloc_mb"` // runtime.MemStats.Sys: total memory from OS
}

ProcessInfo contains pd-agent process resource usage via runtime.MemStats.

type Response 

type Response struct {
	Status string          `json:"status"`         // "ok" or "error"
	Data   json.RawMessage `json:"data,omitempty"` // handler-specific payload
	Error  string          `json:"error,omitempty"`
}

Response is the JSON envelope returned over NATS reply subjects.

type Router 

type Router struct {
	// contains filtered or unexported fields
}

Router dispatches NATS messages to registered handlers based on the method name extracted from the subject suffix. The lifecycle context is propagated to all handlers so they can be cancelled on agent shutdown.

func NewRouter 

func NewRouter(ctx context.Context) *Router

NewRouter creates a Router with a lifecycle context. When ctx is cancelled, in-flight handlers receive cancellation and should exit promptly.

func (*Router) Handle 

func (r *Router) Handle(method string, fn HandlerFunc)

Handle registers a handler for the given method name. Method names are case-sensitive and must not contain dots.

func (*Router) SubscribeBroadcast 

func (r *Router) SubscribeBroadcast(nc *nats.Conn, subjectPrefix string) (*nats.Subscription, error)

SubscribeBroadcast subscribes to subjectPrefix.> using a plain Subscribe (NOT QueueSubscribe) so that every agent receives every message. This is intentional — broadcast messages like health-check need all agents to respond.

func (*Router) SubscribeDirect 

func (r *Router) SubscribeDirect(nc *nats.Conn, subjectPrefix string) (*nats.Subscription, error)

SubscribeDirect subscribes to subjectPrefix.> for a single agent using a plain Subscribe. The subject prefix includes the agent ID (e.g., groupPrefix.direct.<agentID>), so only this agent receives messages. Kept as a separate method from SubscribeBroadcast for call-site readability.

func (*Router) SubscribeRequests 

func (r *Router) SubscribeRequests(nc *nats.Conn, subjectPrefix, queueGroup string) (*nats.Subscription, error)

SubscribeRequests subscribes to subjectPrefix.> using QueueSubscribe so that only one agent in the queue group handles each request.

type RuntimeInfo 

type RuntimeInfo struct {
	GoVersion    string  `json:"go_version"`
	NumGoroutine int     `json:"num_goroutine"`
	HeapAllocMB  float64 `json:"heap_alloc_mb"`
	HeapInuseMB  float64 `json:"heap_inuse_mb"`
	StackInuseMB float64 `json:"stack_inuse_mb"`
	TotalAllocMB float64 `json:"total_alloc_mb"`
	NumGC        uint32  `json:"num_gc"`
	LastGC       string  `json:"last_gc,omitempty"`
}

RuntimeInfo contains Go runtime metrics.

type SystemInfo 

type SystemInfo struct {
	OS       string `json:"os"`
	Arch     string `json:"arch"`
	NumCPU   int    `json:"num_cpu"`
	Hostname string `json:"hostname"`
}

SystemInfo contains host-level resource info.

type TaskInfo 

type TaskInfo struct {
	Type      string `json:"type"` // "scan" or "enumeration"
	TaskID    string `json:"task_id"`
	StartedAt string `json:"started_at"` // RFC3339 UTC
}

TaskInfo describes an active scan or enumeration.

type WorkHandler 

type WorkHandler func(ctx context.Context, msg jetstream.Msg, work *WorkMessage) error

WorkHandler processes a parsed WorkMessage. The raw jetstream.Msg is supplied so the handler can ack or nak directly.

type WorkMessage 

type WorkMessage struct {
	Type             string   `json:"type"`           // "scan" or "enumeration"
	ScanID           string   `json:"scan_id"`        // scan_id or enumeration_id
	ChunkSubject     string   `json:"chunk_subject"`  // subject filter for chunks
	ChunkConsumer    string   `json:"chunk_consumer"` // shared consumer name
	ChunkCount       int      `json:"chunk_count,omitempty"`
	Config           string   `json:"config,omitempty"`        // base64 scan configuration
	ReportConfig     string   `json:"report_config,omitempty"` // base64 nuclei reporting (-rc) configuration
	HistoryID        int64    `json:"history_id,omitempty"`
	Steps            []string `json:"steps,omitempty"` // enumeration steps
	EnumerationPorts string   `json:"enumeration_ports,omitempty"`
}

WorkMessage notifies agents of a new scan or enumeration to process. All work and chunks share a single group-level stream; consumers scope reads via FilterSubject.

type WorkerPool 

type WorkerPool struct {
	// contains filtered or unexported fields
}

WorkerPool pulls work messages from a JetStream consumer and dispatches them to a handler. Concurrency is bounded by parallelism and the consumer's MaxAckPending.

func NewWorkerPool 

func NewWorkerPool(nc *nats.Conn, streamName, consumerName, groupPrefix string, parallelism int, handler WorkHandler) (*WorkerPool, error)

NewWorkerPool creates a durable pull consumer filtered to groupPrefix.work.>, with MaxAckPending=parallelism so JetStream provides the backpressure.

func (*WorkerPool) ActiveWorkers 

func (wp *WorkerPool) ActiveWorkers() int32

ActiveWorkers returns the number of workers currently processing a message.

func (*WorkerPool) IdleSince 

func (wp *WorkerPool) IdleSince() time.Time

IdleSince returns the last-activity time when the pool is idle, or zero when it is actively processing or has never processed work.

func (*WorkerPool) JS 

func (wp *WorkerPool) JS() jetstream.JetStream

JS returns the underlying JetStream instance.

func (*WorkerPool) LastActivity 

func (wp *WorkerPool) LastActivity() time.Time

LastActivity returns the time of the last work activity, or zero if none.

func (*WorkerPool) Run 

func (wp *WorkerPool) Run(ctx context.Context)

Run starts the worker goroutines and blocks until ctx is cancelled.

func (*WorkerPool) Stop 

func (wp *WorkerPool) Stop()

Stop signals workers to stop and waits for in-flight work to complete.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.