mrva/mrvacommander
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Separate queue and agent logic and refactor

Browse Source
This commit is contained in:
Nicolas Will
2024-06-15 23:12:11 +02:00
parent e107f6cf80
commit 7ea45cb176
5 changed files with 328 additions and 357 deletions
Download Patch File Download Diff File Expand all files Collapse all files

187
pkg/agent/agent.go
View File

@@ -1,20 +1,20 @@
package agent
import (
"fmt"
"log/slog"
"mrvacommander/pkg/codeql"
"mrvacommander/pkg/common"
"mrvacommander/pkg/logger"
"mrvacommander/pkg/qpstore"
"mrvacommander/pkg/queue"
"mrvacommander/pkg/storage"
"mrvacommander/utils"
"log/slog"
"fmt"
"path/filepath"
"os"
"os/exec"
"path/filepath"
"sync"
"github.com/google/uuid"
)
type RunnerSingle struct {
@@ -40,102 +40,101 @@ type Visibles struct {
}
func (c *RunnerSingle) Setup(st *Visibles) {
return
// TODO: implement
}
func (r *RunnerSingle) worker(wid int) {
var job common.AnalyzeJob
// TODO: reimplement this later
/*
var job common.AnalyzeJob
for {
job = <-r.queue.Jobs()
for {
job = <-r.queue.Jobs()
slog.Debug("Picking up job", "job", job, "worker", wid)
slog.Debug("Picking up job", "job", job, "worker", wid)
slog.Debug("Analysis: running", "job", job)
storage.SetStatus(job.QueryPackId, job.NWO, common.StatusQueued)
slog.Debug("Analysis: running", "job", job)
storage.SetStatus(job.QueryPackId, job.NWO, common.StatusQueued)
_, err := RunAnalysis(job)
resultFile, err := RunAnalysis(job)
if err != nil {
continue
}
slog.Debug("Analysis run finished", "job", job)
// TODO: FIX THIS
res := common.AnalyzeResult{
RunAnalysisSARIF: resultFile,
RunAnalysisBQRS: "", // FIXME ?
}
r.queue.Results() <- res
storage.SetStatus(job.QueryPackId, job.NWO, common.StatusSuccess)
storage.SetResult(job.QueryPackId, job.NWO, res)
}
*/
}
// RunAnalysisJob runs a CodeQL analysis job (AnalyzeJob) returning an AnalyzeResult
func RunAnalysisJob(job common.AnalyzeJob) (common.AnalyzeResult, error) {
var result = common.AnalyzeResult{
RequestId: job.RequestId,
ResultCount: 0,
ResultArchiveURL: "",
Status: common.StatusError,
}
// Create a temporary directory
tempDir := filepath.Join(os.TempDir(), uuid.New().String())
if err := os.MkdirAll(tempDir, 0755); err != nil {
return result, fmt.Errorf("failed to create temporary directory: %v", err)
}
defer os.RemoveAll(tempDir)
// Extract the query pack
// TODO: download from the 'job' query pack URL
// utils.downloadFile
queryPackPath := filepath.Join(tempDir, "qp-54674")
utils.UntarGz("qp-54674.tgz", queryPackPath)
// Perform the CodeQL analysis
runResult, err := codeql.RunQuery("google_flatbuffers_db.zip", "cpp", queryPackPath, tempDir)
if err != nil {
return result, fmt.Errorf("failed to run analysis: %w", err)
}
// Generate a ZIP archive containing SARIF and BQRS files
resultsArchive, err := codeql.GenerateResultsZipArchive(runResult)
if err != nil {
return result, fmt.Errorf("failed to generate results archive: %w", err)
}
// TODO: Upload the archive to storage
slog.Info("Results archive size", slog.Int("size", len(resultsArchive)))
slog.Info("Analysis job successful.")
result = common.AnalyzeResult{
RequestId: job.RequestId,
ResultCount: runResult.ResultCount,
ResultArchiveURL: "REPLACE_THIS_WITH_STORED_RESULTS_ARCHIVE", // TODO
Status: common.StatusSuccess,
}
return result, nil
}
// RunWorker runs a worker that processes jobs from queue
func RunWorker(queue queue.Queue, wg *sync.WaitGroup) {
defer wg.Done()
for job := range queue.Jobs() {
slog.Info("Running analysis job", slog.Any("job", job))
result, err := RunAnalysisJob(job)
if err != nil {
slog.Error("Failed to run analysis job", slog.Any("error", err))
continue
}
slog.Debug("Analysis run finished", "job", job)
res := common.AnalyzeResult{}
r.queue.Results() <- res
storage.SetStatus(job.QueryPackId, job.NWO, common.StatusSuccess)
storage.SetResult(job.QueryPackId, job.NWO, res)
slog.Info("Analysis job completed", slog.Any("result", result))
queue.Results() <- result
}
}
func RunAnalysis(job common.AnalyzeJob) (string, error) {
// TODO Add multi-language tests including queryLanguage
// queryPackID, queryLanguage, dbOwner, dbRepo :=
// job.QueryPackId, job.QueryLanguage, job.NWO.Owner, job.NWO.Repo
queryPackID, dbOwner, dbRepo :=
job.QueryPackId, job.NWO.Owner, job.NWO.Repo
serverRoot := os.Getenv("MRVA_SERVER_ROOT")
// Set up derived paths
dbPath := filepath.Join(serverRoot, "var/codeql/dbs", dbOwner, dbRepo)
dbZip := filepath.Join(serverRoot, "codeql/dbs", dbOwner, dbRepo,
fmt.Sprintf("%s_%s_db.zip", dbOwner, dbRepo))
dbExtract := filepath.Join(serverRoot, "var/codeql/dbs", dbOwner, dbRepo)
queryPack := filepath.Join(serverRoot,
"var/codeql/querypacks", fmt.Sprintf("qp-%d.tgz", queryPackID))
queryExtract := filepath.Join(serverRoot,
"var/codeql/querypacks", fmt.Sprintf("qp-%d", queryPackID))
queryOutDir := filepath.Join(serverRoot,
"var/codeql/sarif/localrun", dbOwner, dbRepo)
queryOutFile := filepath.Join(queryOutDir,
fmt.Sprintf("%s_%s.sarif", dbOwner, dbRepo))
// Prepare directory, extract database
if err := os.MkdirAll(dbExtract, 0755); err != nil {
slog.Error("Failed to create DB directory %s: %v", dbExtract, err)
return "", err
}
if err := utils.UnzipFile(dbZip, dbExtract); err != nil {
slog.Error("Failed to unzip DB", dbZip, err)
return "", err
}
// Prepare directory, extract query pack
if err := os.MkdirAll(queryExtract, 0755); err != nil {
slog.Error("Failed to create query pack directory %s: %v", queryExtract, err)
return "", err
}
if err := utils.UntarGz(queryPack, queryExtract); err != nil {
slog.Error("Failed to extract querypack %s: %v", queryPack, err)
return "", err
}
// Prepare query result directory
if err := os.MkdirAll(queryOutDir, 0755); err != nil {
slog.Error("Failed to create query result directory %s: %v", queryOutDir, err)
return "", err
}
// Run database analyze
cmd := exec.Command("codeql", "database", "analyze",
"--format=sarif-latest", "--rerun", "--output", queryOutFile,
"-j8", dbPath, queryExtract)
cmd.Dir = serverRoot
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
if err := cmd.Run(); err != nil {
slog.Error("codeql database analyze failed:", "error", err, "job", job)
storage.SetStatus(job.QueryPackId, job.NWO, common.StatusError)
return "", err
}
// Return result path
return queryOutFile, nil
}

2
pkg/codeql/codeql.go
View File

@@ -16,7 +16,6 @@ import (
"gopkg.in/yaml.v3"
)
// Helper Functions
func contains(slice []string, item string) bool {
for _, s := range slice {
if s == item {
@@ -26,7 +25,6 @@ func contains(slice []string, item string) bool {
return false
}
// Main Functions
func getCodeQLCLIPath() (string, error) {
// get the CODEQL_CLI_PATH environment variable
codeqlCliPath := os.Getenv("CODEQL_CLI_PATH")

163
pkg/queue/rabbitmq.go Normal file
View File

@@ -0,0 +1,163 @@
package queue
import (
"mrvacommander/pkg/common"
"mrvacommander/pkg/storage"
"context"
"encoding/json"
"fmt"
"log"
"time"
amqp "github.com/rabbitmq/amqp091-go"
"golang.org/x/exp/slog"
)
type RabbitMQQueue struct {
jobs chan common.AnalyzeJob
results chan common.AnalyzeResult
conn *amqp.Connection
channel *amqp.Channel
}
func InitializeRabbitMQQueue(
host string,
port int16,
user string,
password string,
) (*RabbitMQQueue, error) {
const (
tryCount = 5
retryDelaySec = 3
jobsQueueName = "tasks"
resultsQueueName = "results"
)
var conn *amqp.Connection
var err error
rabbitMQURL := fmt.Sprintf("amqp://%s:%s@%s:%d/", user, password, host, port)
for i := 0; i < tryCount; i++ {
slog.Info("Attempting to connect to RabbitMQ", slog.Int("attempt", i+1))
conn, err = amqp.Dial(rabbitMQURL)
if err != nil {
slog.Warn("Failed to connect to RabbitMQ", "error", err)
if i < tryCount-1 {
slog.Info("Retrying", "seconds", retryDelaySec)
time.Sleep(retryDelaySec * time.Second)
}
} else {
// successfully connected to RabbitMQ
break
}
}
if err != nil {
return nil, fmt.Errorf("failed to connect: %w", err)
}
slog.Info("Connected to RabbitMQ")
ch, err := conn.Channel()
if err != nil {
conn.Close()
return nil, fmt.Errorf("failed to open a channel: %w", err)
}
_, err = ch.QueueDeclare(jobsQueueName, false, false, false, true, nil)
if err != nil {
conn.Close()
return nil, fmt.Errorf("failed to declare tasks queue: %w", err)
}
_, err = ch.QueueDeclare(resultsQueueName, false, false, false, true, nil)
if err != nil {
conn.Close()
return nil, fmt.Errorf("failed to declare results queue: %w", err)
}
err = ch.Qos(1, 0, false)
if err != nil {
conn.Close()
return nil, fmt.Errorf("failed to set QoS: %w", err)
}
result := RabbitMQQueue{
conn: conn,
channel: ch,
jobs: make(chan common.AnalyzeJob),
results: make(chan common.AnalyzeResult),
}
slog.Info("Starting tasks consumer")
go result.ConsumeJobs(jobsQueueName)
slog.Info("Starting results publisher")
go result.PublishResults(resultsQueueName)
return &result, nil
}
func (q *RabbitMQQueue) Jobs() chan common.AnalyzeJob {
return q.jobs
}
func (q *RabbitMQQueue) Results() chan common.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()
}
func (q *RabbitMQQueue) ConsumeJobs(queueName string) {
msgs, err := q.channel.Consume(queueName, "", true, false, false, false, nil)
if err != nil {
slog.Error("failed to register a consumer", slog.Any("error", err))
}
for msg := range msgs {
job := common.AnalyzeJob{}
err := json.Unmarshal(msg.Body, &job)
if err != nil {
slog.Error("failed to unmarshal job", slog.Any("error", err))
continue
}
q.jobs <- job
}
close(q.jobs)
}
func (q *RabbitMQQueue) PublishResults(queueName string) {
for result := range q.results {
q.publishResult(queueName, result)
}
}
func (q *RabbitMQQueue) publishResult(queueName string, result interface{}) {
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
resultBytes, err := json.Marshal(result)
if err != nil {
slog.Error("failed to marshal result", slog.Any("error", err))
return
}
slog.Info("Publishing result", slog.String("result", string(resultBytes)))
err = q.channel.PublishWithContext(ctx, "", queueName, false, false,
amqp.Publishing{
ContentType: "application/json",
Body: resultBytes,
})
if err != nil {
slog.Error("failed to publish result", slog.Any("error", err))
}
}