From e1f3686f4bea6a3467b2ba7f3e78585cd52db0ce Mon Sep 17 00:00:00 2001
From: Kinara Shah <kinara@rancher.com>
Date: Tue, 20 Aug 2024 14:58:30 -0700
Subject: [PATCH] Fixes for generating data.json

- remove rke2/k3s keys so data only has required info for rke1
- decide kdm url based on the tag version
---
 .github/workflows/workflow.yaml |  1 +
 codegen/codegen.go              | 47 +++++++++++++++++++----
 codegen/codegen_test.go         | 66 +++++++++++++++++++++++++++++++++
 3 files changed, 107 insertions(+), 7 deletions(-)
 create mode 100644 codegen/codegen_test.go

diff --git a/.github/workflows/workflow.yaml b/.github/workflows/workflow.yaml
index 00ad41bdc..0acf834e6 100644
--- a/.github/workflows/workflow.yaml
+++ b/.github/workflows/workflow.yaml
@@ -43,6 +43,7 @@ jobs:
           ls -lR build/bin
         env:
           CROSS: 1
+          TAG: ${{ startsWith(github.ref, 'refs/tags/') && github.ref_name || '' }}
       - name: Generate sha256 file for release
         if: github.event_name == 'push' && github.ref_type == 'tag'
         run: |
diff --git a/codegen/codegen.go b/codegen/codegen.go
index 07db3e31f..b5176cbde 100644
--- a/codegen/codegen.go
+++ b/codegen/codegen.go
@@ -1,25 +1,27 @@
 package main
 
 import (
+	"encoding/json"
 	"fmt"
 	"io"
 	"net/http"
 	"os"
+	"strings"
 
+	"github.com/coreos/go-semver/semver"
 	"github.com/rancher/rke/metadata"
 )
 
 const (
-	defaultURL = "https://releases.rancher.com/kontainer-driver-metadata/dev-v2.7/data.json"
-	dataFile   = "data/data.json"
+	defaultDevURL     = "https://releases.rancher.com/kontainer-driver-metadata/dev-v2.7/data.json"
+	defaultReleaseURL = "https://releases.rancher.com/kontainer-driver-metadata/release-v2.7/data.json"
+	dataFile          = "data/data.json"
 )
 
-// Codegen fetch data.json from https://releases.rancher.com/kontainer-driver-metadata/release-v2.7/data.json and generates bindata
+// Codegen fetch data.json from defaultURL or defaultReleaseURL and generates bindata
 func main() {
-	u := os.Getenv(metadata.RancherMetadataURLEnv)
-	if u == "" {
-		u = defaultURL
-	}
+	u := getURL()
+	fmt.Printf("Reading data from %s \n", u)
 	data, err := http.Get(u)
 	if err != nil {
 		panic(fmt.Errorf("failed to fetch data.json from kontainer-driver-metadata repository"))
@@ -31,8 +33,39 @@ func main() {
 		panic(err)
 	}
 
+	var jsonData map[string]interface{}
+	err = json.Unmarshal(b, &jsonData)
+	if err != nil {
+		panic(fmt.Errorf("failed to parse JSON: %v", err))
+	}
+	// rke doesn't need info about rke2 and k3s versions
+	delete(jsonData, "rke2")
+	delete(jsonData, "k3s")
+
+	b, err = json.Marshal(jsonData)
+	if err != nil {
+		panic(fmt.Errorf("failed to marshal json data: %v", err))
+	}
+
 	fmt.Println("Writing data")
 	if err := os.WriteFile(dataFile, b, 0755); err != nil {
 		return
 	}
 }
+
+func getURL() string {
+	u := os.Getenv(metadata.RancherMetadataURLEnv)
+	if u == "" {
+		u = defaultDevURL
+		tag := os.Getenv("TAG")
+		if strings.HasPrefix(tag, "v") {
+			tag = tag[1:]
+		}
+		if v, err := semver.NewVersion(tag); err == nil {
+			if v.PreRelease == "" && v.String() != "" {
+				u = defaultReleaseURL
+			}
+		}
+	}
+	return u
+}
diff --git a/codegen/codegen_test.go b/codegen/codegen_test.go
new file mode 100644
index 000000000..0d43ddb88
--- /dev/null
+++ b/codegen/codegen_test.go
@@ -0,0 +1,66 @@
+package main
+
+import (
+	"os"
+	"testing"
+
+	"github.com/rancher/rke/metadata"
+)
+
+func TestGetURL(t *testing.T) {
+	tests := []struct {
+		name        string
+		envVar      string
+		tag         string
+		expectedURL string
+	}{
+		{
+			name:        "No Metadata URL and TAG is release version",
+			envVar:      "",
+			tag:         "v1.0.0",
+			expectedURL: defaultReleaseURL,
+		},
+		{
+			name:        "No Metadata URL and TAG is pre-release version",
+			envVar:      "",
+			tag:         "v1.0.0-alpha",
+			expectedURL: defaultDevURL,
+		},
+		{
+			name:        "Metadata URL set",
+			envVar:      "https://example.com",
+			tag:         "v1.0.0",
+			expectedURL: "https://example.com",
+		},
+		{
+			name:        "Invalid TAG",
+			envVar:      "",
+			tag:         "invalid-tag",
+			expectedURL: defaultDevURL,
+		},
+		{
+			name:        "No TAG",
+			envVar:      "",
+			tag:         "",
+			expectedURL: defaultDevURL,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Set the environment variables
+			os.Setenv(metadata.RancherMetadataURLEnv, tt.envVar)
+			os.Setenv("TAG", tt.tag)
+			defer func() {
+				os.Unsetenv(metadata.RancherMetadataURLEnv)
+				os.Unsetenv("TAG")
+			}()
+
+			result := getURL()
+
+			if result != tt.expectedURL {
+				t.Errorf("expected %s, got %s", tt.expectedURL, result)
+			}
+		})
+	}
+}