From 3fde37726e0c80e32cb33b23b703a4c2bcc6e74c Mon Sep 17 00:00:00 2001
From: Daniel Hrabovcak <thespiritxiii@gmail.com>
Date: Thu, 14 Sep 2023 09:48:55 -0400
Subject: [PATCH] Add Authorization header support into the example app

---
 examples/instrumentation/go-synthetic/auth.go | 73 +++++++++++++++++++
 examples/instrumentation/go-synthetic/main.go |  9 ++-
 2 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/examples/instrumentation/go-synthetic/auth.go b/examples/instrumentation/go-synthetic/auth.go
index 97c8c7381a..f60fbae3d2 100644
--- a/examples/instrumentation/go-synthetic/auth.go
+++ b/examples/instrumentation/go-synthetic/auth.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"errors"
 	"flag"
 	"net/http"
 )
@@ -37,3 +38,75 @@ func (c *basicAuthConfig) handle(handler http.Handler) http.Handler {
 		http.Error(w, "Unauthorized", http.StatusUnauthorized)
 	})
 }
+
+type authorizationConfig struct {
+	scheme     string
+	parameters string
+}
+
+func newAuthorizationConfigFromFlags() *authorizationConfig {
+	c := &authorizationConfig{}
+	flag.StringVar(&c.scheme, "auth-scheme", "", "Authorization header scheme")
+	flag.StringVar(&c.parameters, "auth-parameters", "", "Data to require in the Authorization header")
+	return c
+}
+
+func (c *authorizationConfig) isEnabled() bool {
+	return c.scheme != "" || c.parameters != ""
+}
+
+func (c *authorizationConfig) validate() error {
+	var errs []error
+	if c.scheme == "" && c.parameters != "" {
+		errs = append(errs, errors.New("must specify --auth-scheme when using --auth-parameters"))
+	}
+	if c.scheme == "Basic" {
+		errs = append(errs, errors.New("use --basic-auth flags to specify BasicAuth"))
+	}
+	return errors.Join(errs...)
+}
+
+func (c *authorizationConfig) handle(handler http.Handler) http.Handler {
+	if !c.isEnabled() {
+		return handler
+	}
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		auth := r.Header.Get("Authorization")
+		expected := c.scheme + " " + c.parameters
+		if auth == expected {
+			handler.ServeHTTP(w, r)
+			return
+		}
+
+		w.Header().Set("WWW-Authenticate", c.scheme+` realm="restricted", charset="UTF-8"`)
+		http.Error(w, "Unauthorized", http.StatusUnauthorized)
+	})
+}
+
+type httpClientConfig struct {
+	basicAuth *basicAuthConfig
+	auth      *authorizationConfig
+}
+
+func newHttpClientConfigFromFlags() *httpClientConfig {
+	return &httpClientConfig{
+		basicAuth: newBasicAuthConfigFromFlags(),
+		auth:      newAuthorizationConfigFromFlags(),
+	}
+}
+
+func (c *httpClientConfig) validate() error {
+	var errs []error
+	if c.basicAuth.isEnabled() && c.auth.isEnabled() {
+		errs = append(errs, errors.New("cannot specify both --basic-auth and --auth flags"))
+	}
+	if err := c.auth.validate(); err != nil {
+		errs = append(errs, err)
+	}
+	return errors.Join(errs...)
+}
+
+func (c *httpClientConfig) handle(handler http.Handler) http.Handler {
+	return c.auth.handle(c.basicAuth.handle(handler))
+}
diff --git a/examples/instrumentation/go-synthetic/main.go b/examples/instrumentation/go-synthetic/main.go
index e59dfc5808..47c3a2101b 100644
--- a/examples/instrumentation/go-synthetic/main.go
+++ b/examples/instrumentation/go-synthetic/main.go
@@ -172,9 +172,14 @@ var (
 )
 
 func main() {
-	basicAuthConfig := newBasicAuthConfigFromFlags()
+	httpClientConfig := newHttpClientConfigFromFlags()
 	flag.Parse()
 
+	if err := httpClientConfig.validate(); err != nil {
+		log.Println("Invalid HTTP client config flags:", err)
+		os.Exit(1)
+	}
+
 	metrics := prometheus.NewRegistry()
 	metrics.MustRegister(
 		collectors.NewGoCollector(collectors.WithGoCollectorRuntimeMetrics(collectors.MetricsAll)),
@@ -223,7 +228,7 @@ func main() {
 
 		server := &http.Server{
 			Addr:    *addr,
-			Handler: basicAuthConfig.handle(mux),
+			Handler: httpClientConfig.handle(mux),
 		}
 
 		g.Add(func() error {