From 12589f089389478a4aeb1cfb8d14365cf2e5e4c2 Mon Sep 17 00:00:00 2001
From: Tarun Koyalwar <tarun@projectdiscovery.io>
Date: Sat, 11 May 2024 00:00:33 +0530
Subject: [PATCH] add scan-id validation

---
 internal/pdcp/writer.go   | 13 +++++++++++--
 internal/runner/runner.go |  4 +++-
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/internal/pdcp/writer.go b/internal/pdcp/writer.go
index 6918d76100..a76d4ea31a 100644
--- a/internal/pdcp/writer.go
+++ b/internal/pdcp/writer.go
@@ -9,6 +9,7 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"regexp"
 	"sync/atomic"
 	"time"
 
@@ -27,9 +28,13 @@ const (
 	appendEndpoint = "/v1/scans/%s/import"
 	flushTimer     = time.Duration(1) * time.Minute
 	MaxChunkSize   = 1024 * 1024 * 4 // 4 MB
+	xidRe          = `^[a-z0-9]{20}$`
 )
 
-var _ output.Writer = &UploadWriter{}
+var (
+	xidRegex               = regexp.MustCompile(xidRe)
+	_        output.Writer = &UploadWriter{}
+)
 
 // UploadWriter is a writer that uploads its output to pdcp
 // server to enable web dashboard and more
@@ -87,8 +92,12 @@ func NewUploadWriter(ctx context.Context, creds *pdcpauth.PDCPCredentials) (*Upl
 }
 
 // SetScanID sets the scan id for the upload writer
-func (u *UploadWriter) SetScanID(id string) {
+func (u *UploadWriter) SetScanID(id string) error {
+	if !xidRegex.MatchString(id) {
+		return fmt.Errorf("invalid scan id provided")
+	}
 	u.scanID = id
+	return nil
 }
 
 // SetScanName sets the scan name for the upload writer
diff --git a/internal/runner/runner.go b/internal/runner/runner.go
index d0c2a9e250..d085c1bafb 100644
--- a/internal/runner/runner.go
+++ b/internal/runner/runner.go
@@ -414,7 +414,9 @@ func (r *Runner) setupPDCPUpload(writer output.Writer) output.Writer {
 		return writer
 	}
 	if r.options.ScanID != "" {
-		uploadWriter.SetScanID(r.options.ScanID)
+		if err := uploadWriter.SetScanID(r.options.ScanID); err != nil {
+			gologger.Fatal().Msgf("failed to set scan id: %s", err)
+		}
 	}
 	if r.options.ScanName != "" {
 		uploadWriter.SetScanName(r.options.ScanName)