Add prober flags for passing extra requests. (#395)

Signed-off-by: Simon Kent <simek@google.com> Signed-off-by: Simon Kent <simek@google.com>
sigstore · Oct 3, 2022 · eceb1be · eceb1be
1 parent 8e0759e
commit eceb1be
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 31 deletions.
diff --git a/cmd/prober/endpoints.go b/cmd/prober/endpoints.go
@@ -20,47 +20,47 @@ var (
 )
 
 type ReadProberCheck struct {
-	endpoint string
-	method   string
-	body     string
-	queries  map[string]string
+	Endpoint string            `json:"endpoint"`
+	Method   string            `json:"method"`
+	Body     string            `json:"body"`
+	Queries  map[string]string `json:"queries"`
 }
 
 var RekorEndpoints = []ReadProberCheck{
 	{
-		endpoint: "/api/v1/log/publicKey",
-		method:   GET,
+		Endpoint: "/api/v1/log/publicKey",
+		Method:   GET,
 	}, {
-		endpoint: "/api/v1/log",
-		method:   GET,
+		Endpoint: "/api/v1/log",
+		Method:   GET,
 	}, {
-		endpoint: "/api/v1/log/entries",
-		method:   GET,
-		queries:  map[string]string{"logIndex": "10"},
+		Endpoint: "/api/v1/log/entries",
+		Method:   GET,
+		Queries:  map[string]string{"logIndex": "10"},
 	}, {
-		endpoint: "/api/v1/log/proof",
-		method:   GET,
-		queries:  map[string]string{"firstSize": "10", "lastSize": "20"},
+		Endpoint: "/api/v1/log/proof",
+		Method:   GET,
+		Queries:  map[string]string{"firstSize": "10", "lastSize": "20"},
 	}, {
-		endpoint: "/api/v1/log/entries/retrieve",
-		method:   POST,
-		body:     "{\"hash\":\"sha256:2bd37672a9e472c79c64f42b95e362db16870e28a90f3b17fee8faf952e79b4b\"}",
+		Endpoint: "/api/v1/log/entries/retrieve",
+		Method:   POST,
+		Body:     "{\"hash\":\"sha256:2bd37672a9e472c79c64f42b95e362db16870e28a90f3b17fee8faf952e79b4b\"}",
 	}, {
-		endpoint: "/api/v1/index/retrieve",
-		method:   POST,
-		body:     "{\"hash\":\"sha256:2bd37672a9e472c79c64f42b95e362db16870e28a90f3b17fee8faf952e79b4b\"}",
+		Endpoint: "/api/v1/index/retrieve",
+		Method:   POST,
+		Body:     "{\"hash\":\"sha256:2bd37672a9e472c79c64f42b95e362db16870e28a90f3b17fee8faf952e79b4b\"}",
 	},
 }
 
 var FulcioEndpoints = []ReadProberCheck{
 	{
-		endpoint: "/api/v1/rootCert",
-		method:   GET,
+		Endpoint: "/api/v1/rootCert",
+		Method:   GET,
 	}, {
-		endpoint: "/api/v2/configuration",
-		method:   GET,
+		Endpoint: "/api/v2/configuration",
+		Method:   GET,
 	}, {
-		endpoint: "/api/v2/trustBundle",
-		method:   GET,
+		Endpoint: "/api/v2/trustBundle",
+		Method:   GET,
 	},
 }
diff --git a/cmd/prober/prober.go b/cmd/prober/prober.go
@@ -17,6 +17,7 @@ package main
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"flag"
 	"fmt"
 	"log"
@@ -51,7 +52,26 @@ func init() {
 	flag.BoolVar(&oneTime, "one-time", false, "Whether to run only one time and exit.")
 	flag.BoolVar(&runWriteProber, "write-prober", false, " [Kubernetes only] run the probers for the write endpoints.")
 
+	var rekorRequestsJson string
+	flag.StringVar(&rekorRequestsJson, "rekor-requests", "[]", "Additional rekor requests (JSON array.)")
+
+	var fulcioRequestsJson string
+	flag.StringVar(&fulcioRequestsJson, "fulcio-requests", "[]", "Additional fulcio requests (JSON array.)")
+
 	flag.Parse()
+
+	var rekorFlagRequests []ReadProberCheck
+	if err := json.Unmarshal([]byte(rekorRequestsJson), &rekorFlagRequests); err != nil {
+		log.Fatal("Failed to parse rekor-requests: ", err)
+	}
+
+	var fulcioFlagRequests []ReadProberCheck
+	if err := json.Unmarshal([]byte(fulcioRequestsJson), &fulcioFlagRequests); err != nil {
+		log.Fatal("Failed to parse fulcio-requests: ", err)
+	}
+
+	RekorEndpoints = append(RekorEndpoints, rekorFlagRequests...)
+	FulcioEndpoints = append(FulcioEndpoints, fulcioFlagRequests...)
 }
 
 func main() {
@@ -84,13 +104,13 @@ func runProbers(ctx context.Context, freq int, runOnce bool) {
 		for _, r := range RekorEndpoints {
 			if err := observeRequest(rekorURL, r); err != nil {
 				hasErr = true
-				fmt.Printf("error running request %s: %v\n", r.endpoint, err)
+				fmt.Printf("error running request %s: %v\n", r.Endpoint, err)
 			}
 		}
 		for _, r := range FulcioEndpoints {
 			if err := observeRequest(fulcioURL, r); err != nil {
 				hasErr = true
-				fmt.Printf("error running request %s: %v\n", r.endpoint, err)
+				fmt.Printf("error running request %s: %v\n", r.Endpoint, err)
 			}
 		}
 		if runWriteProber {
@@ -135,19 +155,19 @@ func observeRequest(host string, r ReadProberCheck) error {
 	}
 	defer resp.Body.Close()
 
-	exportDataToPrometheus(resp, host, r.endpoint, r.method, latency)
+	exportDataToPrometheus(resp, host, r.Endpoint, r.Method, latency)
 	return nil
 }
 
 func httpRequest(host string, r ReadProberCheck) (*http.Request, error) {
-	req, err := http.NewRequest(r.method, host+r.endpoint, bytes.NewBuffer([]byte(r.body)))
+	req, err := http.NewRequest(r.Method, host+r.Endpoint, bytes.NewBuffer([]byte(r.Body)))
 	if err != nil {
 		return nil, err
 	}
 
 	setHeaders(req, "")
 	q := req.URL.Query()
-	for k, v := range r.queries {
+	for k, v := range r.Queries {
 		q.Add(k, v)
 	}
 	req.URL.RawQuery = q.Encode()