Add connector service and http handlers

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

pkg/probo/connector_service.go

@@ -0,0 +1,75 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package probo
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/getprobo/probo/pkg/connector"
+	"github.com/getprobo/probo/pkg/coredata"
+	"github.com/getprobo/probo/pkg/gid"
+	"go.gearno.de/kit/pg"
+)
+
+type (
+	ConnectorService struct {
+		svc *TenantService
+	}
+
+	CreateOrUpdateConnectorRequest struct {
+		OrganizationID gid.GID
+		Name           string
+		Type           string
+		Connection     connector.Connection
+	}
+)
+
+func (s *ConnectorService) CreateOrUpdate(ctx context.Context, req CreateOrUpdateConnectorRequest) (*coredata.Connector, error) {
+	connectorID, err := gid.NewGID(s.svc.scope.GetTenantID(), coredata.ConnectorEntityType)
+	if err != nil {
+		return nil, fmt.Errorf("cannot create connector global id: %w", err)
+	}
+
+	now := time.Now()
+
+	connector := &coredata.Connector{
+		ID:             connectorID,
+		OrganizationID: req.OrganizationID,
+		Name:           req.Name,
+		Type:           req.Type,
+		Connection:     req.Connection,
+		CreatedAt:      now,
+		UpdatedAt:      now,
+	}
+
+	err = s.svc.pg.WithConn(
+		ctx,
+		func(conn pg.Conn) error {
+			if err := connector.Upsert(ctx, conn, s.svc.scope); err != nil {
+				return fmt.Errorf("cannot upsert connector: %w", err)
+			}
+
+			return nil
+		},
+	)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return connector, nil
+}

pkg/probo/service.go

@@ -49,6 +49,7 @@ type (
 		Controls                *ControlService
 		Risks                   *RiskService
 		VendorComplianceReports *VendorComplianceReportService
+		Connectors              *ConnectorService
 	}
 )
 
@@ -90,5 +91,6 @@ func (s *Service) WithTenant(tenantID gid.TenantID) *TenantService {
 	tenantService.Controls = &ControlService{svc: tenantService}
 	tenantService.Risks = &RiskService{svc: tenantService}
 	tenantService.VendorComplianceReports = &VendorComplianceReportService{svc: tenantService}
+	tenantService.Connectors = &ConnectorService{svc: tenantService}
 	return tenantService
 }

pkg/server/api/api.go

