From c759d37fd388c2e16ef6ef4bf770e46d117c59ea Mon Sep 17 00:00:00 2001
From: Neeraj Gupta <254676+ua741@users.noreply.github.com>
Date: Fri, 16 Aug 2024 11:40:31 +0530
Subject: [PATCH] Add entity type validation

---
 server/ente/userentity/entity.go | 15 ++++++++++++++-
 server/pkg/api/userentity.go     |  8 ++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/server/ente/userentity/entity.go b/server/ente/userentity/entity.go
index 4d75a04c77..fed275238f 100644
--- a/server/ente/userentity/entity.go
+++ b/server/ente/userentity/entity.go
@@ -1,6 +1,10 @@
 package userentity
 
-import "github.com/ente-io/museum/ente/base"
+import (
+	"fmt"
+	"github.com/ente-io/museum/ente"
+	"github.com/ente-io/museum/ente/base"
+)
 
 type EntityType string
 
@@ -11,6 +15,15 @@ const (
 	Profile EntityType = "profile"
 )
 
+func (et EntityType) IsValid() error {
+
+	switch et {
+	case Location, Person, Profile:
+		return nil
+	}
+	return ente.NewBadRequestWithMessage(fmt.Sprintf("Invalid EntityType: %s", et))
+}
+
 func (et EntityType) GetNewID() (*string, error) {
 	return base.NewID(string(et))
 }
diff --git a/server/pkg/api/userentity.go b/server/pkg/api/userentity.go
index 004eda6137..492e9c3bd9 100644
--- a/server/pkg/api/userentity.go
+++ b/server/pkg/api/userentity.go
@@ -25,6 +25,10 @@ func (h *UserEntityHandler) CreateKey(c *gin.Context) {
 			stacktrace.Propagate(ente.ErrBadRequest, fmt.Sprintf("Request binding failed %s", err)))
 		return
 	}
+	if err := request.Type.IsValid(); err != nil {
+		handler.Error(c, stacktrace.Propagate(err, "Invalid EntityType"))
+		return
+	}
 	err := h.Controller.CreateKey(c, request)
 	if err != nil {
 		handler.Error(c, stacktrace.Propagate(err, "Failed to create CreateKey"))
@@ -57,6 +61,10 @@ func (h *UserEntityHandler) CreateEntity(c *gin.Context) {
 			stacktrace.Propagate(ente.ErrBadRequest, fmt.Sprintf("Request binding failed %s", err)))
 		return
 	}
+	if err := request.Type.IsValid(); err != nil {
+		handler.Error(c, stacktrace.Propagate(err, "Invalid EntityType"))
+		return
+	}
 	resp, err := h.Controller.CreateEntity(c, request)
 	if err != nil {
 		handler.Error(c, stacktrace.Propagate(err, "Failed to create CreateEntity"))