44	"context" 
55	"fmt" 
66	"path/filepath" 
7+ 	"regexp" 
78	"strings" 
89	"sync" 
910	"text/template" 
@@ -21,7 +22,11 @@ import (
2122
2223	"github.com/argoproj/argo-cd/v2/pkg/apiclient/application" 
2324	"github.com/argoproj/argo-cd/v2/pkg/apis/application/v1alpha1" 
24- 	"gopkg.in/yaml.v2" 
25+ 
26+ 	"github.com/goccy/go-yaml" 
27+ 	"github.com/goccy/go-yaml/ast" 
28+ 	"github.com/goccy/go-yaml/parser" 
29+ 	"github.com/goccy/go-yaml/token" 
2530)
2631
2732// Stores some statistics about the results of a run 
@@ -459,8 +464,7 @@ func marshalParamsOverride(app *v1alpha1.Application, originalData []byte) ([]by
459464		if  strings .HasPrefix (app .Annotations [common .WriteBackTargetAnnotation ], common .HelmPrefix ) {
460465			images  :=  GetImagesAndAliasesFromApplication (app )
461466
462- 			helmNewValues  :=  yaml.MapSlice {}
463- 			err  =  yaml .Unmarshal (originalData , & helmNewValues )
467+ 			helmNewValues , err  :=  parser .ParseBytes (originalData , parser .ParseComments )
464468			if  err  !=  nil  {
465469				return  nil , err 
466470			}
@@ -488,7 +492,7 @@ func marshalParamsOverride(app *v1alpha1.Application, originalData []byte) ([]by
488492					if  helmParamVersion  ==  nil  {
489493						return  nil , fmt .Errorf ("%s parameter not found" , helmAnnotationParamVersion )
490494					}
491- 					err  =  setHelmValue (& helmNewValues , helmAnnotationParamVersion , helmParamVersion .Value )
495+ 					err  =  setHelmValue (helmNewValues , helmAnnotationParamVersion , helmParamVersion .Value )
492496					if  err  !=  nil  {
493497						return  nil , fmt .Errorf ("failed to set image parameter version value: %v" , err )
494498					}
@@ -499,13 +503,13 @@ func marshalParamsOverride(app *v1alpha1.Application, originalData []byte) ([]by
499503					return  nil , fmt .Errorf ("%s parameter not found" , helmAnnotationParamName )
500504				}
501505
502- 				err  =  setHelmValue (& helmNewValues , helmAnnotationParamName , helmParamName .Value )
506+ 				err  =  setHelmValue (helmNewValues , helmAnnotationParamName , helmParamName .Value )
503507				if  err  !=  nil  {
504508					return  nil , fmt .Errorf ("failed to set image parameter name value: %v" , err )
505509				}
506510			}
507511
508- 			override ,  err   =   yaml . Marshal (helmNewValues )
512+ 			override   =  [] byte (helmNewValues . String () )
509513		} else  {
510514			var  params  helmOverride 
511515			newParams  :=  helmOverride {
@@ -561,76 +565,118 @@ func mergeKustomizeOverride(t *kustomizeOverride, o *kustomizeOverride) {
561565	}
562566}
563567
564- // Check if a key exists in a MapSlice and return its index and value 
565- func  findHelmValuesKey (m  yaml.MapSlice , key  string ) (int , bool ) {
566- 	for  i , item  :=  range  m  {
567- 		if  item .Key  ==  key  {
568- 			return  i , true 
568+ func  splitKeys (input  string ) []string  {
569+ 	// Regular expression to match quoted and unquoted segments 
570+ 	re  :=  regexp .MustCompile (`'([^']*)'|"([^"]*)"|([^.".]+)` )
571+ 	matches  :=  re .FindAllStringSubmatch (input , - 1 )
572+ 
573+ 	var  result  []string 
574+ 	for  _ , match  :=  range  matches  {
575+ 		if  match [1 ] !=  ""  { // Single-quoted string 
576+ 			result  =  append (result , match [1 ])
577+ 		} else  if  match [2 ] !=  ""  { // Double-quoted string 
578+ 			result  =  append (result , match [2 ])
579+ 		} else  if  match [3 ] !=  ""  { // Unquoted segment 
580+ 			result  =  append (result , match [3 ])
569581		}
570582	}
571- 	return  - 1 , false 
583+ 
584+ 	return  result 
572585}
573586
574587// set value of the parameter passed from the annotations. 
575- func  setHelmValue (currentValues  * yaml.MapSlice , key  string , value  interface {}) error  {
576- 	// Check if the full key exists 
577- 	if  idx , found  :=  findHelmValuesKey (* currentValues , key ); found  {
578- 		(* currentValues )[idx ].Value  =  value 
579- 		return  nil 
588+ func  setHelmValue (file  * ast.File , keyPath , value  string ) error  {
589+ 	path  :=  splitKeys (keyPath )
590+ 	if  len (path ) ==  0  {
591+ 		return  fmt .Errorf ("empty key provided" )
580592	}
581593
582- 	var  err  error 
583- 	keys  :=  strings .Split (key , "." )
584- 	current  :=  currentValues 
585- 	var  parent  * yaml.MapSlice 
586- 	parentIdx  :=  - 1 
587- 
588- 	for  i , k  :=  range  keys  {
589- 		if  idx , found  :=  findHelmValuesKey (* current , k ); found  {
590- 			if  i  ==  len (keys )- 1  {
591- 				// If we're at the final key, set the value and return 
592- 				(* current )[idx ].Value  =  value 
593- 				return  nil 
594- 			} else  {
595- 				// Navigate deeper into the map 
596- 				if  nestedMap , ok  :=  (* current )[idx ].Value .(yaml.MapSlice ); ok  {
597- 					parent  =  current 
598- 					parentIdx  =  idx 
599- 					current  =  & nestedMap 
600- 				} else  {
601- 					return  fmt .Errorf ("unexpected type %T for key %s" , (* current )[idx ].Value , k )
594+ 	var  mapping  * ast.MappingNode 
595+ 	if  file .Docs [0 ].Body  ==  nil  {
596+ 		tk  :=  token .New ("$" , "$" , & token.Position {})
597+ 		mapping  =  ast .Mapping (tk , false )
598+ 		file .Docs [0 ].Body  =  mapping 
599+ 	} else  {
600+ 		mapping , _  =  file .Docs [0 ].Body .(* ast.MappingNode )
601+ 		if  mapping  ==  nil  {
602+ 			return  fmt .Errorf ("yaml is invalid" )
603+ 		}
604+ 	}
605+ 
606+ 	// Traverse the path 
607+ 	var  lastNode  * ast.MappingValueNode 
608+ 	for  index , key  :=  range  path  {
609+ 		found  :=  false 
610+ 		var  currentNode  * ast.MappingValueNode 
611+ 
612+ 		for  _ , v  :=  range  mapping .Values  {
613+ 			if  v .Key .GetToken ().Value  ==  key  {
614+ 				currentNode  =  v 
615+ 				if  index  ==  len (path )- 1  {
616+ 					lastNode  =  currentNode 
617+ 					found  =  true 
618+ 					break 
619+ 				}
620+ 				// Move deeper into the structure 
621+ 				if  nextMapping , ok  :=  v .Value .(* ast.MappingNode ); ok  {
622+ 					mapping  =  nextMapping 
623+ 					found  =  true 
624+ 					break 
602625				}
603626			}
604- 		} else  {
605- 			newCurrent  :=  yaml.MapSlice {}
606- 			var  newParent  yaml.MapSlice 
627+ 		}
607628
608- 			if  i  ==  len (keys )- 1  {
609- 				newParent  =  append (* current , yaml.MapItem {Key : k , Value : value })
610- 			} else  {
611- 				newParent  =  append (* current , yaml.MapItem {Key : k , Value : newCurrent })
612- 			}
629+ 		// If key does not exist, create it 
630+ 		if  ! found  {
631+ 			// Create a token with proper position (assuming default line/column) 
632+ 			keyToken  :=  token .New (key , key , & token.Position {Column : index * 2  +  1 })
633+ 			newKey  :=  ast .String (keyToken ) // Create key node 
634+ 			mappingToken  :=  token .New (key , key , & token.Position {})
635+ 			newMapping  :=  ast .Mapping (mappingToken , false ) // Create empty mapping 
613636
614- 			if  parent  ==  nil  {
615- 				* currentValues  =  newParent 
616- 			} else  {
617- 				// if parentIdx has not been set (parent element is also new), set it to the last element 
618- 				if  parentIdx  ==  - 1  {
619- 					parentIdx  =  len (* parent ) -  1 
620- 					if  parentIdx  <  0  {
621- 						parentIdx  =  0 
622- 					}
637+ 			if  currentNode  !=  nil  {
638+ 				comment  :=  currentNode .Value .GetComment ()
639+ 				currentNode .Value  =  newMapping 
640+ 				err  :=  currentNode .Value .SetComment (comment )
641+ 				if  err  !=  nil  {
642+ 					return  err 
623643				}
624- 				(* parent )[parentIdx ].Value  =  newParent 
644+ 				lastNode  =  currentNode 
645+ 			} else  {
646+ 				// Add the new mapping to the parent mapping 
647+ 				lastNode  =  ast .MappingValue (mappingToken , newKey , newMapping )
648+ 				mapping .Values  =  append (mapping .Values , lastNode )
625649			}
626- 
627- 			parent  =  & newParent 
628- 			current  =  & newCurrent 
629- 			parentIdx  =  - 1 
650+ 			mapping  =  newMapping 
630651		}
631652	}
632653
633- 	return  err 
654+ 	if  lastNode  ==  nil  {
655+ 		return  fmt .Errorf ("key not found" )
656+ 	}
657+ 
658+ 	var  valueToken  * token.Token 
659+ 	if  token .IsNeedQuoted (value ) {
660+ 		valueToken  =  token .SingleQuote (value , value , & token.Position {})
661+ 	} else  {
662+ 		valueToken  =  token .New (value , value , & token.Position {})
663+ 	}
664+ 	newValue  :=  ast .String (valueToken )
665+ 	comment  :=  lastNode .Value .GetComment ()
666+ 	if  comment  ==  nil  {
667+ 		comment  =  lastNode .Key .GetComment ()
668+ 	}
669+ 	lastNode .Value  =  newValue 
670+ 	err  :=  lastNode .Key .SetComment (nil )
671+ 	if  err  !=  nil  {
672+ 		return  err 
673+ 	}
674+ 	err  =  lastNode .Value .SetComment (comment )
675+ 	if  err  !=  nil  {
676+ 		return  err 
677+ 	}
678+ 
679+ 	return  nil 
634680}
635681
636682func  getWriteBackConfig (app  * v1alpha1.Application , kubeClient  * kube.ImageUpdaterKubernetesClient , argoClient  ArgoCD ) (* WriteBackConfig , error ) {
0 commit comments