fix: Add debug pprof endpoints on main router

* Make debug endpoint available only on localhost * Update example mock manager * Update swagger docs Signed-off-by: Mahendra Paipuri <[email protected]>
mahendrapaipuri · Jun 20, 2024 · 618160c · 618160c
1 parent 5059c38
commit 618160c
Show file tree

Hide file tree

Showing 7 changed files with 54 additions and 10 deletions.
diff --git a/examples/mock_resource_manager/cmd/mock_ceems_server/main.go b/examples/mock_resource_manager/cmd/mock_ceems_server/main.go
@@ -7,7 +7,12 @@ import (
 	"os"
 
 	_ "github.com/mahendrapaipuri/ceems/examples/mock_resource_manager/pkg/resource"
+
 	"github.com/mahendrapaipuri/ceems/pkg/api/cli"
+
+	// If existing schedulers in CEEMS are needed, they need to be imported too
+	// For instance to import slurm manager, following import statement must be added
+	_ "github.com/mahendrapaipuri/ceems/pkg/api/resource/slurm"
 )
 
 // Main entry point for `usagestats` app

diff --git a/examples/mock_resource_manager/pkg/resource/mock_manager.go b/examples/mock_resource_manager/pkg/resource/mock_manager.go
@@ -53,11 +53,12 @@ func NewMockManager(cluster models.Cluster, logger log.Logger) (resource.Fetcher
 	}, nil
 }
 
-// Add the logic here to get compute units from resource manager and return slice of Unit structs
+// Add the logic here to get compute units from resource manager and return slice of
+// ClusterUnits structs
 //
 // When making Unit stucts, ensure to format the datetime using base.DatetimeLayout
 // Also ensure to set StartTS and EndTS fields to start and end times in unix milliseconds epoch
