From 7727de5142b131e7863ef9b2c6d2ccddaa8b2c32 Mon Sep 17 00:00:00 2001
From: William Dumont <william.dumont@grafana.com>
Date: Wed, 22 May 2024 11:25:51 +0200
Subject: [PATCH] add error checks

---
 .../component/prometheus/relabel/relabel.go   |  5 ++++-
 .../livedebugging/debug_stream_manager.go     | 21 ++++++++++++++++---
 .../debug_stream_manager_test.go              |  9 ++++++++
 internal/web/api/api.go                       | 13 ++++++++++++
 4 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/internal/component/prometheus/relabel/relabel.go b/internal/component/prometheus/relabel/relabel.go
index 7cff5c4fd7..015983fc0d 100644
--- a/internal/component/prometheus/relabel/relabel.go
+++ b/internal/component/prometheus/relabel/relabel.go
@@ -23,9 +23,11 @@ import (
 	"go.uber.org/atomic"
 )
 
+const name = "prometheus.relabel"
+
 func init() {
 	component.Register(component.Registration{
-		Name:      "prometheus.relabel",
+		Name:      name,
 		Stability: featuregate.StabilityGenerallyAvailable,
 		Args:      Arguments{},
 		Exports:   Exports{},
@@ -107,6 +109,7 @@ func New(o component.Options, args Arguments) (*Component, error) {
 	if err != nil {
 		return nil, err
 	}
+	debuggingStreamHandler.(livedebugging.DebugStreamManager).Register(name)
 
 	data, err := o.GetServiceData(labelstore.ServiceName)
 	if err != nil {
diff --git a/internal/service/livedebugging/debug_stream_manager.go b/internal/service/livedebugging/debug_stream_manager.go
index a457a25f0e..716e39599f 100644
--- a/internal/service/livedebugging/debug_stream_manager.go
+++ b/internal/service/livedebugging/debug_stream_manager.go
@@ -4,6 +4,10 @@ import "sync"
 
 // DebugStreamManager defines the operations for managing debug streams.
 type DebugStreamManager interface {
+	// Register a component by name
+	Register(componentName string)
+	// IsRegistered returns true if the component name was registered
+	IsRegistered(componentName string) bool
 	// Stream streams data for a given componentID.
 	Stream(componentID string, data string)
 	// SetStream assigns a debug stream callback to a componentID.
@@ -13,14 +17,16 @@ type DebugStreamManager interface {
 }
 
 type debugStreamManager struct {
-	loadMut sync.RWMutex
-	streams map[string]map[string]func(string)
+	loadMut              sync.RWMutex
+	streams              map[string]map[string]func(string)
+	registeredComponents map[string]struct{}
 }
 
 // NewDebugStreamManager creates a new instance of DebugStreamManager.
 func NewDebugStreamManager() *debugStreamManager {
 	return &debugStreamManager{
-		streams: make(map[string]map[string]func(string)),
+		streams:              make(map[string]map[string]func(string)),
+		registeredComponents: make(map[string]struct{}),
 	}
 }
 
@@ -48,3 +54,12 @@ func (s *debugStreamManager) DeleteStream(streamID string, componentID string) {
 	defer s.loadMut.Unlock()
 	delete(s.streams[componentID], streamID)
 }
+
+func (s *debugStreamManager) Register(componentName string) {
+	s.registeredComponents[componentName] = struct{}{}
+}
+
+func (s *debugStreamManager) IsRegistered(componentName string) bool {
+	_, exist := s.registeredComponents[componentName]
+	return exist
+}
diff --git a/internal/service/livedebugging/debug_stream_manager_test.go b/internal/service/livedebugging/debug_stream_manager_test.go
index 5f18ca0b60..19f6ff591c 100644
--- a/internal/service/livedebugging/debug_stream_manager_test.go
+++ b/internal/service/livedebugging/debug_stream_manager_test.go
@@ -6,6 +6,15 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestRegister(t *testing.T) {
+	manager := NewDebugStreamManager()
+	require.False(t, manager.IsRegistered("type1"))
+	manager.Register("type1")
+	require.True(t, manager.IsRegistered("type1"))
+	// registering a component name that has already been registered does not do anything
+	require.NotPanics(t, func() { manager.Register("type1") })
+}
+
 func TestStream(t *testing.T) {
 	manager := NewDebugStreamManager()
 	componentID := "component1"
diff --git a/internal/web/api/api.go b/internal/web/api/api.go
index 8edfd2b9f7..a9d5ba9a6c 100644
--- a/internal/web/api/api.go
+++ b/internal/web/api/api.go
@@ -6,6 +6,7 @@ package api
 
 import (
 	"encoding/json"
+	"fmt"
 	"math/rand"
 	"net/http"
 	"path"
@@ -119,6 +120,18 @@ func (a *AlloyAPI) startDebugStream() http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		vars := mux.Vars(r)
 		componentID := vars["id"]
+		requestedComponent := component.ParseID(componentID)
+
+		component, err := a.alloy.GetComponent(requestedComponent, component.InfoOptions{})
+		if err != nil {
+			http.NotFound(w, r)
+			return
+		}
+
+		if !a.debuggingStreamHandler.IsRegistered(component.ComponentName) {
+			http.Error(w, fmt.Sprintf("Live debugging is not supported by %s", component.ComponentName), http.StatusInternalServerError)
+			return
+		}
 
 		// Buffer of 1000 entries to handle load spikes and prevent this functionality from eating up too much memory.
 		// TODO: in the future we may want to make this value configurable to handle heavy load