@@ -18,6 +18,7 @@ import (
 	"errors"
 	"net/http"
 
+	"github.com/getprobo/probo/pkg/connector"
 	"github.com/getprobo/probo/pkg/probo"
 	console_v1 "github.com/getprobo/probo/pkg/server/api/console/v1"
 	"github.com/getprobo/probo/pkg/usrmgr"
@@ -32,6 +33,7 @@ type (
 		Probo             *probo.Service
 		Usrmgr            *usrmgr.Service
 		Auth              console_v1.AuthConfig
+		ConnectorRegistry *connector.ConnectorRegistry
 	}
 
 	Server struct {
@@ -101,7 +103,7 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	router.Use(cors.Handler(corsOpts))
 
 	// Mount the console API with authentication
-	router.Mount("/console/v1", console_v1.NewMux(s.cfg.Probo, s.cfg.Usrmgr, s.cfg.Auth))
+	router.Mount("/console/v1", console_v1.NewMux(s.cfg.Probo, s.cfg.Usrmgr, s.cfg.Auth, s.cfg.ConnectorRegistry))
 
 	router.ServeHTTP(w, r)
 }

pkg/server/api/console/v1/resolver.go

@@ -28,6 +28,7 @@ import (
 	"github.com/99designs/gqlgen/graphql/handler/extension"
 	"github.com/99designs/gqlgen/graphql/handler/transport"
 	"github.com/99designs/gqlgen/graphql/playground"
+	"github.com/getprobo/probo/pkg/connector"
 	"github.com/getprobo/probo/pkg/coredata"
 	"github.com/getprobo/probo/pkg/gid"
 	"github.com/getprobo/probo/pkg/probo"
@@ -72,7 +73,7 @@ func UserFromContext(ctx context.Context) *coredata.User {
 	return user
 }
 
-func NewMux(proboSvc *probo.Service, usrmgrSvc *usrmgr.Service, authCfg AuthConfig) *chi.Mux {
+func NewMux(proboSvc *probo.Service, usrmgrSvc *usrmgr.Service, authCfg AuthConfig, connectorRegistry *connector.ConnectorRegistry) *chi.Mux {
 	r := chi.NewMux()
 
 	r.Post("/auth/register", SignUpHandler(usrmgrSvc, authCfg))
@@ -82,6 +83,63 @@ func NewMux(proboSvc *probo.Service, usrmgrSvc *usrmgr.Service, authCfg AuthConf
 	r.Post("/auth/forget-password", ForgetPasswordHandler(usrmgrSvc, authCfg))
 	r.Post("/auth/reset-password", ResetPasswordHandler(usrmgrSvc, authCfg))
 
+	r.Get("/connectors/initiate", WithSession(usrmgrSvc, authCfg, func(w http.ResponseWriter, r *http.Request) {
+		session := SessionFromContext(r.Context())
+		if session == nil {
+			panic(fmt.Errorf("session not found"))
+		}
+
+		// TODO: check if current user has access to the organization
+
+		connectorID := r.URL.Query().Get("connector_id")
+		organizationID := r.URL.Query().Get("organization_id")
+
+		redirectURL, err := connectorRegistry.Initiate(r.Context(), connectorID, organizationID, r)
+		if err != nil {
+			panic(fmt.Errorf("cannot initiate connector: %w", err))
+		}
+
+		http.Redirect(w, r, redirectURL, http.StatusSeeOther)
+	}))
+
+	r.Get("/connectors/complete", WithSession(usrmgrSvc, authCfg, func(w http.ResponseWriter, r *http.Request) {
+		session := SessionFromContext(r.Context())
+		if session == nil {
+			panic(fmt.Errorf("session not found"))
+		}
+
+		// TODO: check if current user has access to the organization
+
+		connectorID := r.URL.Query().Get("connector_id")
+		organizationIDString := r.URL.Query().Get("organization_id")
+
+		connection, err := connectorRegistry.Complete(r.Context(), connectorID, organizationIDString, r)
+		if err != nil {
+			panic(fmt.Errorf("failed to complete connector: %w", err))
+		}
+
+		organizationID, err := gid.ParseGID(organizationIDString)
+		if err != nil {
+			panic(fmt.Errorf("failed to parse organization id: %w", err))
+		}
+
+		tenantID := session.ID.TenantID()
+		_, err = proboSvc.WithTenant(tenantID).Connectors.CreateOrUpdate(
+			r.Context(),
+			probo.CreateOrUpdateConnectorRequest{
+				OrganizationID: organizationID,
+				Name:           connectorID,
+				Type:           string(connection.Type()),
+				Connection:     connection,
+			},
+		)
+		if err != nil {
+			panic(fmt.Errorf("failed to create or update connector: %w", err))
+		}
+
+		http.Redirect(w, r, "/foo", http.StatusSeeOther)
+	}))
+
 	r.Get("/", playground.Handler("GraphQL", "/api/console/v1/query"))
 	r.Post("/query", graphqlHandler(proboSvc, usrmgrSvc, authCfg))
 
@@ -139,7 +197,7 @@ func graphqlHandler(proboSvc *probo.Service, usrmgrSvc *usrmgr.Service, authCfg
 		},
 	)
 
-	return func(w http.ResponseWriter, r *http.Request) {
+	return WithSession(usrmgrSvc, authCfg, func(w http.ResponseWriter, r *http.Request) {
 		ctx := r.Context()
 
 		// Hack to capture the panic value, because gqlgen execute resolver in a different goroutine.
@@ -147,6 +205,18 @@ func graphqlHandler(proboSvc *probo.Service, usrmgrSvc *usrmgr.Service, authCfg
 		var panicValue any
 		ctx = context.WithValue(ctx, panicValueContextKey, &panicValue)
 
+		srv.ServeHTTP(w, r.WithContext(ctx))
+
+		if panicValue != nil {
+			panic(panicValue)
+		}
+	})
+}
+
+func WithSession(usrmgrSvc *usrmgr.Service, authCfg AuthConfig, next http.HandlerFunc) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		ctx := r.Context()
+
 		cookieValue, err := securecookie.Get(r, securecookie.DefaultConfig(
 			authCfg.CookieName,
 			authCfg.CookieSecret,
@@ -156,7 +226,7 @@ func graphqlHandler(proboSvc *probo.Service, usrmgrSvc *usrmgr.Service, authCfg
 				panic(fmt.Errorf("failed to get session: %w", err))
 			}
 
-			srv.ServeHTTP(w, r)
+			next(w, r)
 			return
 		}
 
@@ -167,7 +237,7 @@ func graphqlHandler(proboSvc *probo.Service, usrmgrSvc *usrmgr.Service, authCfg
 				authCfg.CookieSecret,
 			))
 
-			srv.ServeHTTP(w, r)
+			next(w, r)
 			return
 		}
 
@@ -178,7 +248,7 @@ func graphqlHandler(proboSvc *probo.Service, usrmgrSvc *usrmgr.Service, authCfg
 				authCfg.CookieSecret,
 			))
 
-			srv.ServeHTTP(w, r)
+			next(w, r)
 			return
 		}
 
@@ -189,7 +259,7 @@ func graphqlHandler(proboSvc *probo.Service, usrmgrSvc *usrmgr.Service, authCfg
 				authCfg.CookieSecret,
 			))
 
-			srv.ServeHTTP(w, r)
+			next(w, r)
 			return
 		}
 