-func (s *mockManager) Fetch(start time.Time, end time.Time) ([]models.ClusterUnits, error) {
+func (s *mockManager) FetchUnits(start time.Time, end time.Time) ([]models.ClusterUnits, error) {
 	return []models.ClusterUnits{
 		{
 			Cluster: models.Cluster{
@@ -74,3 +75,33 @@ func (s *mockManager) Fetch(start time.Time, end time.Time) ([]models.ClusterUni
 		},
 	}, nil
 }
+
+// Add the logic here to get users and projects/accounts/tenants/namespaces from
+// resource manager
+func (s *mockManager) FetchUsersProjects(current time.Time) ([]models.ClusterUsers, []models.ClusterProjects, error) {
+	return []models.ClusterUsers{
+			{
+				Cluster: models.Cluster{
+					ID: "mock",
+				},
+				Users: []models.User{
+					{
+						Name:     "usr1",
+						Projects: models.List{"prj1", "prj2"},
+					},
+				},
+			},
+		}, []models.ClusterProjects{
+			{
+				Cluster: models.Cluster{
+					ID: "mock",
+				},
+				Projects: []models.Project{
+					{
+						Name:  "usr1",
+						Users: models.List{"prj1", "prj2"},
+					},
+				},
+			},
+		}, nil
+}
diff --git a/pkg/api/http/docs/docs.go b/pkg/api/http/docs/docs.go
@@ -1418,7 +1418,7 @@ var SwaggerInfo = &swag.Spec{
 	BasePath:         "",
 	Schemes:          []string{},
 	Title:            "CEEMS API",
-	Description:      "OpenAPI specification (OAS) for the CEEMS REST API.\n\nSee the Interactive Docs to try CEEMS API methods without writing code, and get\nthe complete schema of resources exposed by the API.\n\nIf basic auth is enabled, all the endpoints require authentication.\n\nAll the endpoints, except `health` and `demo`, must send a user-agent header.\n\nTimestamps must be specified in milliseconds, unless otherwise specified.",
+	Description:      "OpenAPI specification (OAS) for the CEEMS REST API.\n\nSee the Interactive Docs to try CEEMS API methods without writing code, and get\nthe complete schema of resources exposed by the API.\n\nIf basic auth is enabled, all the endpoints require authentication.\n\nAll the endpoints, except `health`, `swagger`, `debug` and `demo`,\nmust send a user-agent header.\n\nTimestamps must be specified in milliseconds, unless otherwise specified.",
 	InfoInstanceName: "swagger",
 	SwaggerTemplate:  docTemplate,
 	LeftDelim:        "{{",

diff --git a/pkg/api/http/docs/swagger.json b/pkg/api/http/docs/swagger.json
@@ -1,7 +1,7 @@
 {
     "swagger": "2.0",
     "info": {
-        "description": "OpenAPI specification (OAS) for the CEEMS REST API.\n\nSee the Interactive Docs to try CEEMS API methods without writing code, and get\nthe complete schema of resources exposed by the API.\n\nIf basic auth is enabled, all the endpoints require authentication.\n\nAll the endpoints, except `health` and `demo`, must send a user-agent header.\n\nTimestamps must be specified in milliseconds, unless otherwise specified.",
+        "description": "OpenAPI specification (OAS) for the CEEMS REST API.\n\nSee the Interactive Docs to try CEEMS API methods without writing code, and get\nthe complete schema of resources exposed by the API.\n\nIf basic auth is enabled, all the endpoints require authentication.\n\nAll the endpoints, except `health`, `swagger`, `debug` and `demo`,\nmust send a user-agent header.\n\nTimestamps must be specified in milliseconds, unless otherwise specified.",
         "title": "CEEMS API",
         "contact": {
             "name": "Mahendra Paipuri",

diff --git a/pkg/api/http/docs/swagger.yaml b/pkg/api/http/docs/swagger.yaml
@@ -393,7 +393,8 @@ info:
 
     If basic auth is enabled, all the endpoints require authentication.
 
-    All the endpoints, except `health` and `demo`, must send a user-agent header.
+    All the endpoints, except `health`, `swagger`, `debug` and `demo`,
+    must send a user-agent header.
 
     Timestamps must be specified in milliseconds, unless otherwise specified.
   license:

diff --git a/pkg/api/http/middleware.go b/pkg/api/http/middleware.go
@@ -20,6 +20,11 @@ const (
 	ceemsUserHeader     = "X-Ceems-User" // Special header that will be included in requests from CEEMS LB
 )
 
+// Debug end point regex match
+var (
+	debugEndpoints = regexp.MustCompile("/debug/(.*)")
+)
+
 // Define our struct
 type authenticationMiddleware struct {
 	logger          log.Logger
@@ -47,7 +52,8 @@ func (amw *authenticationMiddleware) Middleware(next http.Handler) http.Handler
 		// unautorised user cannot access these end points
 		if r.URL.Path == "/" ||
 			r.URL.Path == amw.routerPrefix ||
-			amw.whitelistedURLs.MatchString(r.URL.Path) {
+			amw.whitelistedURLs.MatchString(r.URL.Path) ||
+			debugEndpoints.MatchString(r.URL.Path) {
 			goto end
 		}
 

diff --git a/pkg/api/http/server.go b/pkg/api/http/server.go
@@ -202,8 +202,8 @@ func NewCEEMSServer(c *Config) (*CEEMSServer, func(), error) {
 	// A demo end point that returns mocked data for units and/or usage tables
 	subRouter.HandleFunc("/demo/{resource:(?:units|usage)}", server.demo).Methods(http.MethodGet)
 
-	// pprof debug end points
-	subRouter.PathPrefix("/debug/").Handler(http.DefaultServeMux)
+	// pprof debug end points. Expose them only on localhost
+	router.PathPrefix("/debug/").Handler(http.DefaultServeMux).Host("localhost")
 
 	subRouter.PathPrefix("/swagger/").Handler(httpSwagger.Handler(
 		httpSwagger.URL("doc.json"), // The url pointing to API definition
@@ -223,7 +223,7 @@ func NewCEEMSServer(c *Config) (*CEEMSServer, func(), error) {
 	amw := authenticationMiddleware{
 		logger:          c.Logger,
 		routerPrefix:    routePrefix,
-		whitelistedURLs: regexp.MustCompile(fmt.Sprintf("%s(swagger|debug|health|demo)(.*)", routePrefix)),
+		whitelistedURLs: regexp.MustCompile(fmt.Sprintf("%s(swagger|health|demo)(.*)", routePrefix)),
 		db:              dbConn,
 		adminUsers:      adminUsers,
 	}
@@ -242,7 +242,8 @@ func NewCEEMSServer(c *Config) (*CEEMSServer, func(), error) {
 //	@description
 //	@description	If basic auth is enabled, all the endpoints require authentication.
 //	@description
-//	@description	All the endpoints, except `health` and `demo`, must send a user-agent header.
+//	@description	All the endpoints, except `health`, `swagger`, `debug` and `demo`,
+//	@description	must send a user-agent header.
 //	@description
 //	@description				Timestamps must be specified in milliseconds, unless otherwise specified.
 //