From 28e53065bc070d4edae80fff3d9b425737f54ac1 Mon Sep 17 00:00:00 2001
From: Yash Mehrotra <yashmehrotra95@gmail.com>
Date: Wed, 2 Oct 2024 23:24:28 +0530
Subject: [PATCH] fix: set user object correctly in kratos auth

---
 agent/agent.go |  2 +-
 auth/kratos.go | 25 +++++++++++++------------
 db/agents.go   |  4 +++-
 3 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/agent/agent.go b/agent/agent.go
index 52381818e..3ae5a0e61 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -19,7 +19,7 @@ func generateAgent(ctx context.Context, body api.GenerateAgentRequest) (*api.Gen
 		return nil, fmt.Errorf("failed to generate username and password: %w", err)
 	}
 
-	person, err := db.CreatePerson(ctx, username, fmt.Sprintf("%s@local", username), "agent")
+	person, err := db.CreatePerson(ctx, username, fmt.Sprintf("%s@local", username), db.PersonTypeAgent)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create a new person: %w", err)
 	}
diff --git a/auth/kratos.go b/auth/kratos.go
index 97a419a17..c41a743fb 100644
--- a/auth/kratos.go
+++ b/auth/kratos.go
@@ -12,6 +12,7 @@ import (
 	"github.com/flanksource/commons/rand"
 	"github.com/flanksource/duty/context"
 	"github.com/flanksource/duty/models"
+	"github.com/flanksource/incident-commander/db"
 	"github.com/google/uuid"
 	"github.com/labstack/echo/v4"
 	client "github.com/ory/client-go"
@@ -164,20 +165,20 @@ func (k *kratosMiddleware) Session(next echo.HandlerFunc) echo.HandlerFunc {
 			return c.String(http.StatusUnauthorized, "Authorization Error")
 		}
 
-		person := models.Person{ID: uid}
-
-		if traits, ok := session.Identity.GetTraits().(map[string]any); ok {
-			if e, ok := traits["email"].(string); ok {
-				person.Email = e
-			}
+		person, err := db.GetUserByID(ctx, uid.String())
+		if err != nil {
+			ctx.GetSpan().RecordError(err)
+			return c.String(http.StatusInternalServerError, "Error fetching user details from database")
+		}
 
-			switch v := traits["name"].(type) {
-			case map[string]string:
-				person.Name = fmt.Sprintf("%s %s", v["first"], v["last"])
+		if person.Type == db.PersonTypeAgent {
+			agent, err := db.FindAgent(ctx, person.Name)
+			if err != nil {
+				ctx.GetSpan().RecordError(err)
+				return c.String(http.StatusInternalServerError, "Error fetching agent details from database")
 			}
-
-			if agent, ok := traits["agent"].(models.Agent); ok {
-				ctx = ctx.WithAgent(agent)
+			if agent != nil {
+				ctx = ctx.WithAgent(*agent)
 			}
 		}
 
diff --git a/db/agents.go b/db/agents.go
index 9adba0cca..421b4385a 100644
--- a/db/agents.go
+++ b/db/agents.go
@@ -11,6 +11,8 @@ import (
 	"gorm.io/gorm"
 )
 
+const PersonTypeAgent = "agent"
+
 func FindFirstAgent(ctx context.Context, names ...string) (*models.Agent, error) {
 	for _, name := range names {
 		agent, err := FindAgent(ctx, name)
@@ -69,7 +71,7 @@ func GetOrCreateAgent(ctx context.Context, name string) (*models.Agent, error) {
 }
 
 func CreateAgent(ctx context.Context, name string, personID *uuid.UUID, properties map[string]string) error {
-	properties = collections.MergeMap(properties, map[string]string{"type": "agent"})
+	properties = collections.MergeMap(properties, map[string]string{"type": PersonTypeAgent})
 
 	a := models.Agent{
 		Name:       name,