Allow more headers to be added from console

.gitignore

@@ -1,2 +1,3 @@
 /.idea/
 /build/
+.vscode

README.md

@@ -27,6 +27,8 @@ Usage of ./build/cspr-collector:
         enable http output
   -output-http-host string
         http host to send the csp violations to (default "http://localhost:80/")
+  -output-http-headers string
+        header header key with value to send with the http request. Example "Authorization: ApiKey <APIKEY>"
   -output-stdout
         enable stdout output
 ```
@@ -98,6 +100,12 @@ curl -X POST \
 go fmt ./...
 ```
 
+## Health check endpoint
+endpoint with health response "OK"
+```bash
+curl http://localhost:8080/health
+```
+
 ## License
 
 The MIT License (MIT). Please see the [license file](LICENSE.md) for more information.

cmd/cspr-collector/main.go

@@ -3,11 +3,12 @@ package main
 import (
 	"context"
 	"flag"
-	cspr "github.com/mhilker/cspr-collector"
 	"log"
 	"net/http"
 	"os"
 	"os/signal"
+
+	cspr "github.com/mhilker/cspr-collector"
 )
 
 var (
@@ -16,6 +17,7 @@ var (
 	OutputStdout      = flag.Bool("output-stdout", false, "enable stdout output")
 	OutputHTTPEnabled = flag.Bool("output-http", false, "enable http output")
 	OutputHTTPHost    = flag.String("output-http-host", "http://localhost:80/", "http host to send the csp violations to")
+	OutputHTTPHeaders = flag.String("output-http-headers", "", "additional headers for HTTP output")
 	OutputESEnabled   = flag.Bool("output-es", false, "enable elasticsearch output")
 	OutputESHost      = flag.String("output-es-host", "http://localhost:9200/", "elasticsearch host to send the csp violations to")
 	OutputESIndex     = flag.String("output-es-index", "cspr-violations", "elasticsearch index to save the csp violations in")
@@ -67,7 +69,12 @@ func NewOutput() *cspr.CombinedOutput {
 
 	if *OutputHTTPEnabled {
 		log.Printf("Enable HTTP Output.")
-		outputs = append(outputs, &cspr.HTTPOutput{Url: *OutputHTTPHost})
+		headers := cspr.ParseHeaders(*OutputHTTPHeaders)
+
+		outputs = append(outputs, &cspr.HTTPOutput{
+			Url:     *OutputHTTPHost,
+			Headers: headers,
+		})
 	}
 
 	if *OutputESEnabled {

collector.go

@@ -19,6 +19,11 @@ type Collector struct {
 }
 
 func (c *Collector) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if r.RequestURI == "/health" {
+		c.handleHealth(w)
+		return
+	}
+
 	if r.RequestURI != "/" {
 		message := fmt.Sprintf("Path \"%s\" not found.", r.RequestURI)
 		c.response(w, http.StatusNotFound, message)
@@ -58,7 +63,6 @@ func (c *Collector) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	c.WorkQueue <- data
 
 	c.response(w, http.StatusCreated, "")
-	return
 }
 
 func (c *Collector) response(w http.ResponseWriter, status int, message string) {
@@ -82,3 +86,8 @@ func (c *Collector) response(w http.ResponseWriter, status int, message string)
 		log.Fatal(err)
 	}
 }
+
+func (c *Collector) handleHealth(w http.ResponseWriter) {
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte("OK"))
+}

dispatcher.go

@@ -31,15 +31,20 @@ func (d *Dispatcher) Run() {
 }
 
 func (d *Dispatcher) start() {
-	for {
-		select {
-		case work := <-d.WorkQueue:
-			log.Print("Received work request.")
-			go func() {
-				worker := <-d.WorkerQueue
-				log.Print("Dispatching work request.")
-				worker <- work
-			}()
-		}
+	for work := range d.WorkQueue {
+		log.Print("Received work request.")
+		go func(w interface{}) {
+			worker := <-d.WorkerQueue
+			log.Print("Dispatching work request.")
+
+			// Perform a type assertion
+			typedWork, ok := w.(CSPRequest)
+			if !ok {
+				log.Print("Invalid work request type")
+				return
+			}
+
+			worker <- typedWork
+		}(work)
 	}
 }

output_http.go

@@ -5,10 +5,12 @@ import (
 	"encoding/json"
 	"log"
 	"net/http"
+	"strings"
 )
 
 type HTTPOutput struct {
-	Url string
+	Url     string
+	Headers map[string]string
 }
 
 func (o *HTTPOutput) Write(data []CSPRequest) {
@@ -24,6 +26,10 @@ func (o *HTTPOutput) Write(data []CSPRequest) {
 			log.Print(err.Error())
 			return
 		}
+
+		for key, value := range o.Headers {
+			request.Header.Set(key, value)
+		}
 		request.Header.Set("Content-Type", "application/json")
 
 		client := &http.Client{}
@@ -37,3 +43,19 @@ func (o *HTTPOutput) Write(data []CSPRequest) {
 		log.Print("Response Status:", response.Status)
 	}
 }
+
+func ParseHeaders(headerString string) map[string]string {
+	headers := make(map[string]string)
+	pairs := strings.Split(headerString, ",")
+
+	for _, pair := range pairs {
+		split := strings.SplitN(pair, ":", 2)
+		if len(split) == 2 {
+			key := strings.TrimSpace(split[0])
+			value := strings.TrimSpace(split[1])
+			headers[key] = value
+		}
+	}
+
+	return headers
+}