Add Config struct with DisableConfigEndpoint option to profutils profiling server. Register config section under "prof" key with pflags generation and conditionally skip the /config HTTP handler when disabled. (#7016)

* Add Config struct with DisableConfigEndpoint option to profutils profiling server. Register config section under "prof" key with pflags generation and conditionally skip the /config HTTP handler when disabled.

Signed-off-by: Haytham Abuelfutuh <haytham@afutuh.com>

* Fix tests

Signed-off-by: Haytham Abuelfutuh <haytham@afutuh.com>

* Fix config registration

Signed-off-by: Haytham Abuelfutuh <haytham@afutuh.com>

* Fix unit tests

Signed-off-by: Haytham Abuelfutuh <haytham@afutuh.com>

* Fix test

Signed-off-by: Haytham Abuelfutuh <haytham@afutuh.com>

---------

Signed-off-by: Haytham Abuelfutuh <haytham@afutuh.com>

Author: EngHabu

flyteadmin/auth/cookie_manager_test.go

@@ -315,6 +315,7 @@ func TestExtractAccessTokenFromCookies(t *testing.T) {
 		t.Logf("Cookie 1 raw decoded (first 50 bytes): %s", string(decoded[:50]))
 	}
 	s := securecookie.New(hashKey, blockKey)
+	s.MaxAge(0) // Disable timestamp validation so hardcoded cookies don't expire
 	var firstHalf string
 	err = s.Decode("flyte_at_1", cookieValue1, &firstHalf)
 	if err != nil {
@@ -323,6 +324,7 @@ func TestExtractAccessTokenFromCookies(t *testing.T) {
 		err2 := s2.Decode("flyte_at_1", cookieValue1, &firstHalf)
 		t.Logf("Error with validation enabled: %v", err2)
 	}
+
 	require.NoError(t, err)
 	t.Logf("First half of access token: %s", firstHalf)
 

flytestdlib/profutils/config.go

@@ -0,0 +1,20 @@
+package profutils
+
+import "github.com/flyteorg/flyte/flytestdlib/config"
+
+//go:generate pflags Config --default-var=defaultConfig
+
+const configSectionKey = "prof"
+
+var (
+	configSection = config.MustRegisterSection(configSectionKey, defaultConfig)
+	defaultConfig = &Config{}
+)
+
+type Config struct {
+	DisableConfigEndpoint bool `config:"DisableConfigEndpoint"`
+}
+
+func GetConfig() *Config {
+	return configSection.GetConfig().(*Config)
+}

flytestdlib/profutils/config_flags.go

@@ -110,14 +110,17 @@ func StartProfilingServer(ctx context.Context, pprofPort int) error {
 	return srv.ListenAndServe()
 }
 
-func configureGlobalHTTPHandler(handlers map[string]http.Handler) error {
+func configureGlobalHTTPHandler(cfg *Config, handlers map[string]http.Handler) error {
 	if handlers == nil {
 		handlers = map[string]http.Handler{}
 	}
+
 	handlers[metricsPath] = promhttp.Handler()
 	handlers[healthcheck] = http.HandlerFunc(healtcheckHandler)
 	handlers[versionPath] = http.HandlerFunc(versionHandler)
-	handlers[configPath] = http.HandlerFunc(configHandler)
+	if !cfg.DisableConfigEndpoint {
+		handlers[configPath] = http.HandlerFunc(configHandler)
+	}
 
 	for p, h := range handlers {
 		http.Handle(p, h)
@@ -126,14 +129,15 @@ func configureGlobalHTTPHandler(handlers map[string]http.Handler) error {
 	return nil
 }
 
-// Forwards the call to StartProfilingServer
+// StartProfilingServerWithDefaultHandlers forwards the call to StartProfilingServer
 // Also registers:
 // 1. the prometheus HTTP handler on '/metrics' path shared with the profiling server.
 // 2. A healthcheck (L7) handler on '/healthcheck'.
 // 3. A version handler on '/version' provides information about the specific build.
 // 4. A config handler on '/config' provides a dump of the currently loaded config.
 func StartProfilingServerWithDefaultHandlers(ctx context.Context, pprofPort int, handlers map[string]http.Handler) error {
-	if err := configureGlobalHTTPHandler(handlers); err != nil {
+	cfg := GetConfig()
+	if err := configureGlobalHTTPHandler(cfg, handlers); err != nil {
 		return err
 	}
 

flytestdlib/profutils/config_flags_test.go

@@ -36,7 +36,7 @@ type TestObj struct {
 }
 
 func init() {
-	if err := configureGlobalHTTPHandler(nil); err != nil {
+	if err := configureGlobalHTTPHandler(&Config{}, nil); err != nil {
 		panic(err)
 	}
 }
@@ -77,6 +77,9 @@ func TestConfigHandler(t *testing.T) {
 				"type": "json",
 			},
 		},
+		"prof": map[string]interface{}{
+			"DisableConfigEndpoint": false,
+		},
 	}, m)
 }