HellreigN/backend/internal/grpcsrv/commander/commander.go

package commander

import (
	"context"
	"fmt"
	"io"
	"log"
	"sync"

	"gitea.d3m0k1d.ru/d3m0k1d/HellreigN/backend/internal/models"
	"gitea.d3m0k1d.ru/d3m0k1d/HellreigN/proto/proto"
	"golang.org/x/sync/errgroup"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/stats"
)

// Commander handles command execution on connected agents.
type Commander struct {
	proto.UnimplementedCommanderServer
	tracker *ConnTracker
	jobber  Jobber
}

// Jobber persists job state.
type Jobber interface {
	InitJob(ctx context.Context, agentID string, job models.JobForInsert) (int64, error)
	UpdateJobInDB(ctx context.Context, jid int64, msg models.JobForUpdate) (models.Job, error)
}

func New(jobber Jobber, tracker *ConnTracker) *Commander {
	return &Commander{
		jobber:  jobber,
		tracker: tracker,
	}
}

// Agent represents a connected agent with an active bidirectional stream.
type Agent struct {
	bidi   grpc.BidiStreamingServer[proto.FinishedCommand, proto.Command]
	in     chan *proto.Command
	jobs   map[int64]Job
	jobber Jobber
	ctx    context.Context
	aid    string

	Token    string
	Label    string
	Services []string
}

type JobOut struct {
	fc  models.Job
	err error
}

type Job struct {
	out chan JobOut
}

// ConnTracker tracks connected agents and handles cleanup on disconnect.
// It implements grpc.StatsHandler for disconnect detection.
type ConnTracker struct {
	mu     sync.RWMutex
	agents map[string]*Agent
}

// GetAgentByLabel searches for an agent by its human-readable label.
func (self *ConnTracker) GetAgentByLabel(label string) (agent Agent, ok bool) {
	self.mu.RLock()
	defer self.mu.RUnlock()
	for _, a := range self.agents {
		if a.Label == label {
			return *a, true
		}
	}
	return
}

func NewConnTracker() *ConnTracker {
	return &ConnTracker{
		agents: make(map[string]*Agent),
	}
}

func (t *ConnTracker) Register(aid string, agent *Agent) {
	t.mu.Lock()
	t.agents[aid] = agent
	t.mu.Unlock()
	log.Printf("[conntracker] agent registered: %s", aid)
}

func (t *ConnTracker) Unregister(aid string) {
	t.mu.Lock()
	delete(t.agents, aid)
	t.mu.Unlock()
	log.Printf("[conntracker] agent unregistered: %s", aid)
}

func (t *ConnTracker) GetAgent(aid string) (*Agent, bool) {
	t.mu.RLock()
	defer t.mu.RUnlock()
	a, ok := t.agents[aid]
	return a, ok
}

func (t *ConnTracker) Agents() []*Agent {
	t.mu.RLock()
	defer t.mu.RUnlock()
	result := make([]*Agent, 0, len(t.agents))
	for _, a := range t.agents {
		result = append(result, a)
	}
	return result
}

// grpc.StatsHandler implementation.

func (t *ConnTracker) TagRPC(ctx context.Context, _ *stats.RPCTagInfo) context.Context {
	return ctx
}

func (t *ConnTracker) HandleRPC(ctx context.Context, _ stats.RPCStats) {}

func (t *ConnTracker) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context {
	return ctx
}

func (t *ConnTracker) HandleConn(ctx context.Context, s stats.ConnStats) {
	switch s.(type) {
	case *stats.ConnEnd:
		md, ok := metadata.FromIncomingContext(ctx)
		if !ok {
			return
		}
		aidVals := md["agentid"]
		if len(aidVals) == 0 {
			return
		}
		t.Unregister(aidVals[0])
	}
}

// Stream handles a new agent connection and runs the send/recv loops.
func (c *Commander) Stream(
	bidi grpc.BidiStreamingServer[proto.FinishedCommand, proto.Command],
) error {
	md, ok := metadata.FromIncomingContext(bidi.Context())
	if !ok {
		return fmt.Errorf("no metadata in context")
	}
	aidVals := md["agentid"]
	if len(aidVals) == 0 {
		return fmt.Errorf("agentid metadata missing")
	}
	aid := aidVals[0]

	var label string
	if vals := md["label"]; len(vals) > 0 {
		label = vals[0]
	}

	agent := NewAgent(bidi.Context(), c.jobber, aid, label)
	agent.bidi = bidi

	c.tracker.Register(aid, agent)
	defer c.tracker.Unregister(aid)

	return agent.run()
}

// GetAgent returns the agent by ID. Delegates to the tracker.
func (c *Commander) GetAgent(aid string) (*Agent, bool) {
	return c.tracker.GetAgent(aid)
}

func (a *Agent) AddJob(job models.JobForInsert) (int64, error) {
	jid, err := a.jobber.InitJob(a.ctx, a.aid, job)
	if err != nil {
		return 0, err
	}
	a.jobs[jid] = newJob()
	a.in <- &proto.Command{
		Id:      jid,
		Command: job.Command,
		Stdin:   job.Stdin,
	}
	return jid, nil
}

func (a *Agent) WaitJob(jid int64) (*models.Job, error) {
	result := <-a.jobs[jid].out
	return &result.fc, result.err
}

func (a *Agent) run() error {
	wg := new(errgroup.Group)
	wg.Go(a.recv)
	wg.Go(a.send)
	return wg.Wait()
}

func (a *Agent) recv() error {
	for {
		job, err := func() (job models.Job, err error) {
			msg, err := a.bidi.Recv()
			if err != nil {
				return
			}
			return a.jobber.UpdateJobInDB(a.ctx, msg.Id, models.JobForUpdate{
				Stdout: msg.Stdout,
				Stderr: msg.Stderr,
				Status: msg.Status,
			})
		}()
		if err == io.EOF {
			return nil
		}
		out := a.jobs[job.ID].out
		out <- JobOut{
			fc:  job,
			err: err,
		}
		close(out)
	}
}

func (a *Agent) send() error {
	for job := range a.in {
		if err := a.bidi.Send(job); err != nil {
			return err
		}
	}
	return io.EOF
}

func NewAgent(
	ctx context.Context,
	jobber Jobber,
	aid string,
	label string,
) *Agent {
	return &Agent{
		in:     make(chan *proto.Command, 10),
		jobs:   make(map[int64]Job),
		jobber: jobber,
		ctx:    ctx,
		aid:    aid,
		Label:  label,
		Token:  aid,
	}
}

func newJob() Job {
	return Job{make(chan JobOut, 1)}
}