Fully implement local and container MRVA

This commit is contained in:
Nicolas Will
2024-06-17 13:16:24 +02:00
parent ef7552c43f
commit e0cbc01d21
43 changed files with 1700 additions and 1137 deletions

View File

@@ -1,14 +1,7 @@
package queue
import (
"mrvacommander/pkg/common"
"mrvacommander/pkg/storage"
)
type Queue interface {
Jobs() chan common.AnalyzeJob
Results() chan common.AnalyzeResult
StartAnalyses(analysis_repos *map[common.NameWithOwner]storage.DBLocation,
session_id int,
session_language string)
Jobs() chan AnalyzeJob
Results() chan AnalyzeResult
Close()
}

View File

@@ -1,31 +0,0 @@
package queue
import (
"log/slog"
"mrvacommander/pkg/common"
"mrvacommander/pkg/storage"
)
func (q *QueueSingle) Jobs() chan common.AnalyzeJob {
return q.jobs
}
func (q *QueueSingle) Results() chan common.AnalyzeResult {
return q.results
}
func (q *QueueSingle) StartAnalyses(analysis_repos *map[common.NameWithOwner]storage.DBLocation, session_id int,
session_language string) {
slog.Debug("Queueing codeql database analyze jobs")
for nwo := range *analysis_repos {
info := common.AnalyzeJob{
QueryPackId: session_id,
QueryLanguage: session_language,
NWO: nwo,
}
q.jobs <- info
storage.SetStatus(session_id, nwo, common.StatusQueued)
storage.AddJob(session_id, info)
}
}

View File

@@ -1,13 +1,9 @@
package queue
import (
"mrvacommander/pkg/common"
"mrvacommander/pkg/storage"
"context"
"encoding/json"
"fmt"
"log"
"time"
amqp "github.com/rabbitmq/amqp091-go"
@@ -15,20 +11,20 @@ import (
)
type RabbitMQQueue struct {
jobs chan common.AnalyzeJob
results chan common.AnalyzeResult
jobs chan AnalyzeJob
results chan AnalyzeResult
conn *amqp.Connection
channel *amqp.Channel
}
// InitializeRabbitMQQueue initializes a RabbitMQ queue.
// NewRabbitMQQueue initializes a RabbitMQ queue.
// It returns a pointer to a RabbitMQQueue and an error.
//
// If isAgent is true, the queue is initialized to be used by an agent.
// Otherwise, the queue is initialized to be used by the server.
// The difference in behaviour is that the agent consumes jobs and publishes results,
// while the server publishes jobs and consumes results.
func InitializeRabbitMQQueue(
func NewRabbitMQQueue(
host string,
port int16,
user string,
@@ -94,8 +90,8 @@ func InitializeRabbitMQQueue(
result := RabbitMQQueue{
conn: conn,
channel: ch,
jobs: make(chan common.AnalyzeJob),
results: make(chan common.AnalyzeResult),
jobs: make(chan AnalyzeJob),
results: make(chan AnalyzeResult),
}
if isAgent {
@@ -115,19 +111,14 @@ func InitializeRabbitMQQueue(
return &result, nil
}
func (q *RabbitMQQueue) Jobs() chan common.AnalyzeJob {
func (q *RabbitMQQueue) Jobs() chan AnalyzeJob {
return q.jobs
}
func (q *RabbitMQQueue) Results() chan common.AnalyzeResult {
func (q *RabbitMQQueue) Results() chan AnalyzeResult {
return q.results
}
func (q *RabbitMQQueue) StartAnalyses(analysis_repos *map[common.NameWithOwner]storage.DBLocation, session_id int, session_language string) {
// TODO: Implement
log.Fatal("unimplemented")
}
func (q *RabbitMQQueue) Close() {
q.channel.Close()
q.conn.Close()
@@ -140,7 +131,7 @@ func (q *RabbitMQQueue) ConsumeJobs(queueName string) {
}
for msg := range msgs {
job := common.AnalyzeJob{}
job := AnalyzeJob{}
err := json.Unmarshal(msg.Body, &job)
if err != nil {
slog.Error("failed to unmarshal job", slog.Any("error", err))
@@ -157,7 +148,7 @@ func (q *RabbitMQQueue) PublishResults(queueName string) {
}
}
func (q *RabbitMQQueue) publishResult(queueName string, result common.AnalyzeResult) {
func (q *RabbitMQQueue) publishResult(queueName string, result AnalyzeResult) {
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
@@ -178,7 +169,7 @@ func (q *RabbitMQQueue) publishResult(queueName string, result common.AnalyzeRes
}
}
func (q *RabbitMQQueue) publishJob(queueName string, job common.AnalyzeJob) {
func (q *RabbitMQQueue) publishJob(queueName string, job AnalyzeJob) {
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
@@ -212,7 +203,7 @@ func (q *RabbitMQQueue) ConsumeResults(queueName string) {
}
for msg := range msgs {
result := common.AnalyzeResult{}
result := AnalyzeResult{}
err := json.Unmarshal(msg.Body, &result)
if err != nil {
slog.Error("failed to unmarshal result", slog.Any("error", err))

29
pkg/queue/queue_single.go Normal file
View File

@@ -0,0 +1,29 @@
package queue
type QueueSingle struct {
NumWorkers int
jobs chan AnalyzeJob
results chan AnalyzeResult
}
func NewQueueSingle(numWorkers int) Queue {
q := QueueSingle{
NumWorkers: numWorkers,
jobs: make(chan AnalyzeJob, 10),
results: make(chan AnalyzeResult, 10),
}
return q
}
func (q QueueSingle) Jobs() chan AnalyzeJob {
return q.jobs
}
func (q QueueSingle) Results() chan AnalyzeResult {
return q.results
}
func (q QueueSingle) Close() {
close(q.jobs)
close(q.results)
}

View File

@@ -1,28 +1,28 @@
package queue
import (
"mrvacommander/pkg/artifactstore"
"mrvacommander/pkg/common"
"mrvacommander/pkg/logger"
)
type QueueSingle struct {
NumWorkers int
jobs chan common.AnalyzeJob
results chan common.AnalyzeResult
modules *Visibles
// AnalyzeJob represents a job specifying a repository and a query pack to analyze it with.
// This is the message format that the agent receives from the queue.
// TODO: make query_pack_location query_pack_url with a presigned URL
type AnalyzeJob struct {
Spec common.JobSpec // json:"job_spec"
QueryPackLocation artifactstore.ArtifactLocation // json:"query_pack_location"
QueryLanguage string // json:"query_language"
}
type Visibles struct {
Logger logger.Logger
}
func NewQueueSingle(numWorkers int, v *Visibles) *QueueSingle {
q := QueueSingle{}
q.jobs = make(chan common.AnalyzeJob, 10)
q.results = make(chan common.AnalyzeResult, 10)
q.NumWorkers = numWorkers
q.modules = v
return &q
// AnalyzeResult represents the result of an analysis job.
// This is the message format that the agent sends to the queue.
// Status will only ever be StatusSuccess or StatusError when sent in a result.
// TODO: make result_location result_archive_url with a presigned URL
type AnalyzeResult struct {
Spec common.JobSpec // json:"job_spec"
Status common.Status // json:"status"
ResultCount int // json:"result_count"
ResultLocation artifactstore.ArtifactLocation // json:"result_location"
SourceLocationPrefix string // json:"source_location_prefix"
DatabaseSHA string // json:"database_sha"
}