@@ -202,16 +272,12 @@ func graphqlHandler(proboSvc *probo.Service, usrmgrSvc *usrmgr.Service, authCfg
 		ctx = context.WithValue(ctx, userContextKey, user)
 		ctx = context.WithValue(ctx, userTenantContextKey, &tenantIDs)
 
-		srv.ServeHTTP(w, r.WithContext(ctx))
+		next(w, r.WithContext(ctx))
 
-		if panicValue != nil {
+		// Update session after the handler completes
-			panic(panicValue)
+		if err := usrmgrSvc.UpdateSession(ctx, session); err != nil {
-		}
-
-		if err := usrmgrSvc.UpdateSession(r.Context(), session); err != nil {
 			panic(fmt.Errorf("failed to update session: %w", err))
 		}
-
 	}
 }
 

pkg/server/server.go

@@ -19,6 +19,7 @@ import (
 	"net/http"
 	"strings"
 
+	"github.com/getprobo/probo/pkg/connector"
 	"github.com/getprobo/probo/pkg/probo"
 	"github.com/getprobo/probo/pkg/server/api"
 	console_v1 "github.com/getprobo/probo/pkg/server/api/console/v1"
@@ -33,6 +34,7 @@ type Config struct {
 	Probo             *probo.Service
 	Usrmgr            *usrmgr.Service
 	Auth              console_v1.AuthConfig
+	ConnectorRegistry *connector.ConnectorRegistry
 }
 
 // Server represents the main server that handles both API and frontend requests
@@ -50,6 +52,7 @@ func NewServer(cfg Config) (*Server, error) {
 		Probo:             cfg.Probo,
 		Usrmgr:            cfg.Usrmgr,
 		Auth:              cfg.Auth,
+		ConnectorRegistry: cfg.ConnectorRegistry,
 	}
 	apiServer, err := api.NewServer(apiCfg)
 	if err != nil {