44package nvfractions
55
66import (
7+ "context"
78 "fmt"
89 "strconv"
10+ "strings"
911
1012 v1 "k8s.io/api/core/v1"
13+ "sigs.k8s.io/controller-runtime/pkg/webhook/admission"
1114
1215 "github.com/kai-scheduler/KAI-scheduler/pkg/binder/common"
1316 "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
@@ -19,18 +22,25 @@ import (
1922// gpu-fraction-container-name, and NvFractions annotations) and normalizes
2023// legacy memory annotations to the NvFractions form, without using the shared
2124// GPU configmap.
22- type NvFractions struct {}
25+ type NvFractions struct {
26+ binderServiceAccountUsername string
27+ }
2328
24- func New () * NvFractions {
25- return & NvFractions {}
29+ func New (binderServiceAccountUsername string ) * NvFractions {
30+ return & NvFractions {
31+ binderServiceAccountUsername : binderServiceAccountUsername ,
32+ }
2633}
2734
2835func (p * NvFractions ) Name () string {
2936 return "nvfractions"
3037}
3138
32- func (p * NvFractions ) Validate (pod * v1.Pod ) error {
33- return resources .ValidateGPUFractionRequest (pod )
39+ func (p * NvFractions ) Validate (ctx context.Context , oldPod , pod * v1.Pod ) error {
40+ if err := p .validateDeviceAnnotation (ctx , oldPod , pod ); err != nil {
41+ return err
42+ }
43+ return resources .ValidateGPUFractionRequest (stripNvFractionsDeviceAnnotations (pod ))
3444}
3545
3646func (p * NvFractions ) Mutate (pod * v1.Pod ) error {
@@ -68,3 +78,88 @@ func adjustFractionalMemoryAnnotations(pod *v1.Pod, containerName string) error
6878 pod .Annotations [resources .CalcGpuFractionAnnotationForContainer (containerName )] = memoryQuantity .String ()
6979 return nil
7080}
81+
82+ // validateDeviceAnnotation makes sure that only the binder service account can modify the nvfractions device annotations.
83+ // This is done to prevent malicious actors from modifying the nvfractions device annotations to gain access to the GPU.
84+ func (p * NvFractions ) validateDeviceAnnotation (ctx context.Context , oldPod , pod * v1.Pod ) error {
85+ hasNvFractionsDeviceAnnotationChange := false
86+ for annotationKey := range pod .Annotations {
87+ if ! isNvFractionsDeviceAnnotation (annotationKey ) {
88+ continue
89+ }
90+ if err := validateDeviceAnnotationContainerExists (pod , annotationKey ); err != nil {
91+ return err
92+ }
93+
94+ if oldPod != nil {
95+ hasNvFractionsDeviceAnnotationChange = isAnnotationChanged (annotationKey , oldPod , pod )
96+ } else {
97+ hasNvFractionsDeviceAnnotationChange = true
98+ }
99+
100+ if hasNvFractionsDeviceAnnotationChange {
101+ break
102+ }
103+ }
104+
105+ if ! hasNvFractionsDeviceAnnotationChange {
106+ return nil
107+ }
108+
109+ if len (p .binderServiceAccountUsername ) == 0 {
110+ return fmt .Errorf ("binder service account username is not configured, cannot validate nvfractions device annotations change" )
111+ }
112+
113+ request , err := admission .RequestFromContext (ctx )
114+ if err != nil {
115+ return fmt .Errorf ("failed to extract admission request: %w" , err )
116+ }
117+ if request .UserInfo .Username != p .binderServiceAccountUsername {
118+ return fmt .Errorf ("%s annotations may only be modified by %s" ,
119+ constants .NvFractionsVisibleDevicesSuffix , p .binderServiceAccountUsername )
120+ }
121+
122+ return nil
123+ }
124+
125+ func isAnnotationChanged (annotationKey string , oldPod * v1.Pod , pod * v1.Pod ) bool {
126+ oldAnnotationValue , oldHasNvFractionsDeviceAnnotation := oldPod .Annotations [annotationKey ]
127+ if ! oldHasNvFractionsDeviceAnnotation {
128+ return true
129+ }
130+ newAnnotationValue := pod .Annotations [annotationKey ]
131+ if oldAnnotationValue != newAnnotationValue {
132+ return true
133+ }
134+ return false
135+ }
136+
137+ func stripNvFractionsDeviceAnnotations (pod * v1.Pod ) * v1.Pod {
138+ if pod == nil || len (pod .Annotations ) == 0 {
139+ return pod
140+ }
141+
142+ podCopy := pod .DeepCopy ()
143+ for key := range podCopy .Annotations {
144+ if isNvFractionsDeviceAnnotation (key ) {
145+ delete (podCopy .Annotations , key )
146+ }
147+ }
148+ return podCopy
149+ }
150+
151+ func isNvFractionsDeviceAnnotation (annotationKey string ) bool {
152+ return strings .HasPrefix (annotationKey , constants .NvFractionsAnnotationPrefix ) &&
153+ strings .HasSuffix (annotationKey , constants .NvFractionsVisibleDevicesSuffix )
154+ }
155+
156+ func validateDeviceAnnotationContainerExists (pod * v1.Pod , annotationKey string ) error {
157+ containerName := strings .TrimPrefix (annotationKey , constants .NvFractionsAnnotationPrefix )
158+ containerName = strings .TrimSuffix (containerName , constants .NvFractionsVisibleDevicesSuffix )
159+ for _ , container := range pod .Spec .Containers {
160+ if container .Name == containerName {
161+ return nil
162+ }
163+ }
164+ return fmt .Errorf ("container %s not found in pod spec, but a fractional annotation referencing it was found" , containerName )
165+ }
0 commit comments