From dce16ad393922c913e55326d26e76fcc611a43f6 Mon Sep 17 00:00:00 2001
From: Daniel <dhaavi@users.noreply.github.com>
Date: Wed, 13 Apr 2022 11:18:40 +0200
Subject: [PATCH] Add go profiling debug api endpoints

---
 api/endpoints_debug.go | 87 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 87 insertions(+)

diff --git a/api/endpoints_debug.go b/api/endpoints_debug.go
index 9e81dae..163e948 100644
--- a/api/endpoints_debug.go
+++ b/api/endpoints_debug.go
@@ -2,10 +2,12 @@ package api
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"net/http"
 	"os"
 	"runtime/pprof"
+	"time"
 
 	"github.com/safing/portbase/utils/debug"
 )
@@ -41,6 +43,42 @@ func registerDebugEndpoints() error {
 		return err
 	}
 
+	if err := RegisterEndpoint(Endpoint{
+		Path:        "debug/cpu",
+		Read:        PermitAnyone,
+		DataFunc:    handleCPUProfile,
+		Name:        "Get CPU Profile",
+		Description: "",
+		Parameters: []Parameter{{
+			Method:      http.MethodGet,
+			Field:       "duration",
+			Value:       "10s",
+			Description: "Specify the formatting style. The default is simple markdown formatting.",
+		}},
+	}); err != nil {
+		return err
+	}
+
+	if err := RegisterEndpoint(Endpoint{
+		Path:        "debug/heap",
+		Read:        PermitAnyone,
+		DataFunc:    handleHeapProfile,
+		Name:        "Get Heap Profile",
+		Description: "",
+	}); err != nil {
+		return err
+	}
+
+	if err := RegisterEndpoint(Endpoint{
+		Path:        "debug/allocs",
+		Read:        PermitAnyone,
+		DataFunc:    handleAllocsProfile,
+		Name:        "Get Allocs Profile",
+		Description: "",
+	}); err != nil {
+		return err
+	}
+
 	if err := RegisterEndpoint(Endpoint{
 		Path:        "debug/info",
 		Read:        PermitAnyone,
@@ -90,6 +128,55 @@ func printStack(_ *Request) (msg string, err error) {
 	return "stack printed to stdout", nil
 }
 
+// handleCPUProfile returns the CPU profile.
+func handleCPUProfile(ar *Request) (data []byte, err error) {
+	// Parse duration.
+	duration := 10 * time.Second
+	if durationOption := ar.Request.URL.Query().Get("duration"); durationOption != "" {
+		parsedDuration, err := time.ParseDuration(durationOption)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse duration: %w", err)
+		}
+		duration = parsedDuration
+	}
+
+	// Start CPU profiling.
+	buf := new(bytes.Buffer)
+	if err := pprof.StartCPUProfile(buf); err != nil {
+		return nil, fmt.Errorf("failed to start cpu profile: %w", err)
+	}
+
+	// Wait for the specified duration.
+	select {
+	case <-time.After(duration):
+	case <-ar.Context().Done():
+		pprof.StopCPUProfile()
+		return nil, context.Canceled
+	}
+
+	// Stop CPU profiling and return data.
+	pprof.StopCPUProfile()
+	return buf.Bytes(), nil
+}
+
+// handleHeapProfile returns the Heap profile.
+func handleHeapProfile(ar *Request) (data []byte, err error) {
+	buf := new(bytes.Buffer)
+	if err := pprof.Lookup("heap").WriteTo(buf, 0); err != nil {
+		return nil, fmt.Errorf("failed to write heap profile: %w", err)
+	}
+	return buf.Bytes(), nil
+}
+
+// handleAllocsProfile returns the Allocs profile.
+func handleAllocsProfile(ar *Request) (data []byte, err error) {
+	buf := new(bytes.Buffer)
+	if err := pprof.Lookup("allocs").WriteTo(buf, 0); err != nil {
+		return nil, fmt.Errorf("failed to write allocs profile: %w", err)
+	}
+	return buf.Bytes(), nil
+}
+
 // debugInfo returns the debugging information for support requests.
 func debugInfo(ar *Request) (data []byte, err error) {
 	// Create debug information helper.