From 118cf2836c9bff0d033125c89fbc7661bb23e7e5 Mon Sep 17 00:00:00 2001
From: demoManito <1430482733@qq.com>
Date: Mon, 26 Sep 2022 23:45:54 +0800
Subject: [PATCH] fix relationship.go maybe nil panic

---
 schema/relationship.go | 90 +++++++++++++++++++++++-------------------
 1 file changed, 50 insertions(+), 40 deletions(-)

diff --git a/schema/relationship.go b/schema/relationship.go
index bb8aeb6498..5e1078b156 100644
--- a/schema/relationship.go
+++ b/schema/relationship.go
@@ -152,36 +152,42 @@ func (schema *Schema) buildPolymorphicRelation(relation *Relationship, field *Fi
 		schema.err = fmt.Errorf("invalid polymorphic type %v for %v on field %s, missing field %s", relation.FieldSchema, schema, field.Name, polymorphic+"ID")
 	}
 
-	if schema.err == nil {
-		relation.References = append(relation.References, &Reference{
-			PrimaryValue: relation.Polymorphic.Value,
-			ForeignKey:   relation.Polymorphic.PolymorphicType,
-		})
+	relation.Type = has
 
-		primaryKeyField := schema.PrioritizedPrimaryField
-		if len(relation.foreignKeys) > 0 {
-			if primaryKeyField = schema.LookUpField(relation.foreignKeys[0]); primaryKeyField == nil || len(relation.foreignKeys) > 1 {
-				schema.err = fmt.Errorf("invalid polymorphic foreign keys %+v for %v on field %s", relation.foreignKeys, schema, field.Name)
-			}
-		}
+	if schema.err != nil {
+		return
+	}
 
-		// use same data type for foreign keys
-		if copyableDataType(primaryKeyField.DataType) {
-			relation.Polymorphic.PolymorphicID.DataType = primaryKeyField.DataType
-		}
-		relation.Polymorphic.PolymorphicID.GORMDataType = primaryKeyField.GORMDataType
-		if relation.Polymorphic.PolymorphicID.Size == 0 {
-			relation.Polymorphic.PolymorphicID.Size = primaryKeyField.Size
+	relation.References = append(relation.References, &Reference{
+		PrimaryValue: relation.Polymorphic.Value,
+		ForeignKey:   relation.Polymorphic.PolymorphicType,
+	})
+
+	primaryKeyField := schema.PrioritizedPrimaryField
+	if len(relation.foreignKeys) > 0 {
+		primaryKeyField = schema.LookUpField(relation.foreignKeys[0])
+		if primaryKeyField == nil || len(relation.foreignKeys) > 1 {
+			schema.err = fmt.Errorf("invalid polymorphic foreign keys %+v for %v on field %s", relation.foreignKeys, schema, field.Name)
 		}
+	}
+	if primaryKeyField == nil {
+		return
+	}
 
-		relation.References = append(relation.References, &Reference{
-			PrimaryKey:    primaryKeyField,
-			ForeignKey:    relation.Polymorphic.PolymorphicID,
-			OwnPrimaryKey: true,
-		})
+	// use same data type for foreign keys
+	if copyableDataType(primaryKeyField.DataType) {
+		relation.Polymorphic.PolymorphicID.DataType = primaryKeyField.DataType
+	}
+	relation.Polymorphic.PolymorphicID.GORMDataType = primaryKeyField.GORMDataType
+	if relation.Polymorphic.PolymorphicID.Size == 0 {
+		relation.Polymorphic.PolymorphicID.Size = primaryKeyField.Size
 	}
 
-	relation.Type = has
+	relation.References = append(relation.References, &Reference{
+		PrimaryKey:    primaryKeyField,
+		ForeignKey:    relation.Polymorphic.PolymorphicID,
+		OwnPrimaryKey: true,
+	})
 }
 
 func (schema *Schema) buildMany2ManyRelation(relation *Relationship, field *Field, many2many string) {
@@ -429,7 +435,6 @@ func (schema *Schema) guessRelation(relation *Relationship, field *Field, cgl gu
 			}
 		}
 	} else {
-		var primaryFields []*Field
 		var primarySchemaName = primarySchema.Name
 		if primarySchemaName == "" {
 			primarySchemaName = relation.FieldSchema.Name
@@ -466,32 +471,37 @@ func (schema *Schema) guessRelation(relation *Relationship, field *Field, cgl gu
 		}
 	}
 
-	if len(foreignFields) == 0 {
+	switch {
+	case len(foreignFields) == 0 || len(primaryFields) == 0:
 		reguessOrErr()
 		return
-	} else if len(relation.primaryKeys) > 0 {
+	case len(relation.primaryKeys) > 0:
 		for idx, primaryKey := range relation.primaryKeys {
-			if f := primarySchema.LookUpField(primaryKey); f != nil {
-				if len(primaryFields) < idx+1 {
-					primaryFields = append(primaryFields, f)
-				} else if f != primaryFields[idx] {
-					reguessOrErr()
-					return
-				}
-			} else {
+			f := primarySchema.LookUpField(primaryKey)
+			if f == nil {
+				reguessOrErr()
+				return
+			}
+			if len(primaryFields) < idx+1 {
+				primaryFields = append(primaryFields, f)
+				continue
+			}
+			if f != primaryFields[idx] {
 				reguessOrErr()
 				return
 			}
 		}
-	} else if len(primaryFields) == 0 {
+	default:
 		if len(foreignFields) == 1 && primarySchema.PrioritizedPrimaryField != nil {
 			primaryFields = append(primaryFields, primarySchema.PrioritizedPrimaryField)
-		} else if len(primarySchema.PrimaryFields) == len(foreignFields) {
+			break
+		}
+		if len(primarySchema.PrimaryFields) == len(foreignFields) {
 			primaryFields = append(primaryFields, primarySchema.PrimaryFields...)
-		} else {
-			reguessOrErr()
-			return
+			break
 		}
+		reguessOrErr()
+		return
 	}
 
 	// build references