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 {