From 77a312b5f029bea8da38db77204139019bc0b35f Mon Sep 17 00:00:00 2001
From: Tim Ross <tim.ross@goteleport.com>
Date: Tue, 27 Feb 2024 11:21:03 -0500
Subject: [PATCH] Prevent overflow in sort functions provided to
 slices.SortFunc

Replaces subtraction based comparisions with calls to cmp.Compare
for all ordered types to avoid overflows.
---
 api/types/instance.go               | 10 +++++++++-
 lib/events/azsessions/azsessions.go | 24 ++++++++++++++++++++++--
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/api/types/instance.go b/api/types/instance.go
index 0e344aaf65326..01730ee7c3914 100644
--- a/api/types/instance.go
+++ b/api/types/instance.go
@@ -295,7 +295,15 @@ func (i *InstanceV1) AppendControlLog(entries ...InstanceControlLogEntry) {
 		i.Spec.ControlLog[idx].Time = entry.Time.UTC()
 	}
 	slices.SortFunc(i.Spec.ControlLog, func(a, b InstanceControlLogEntry) int {
-		return int(a.Time.UnixNano() - b.Time.UnixNano())
+		an, bn := a.Time.UnixNano(), b.Time.UnixNano()
+
+		switch {
+		case an < bn:
+			return -1
+		case bn > an:
+			return 1
+		}
+		return 0
 	})
 }
 
diff --git a/lib/events/azsessions/azsessions.go b/lib/events/azsessions/azsessions.go
index 457e1e730d3dd..ae9303623e3b7 100644
--- a/lib/events/azsessions/azsessions.go
+++ b/lib/events/azsessions/azsessions.go
@@ -320,7 +320,17 @@ func (h *Handler) CompleteUpload(ctx context.Context, upload events.StreamUpload
 	// cleaned up before a new attempt
 
 	parts = slices.Clone(parts)
-	slices.SortFunc(parts, func(a, b events.StreamPart) int { return int(a.Number - b.Number) })
+	slices.SortFunc(parts, func(a, b events.StreamPart) int {
+		if a.Number < b.Number {
+			return -1
+		}
+
+		if a.Number > b.Number {
+			return 1
+		}
+
+		return 0
+	})
 
 	partURLs := make([]string, 0, len(parts))
 	for _, part := range parts {
@@ -492,7 +502,17 @@ func (h *Handler) ListParts(ctx context.Context, upload events.StreamUpload) ([]
 		}
 	}
 
-	slices.SortFunc(parts, func(a, b events.StreamPart) int { return int(a.Number - b.Number) })
+	slices.SortFunc(parts, func(a, b events.StreamPart) int {
+		if a.Number < b.Number {
+			return -1
+		}
+
+		if a.Number > b.Number {
+			return 1
+		}
+
+		return 0
+	})
 
 	return parts, nil
 }