diff --git a/controllers/enrollmentkeys.go b/controllers/enrollmentkeys.go
index ccdf0ade..d827f53f 100644
--- a/controllers/enrollmentkeys.go
+++ b/controllers/enrollmentkeys.go
@@ -242,9 +242,8 @@ func updateEnrollmentKey(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	relayId := uuid.Nil
 	if enrollmentKeyBody.Relay != "" {
-		relayId, err = uuid.Parse(enrollmentKeyBody.Relay)
+		_, err = uuid.Parse(enrollmentKeyBody.Relay)
 		if err != nil {
 			slog.Error("error parsing relay id", "error", err)
 			logic.ReturnErrorResponse(w, r, logic.FormatError(err, "badrequest"))
@@ -253,7 +252,7 @@ func updateEnrollmentKey(w http.ResponseWriter, r *http.Request) {
 	}
 	currKey, _ := logic.GetEnrollmentKey(keyId)
 
-	newEnrollmentKey, err := logic.UpdateEnrollmentKey(keyId, relayId, enrollmentKeyBody.Groups)
+	newEnrollmentKey, err := logic.UpdateEnrollmentKey(keyId, &enrollmentKeyBody)
 	if err != nil {
 		slog.Error("failed to update enrollment key", "error", err)
 		logic.ReturnErrorResponse(w, r, logic.FormatError(err, "internal"))
diff --git a/logic/enrollmentkey.go b/logic/enrollmentkey.go
index 801f3492..6a757602 100644
--- a/logic/enrollmentkey.go
+++ b/logic/enrollmentkey.go
@@ -96,14 +96,19 @@ func CreateEnrollmentKey(uses int, expiration time.Time, networks,
 }
 
 // UpdateEnrollmentKey - updates an existing enrollment key's associated relay
-func UpdateEnrollmentKey(keyId string, relayId uuid.UUID, groups []models.TagID) (*models.EnrollmentKey, error) {
+func UpdateEnrollmentKey(keyId string, updates *models.APIEnrollmentKey) (*models.EnrollmentKey, error) {
 	key, err := GetEnrollmentKey(keyId)
 	if err != nil {
 		return nil, err
 	}
 
-	if relayId != uuid.Nil {
-		relayNode, err := GetNodeByID(relayId.String())
+	relayID := uuid.Nil
+	if updates.Relay != "" {
+		relayID = uuid.MustParse(updates.Relay)
+	}
+
+	if relayID != uuid.Nil {
+		relayNode, err := GetNodeByID(relayID.String())
 		if err != nil {
 			return nil, err
 		}
@@ -115,8 +120,13 @@ func UpdateEnrollmentKey(keyId string, relayId uuid.UUID, groups []models.TagID)
 		}
 	}
 
-	key.Relay = relayId
-	key.Groups = groups
+	if relayID != uuid.Nil {
+		updates.AutoAssignGateway = false
+	}
+
+	key.Relay = relayID
+	key.Groups = updates.Groups
+	key.AutoAssignGateway = updates.AutoAssignGateway
 	if err = upsertEnrollmentKey(&key); err != nil {
 		return nil, err
 	}