From 63abfe369875e12da53375346f5fadcc8514edd4 Mon Sep 17 00:00:00 2001
From: Bryan Frimin <bryan@getprobo.com>
Date: Mon, 29 Dec 2025 11:45:49 +0100
Subject: [PATCH] Fix not token secret

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
---
 pkg/probod/probod.go                            | 1 +
 pkg/server/api/api.go                           | 2 ++
 pkg/server/api/connect/v1/api_key_middleware.go | 4 ++--
 pkg/server/api/console/v1/resolver.go           | 9 +--------
 pkg/server/api/mcp/v1/v1_handler.go             | 4 ++--
 pkg/server/server.go                            | 2 ++
 6 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/pkg/probod/probod.go b/pkg/probod/probod.go
index 7ef207791..3267c4517 100644
--- a/pkg/probod/probod.go
+++ b/pkg/probod/probod.go
@@ -426,6 +426,7 @@ func (impl *Implm) Run(
 			BaseURL:           impl.cfg.BaseURL,
 			Agent:             agent,
 			CustomDomainCname: impl.cfg.CustomDomains.CnameTarget,
+			TokenSecret:       impl.cfg.Auth.Cookie.Secret,
 			Logger:            l.Named("http.server"),
 			Cookie: securecookie.Config{
 				Name:     impl.cfg.Auth.Cookie.Name,
diff --git a/pkg/server/api/api.go b/pkg/server/api/api.go
index 978190507..0a9818525 100644
--- a/pkg/server/api/api.go
+++ b/pkg/server/api/api.go
@@ -133,6 +133,7 @@ func NewServer(cfg Config) (*Server, error) {
 			cfg.Logger.Named("mcp.v1"),
 			cfg.Probo,
 			cfg.IAM,
+			cfg.TokenSecret,
 		),
 		slackHandler: slack_v1.NewMux(
 			cfg.Logger.Named("slack.v1"),
@@ -143,6 +144,7 @@ func NewServer(cfg Config) (*Server, error) {
 			cfg.Logger.Named("connect.v1"),
 			cfg.IAM,
 			cfg.Cookie,
+			cfg.TokenSecret,
 			cfg.BaseURL,
 		),
 	}, nil
diff --git a/pkg/server/api/connect/v1/api_key_middleware.go b/pkg/server/api/connect/v1/api_key_middleware.go
index 7bb604202..87f09c9ba 100644
--- a/pkg/server/api/connect/v1/api_key_middleware.go
+++ b/pkg/server/api/connect/v1/api_key_middleware.go
@@ -36,13 +36,13 @@ func APIKeyFromContext(ctx context.Context) *coredata.PersonalAPIKey {
 	return apiKey
 }
 
-func NewAPIKeyMiddleware(svc *iam.Service) func(next http.Handler) http.Handler {
+func NewAPIKeyMiddleware(svc *iam.Service, tokenSecret string) func(next http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(
 			func(w http.ResponseWriter, r *http.Request) {
 				ctx := r.Context()
 
-				tokenValue, err := securetoken.Get(r, "")
+				tokenValue, err := securetoken.Get(r, tokenSecret)
 				if err != nil {
 					next.ServeHTTP(w, r)
 					return
diff --git a/pkg/server/api/console/v1/resolver.go b/pkg/server/api/console/v1/resolver.go
index d388948ae..272fe6ca1 100644
--- a/pkg/server/api/console/v1/resolver.go
+++ b/pkg/server/api/console/v1/resolver.go
@@ -82,7 +82,7 @@ func NewMux(
 	safeRedirect := &saferedirect.SafeRedirect{AllowedHost: baseURL.Host()}
 
 	r.Use(connect_v1.NewSessionMiddleware(iamSvc, cookieConfig))
-	r.Use(connect_v1.NewAPIKeyMiddleware(iamSvc))
+	r.Use(connect_v1.NewAPIKeyMiddleware(iamSvc, tokenSecret))
 
 	config := schema.Config{
 		Resolvers: &Resolver{
@@ -317,13 +317,6 @@ func (r *Resolver) ProboService(ctx context.Context, tenantID gid.TenantID) *pro
 
 func (r *Resolver) MustAuthorize(ctx context.Context, entityID gid.GID, action iam.Action) {
 	identity := connect_v1.IdentityFromContext(ctx)
-	// TODO: Add API key authorization
-	// apiKey := connect_v1.APIKeyFromContext(ctx)
-
-	// var credentialID *gid.GID
-	// if apiKey != nil {
-	// 	credentialID = &apiKey.ID
-	// }
 
 	err := r.iam.Authorizer.Authorize(
 		ctx,
diff --git a/pkg/server/api/mcp/v1/v1_handler.go b/pkg/server/api/mcp/v1/v1_handler.go
index 2e5cbdc38..de36c443c 100644
--- a/pkg/server/api/mcp/v1/v1_handler.go
+++ b/pkg/server/api/mcp/v1/v1_handler.go
@@ -20,7 +20,7 @@ func (r *Resolver) ProboService(ctx context.Context, objectID gid.GID) *probo.Te
 	return r.proboSvc.WithTenant(objectID.TenantID())
 }
 
-func NewMux(logger *log.Logger, proboSvc *probo.Service, iamSvc *iam.Service) *chi.Mux {
+func NewMux(logger *log.Logger, proboSvc *probo.Service, iamSvc *iam.Service, tokenSecret string) *chi.Mux {
 	logger = logger.Named("mcp.v1")
 
 	logger.Info("initializing MCP server")
@@ -52,7 +52,7 @@ func NewMux(logger *log.Logger, proboSvc *probo.Service, iamSvc *iam.Service) *c
 	)
 
 	r := chi.NewMux()
-	r.Use(connect_v1.NewAPIKeyMiddleware(iamSvc))
+	r.Use(connect_v1.NewAPIKeyMiddleware(iamSvc, tokenSecret))
 	r.Handle("/", RequireAPIKeyHandler(logger, handler))
 
 	logger.Info("MCP server initialized successfully")
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 2fade103e..82dea29a4 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -47,6 +47,7 @@ type Config struct {
 	Trust             *trust.Service
 	Slack             *slack.Service
 	Cookie            securecookie.Config
+	TokenSecret       string
 	ConnectorRegistry *connector.ConnectorRegistry
 	Agent             *agents.Agent
 	CustomDomainCname string
@@ -72,6 +73,7 @@ func NewServer(cfg Config) (*Server, error) {
 		Trust:             cfg.Trust,
 		Slack:             cfg.Slack,
 		Cookie:            cfg.Cookie,
+		TokenSecret:       cfg.TokenSecret,
 		ConnectorRegistry: cfg.ConnectorRegistry,
 		CustomDomainCname: cfg.CustomDomainCname,
 		Logger:            cfg.Logger.Named("api"),