Move more of version check to Rego (#1513)

Just something little I wanted to try :)

Signed-off-by: Anders Eknert <anders@styra.com>

Author: anderseknert

cmd/lint.go

@@ -216,11 +216,11 @@ func lint(args []string, params *lintCommandParams) (report.Report, error) {
 		}
 	}
 
-	var regalDir *os.File
-
-	var customRulesDir string
-
-	var configSearchPath string
+	var (
+		regalDir         *os.File
+		customRulesDir   string
+		configSearchPath string
+	)
 
 	m := metrics.New()
 	if params.metrics {

internal/update/update.go

@@ -3,10 +3,9 @@ package update
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
-	"net/http"
-	"net/url"
 	"os"
 	"path/filepath"
 	"strings"
@@ -26,35 +25,35 @@ var updateModule string
 const CheckVersionDisableEnvVar = "REGAL_DISABLE_VERSION_CHECK"
 
 type Options struct {
-	CurrentVersion string
-	CurrentTime    time.Time
-
-	StateDir string
-
+	CurrentTime       time.Time
+	CurrentVersion    string
+	StateDir          string
 	ReleaseServerHost string
 	ReleaseServerPath string
-
-	CTAURLPrefix string
-
-	Debug bool
+	CTAURLPrefix      string
+	Debug             bool
 }
 
 type latestVersionFileContents struct {
 	CheckedAt     time.Time `json:"checked_at"`
 	LatestVersion string    `json:"latest_version"`
 }
 
+type decision struct {
+	NeedsUpdate   bool   `json:"needs_update"`
+	LatestVersion string `json:"latest_version"`
+	CTA           string `json:"cta"`
+}
+
+var query = ast.MustParseBody("result := data.update.check")
+
 func CheckAndWarn(opts Options, w io.Writer) {
-	// this is a shortcut heuristic to avoid and version checking
-	// when in dev/test etc.
+	// this is a shortcut heuristic to avoid version checking when in dev/test etc.
 	if !strings.HasPrefix(opts.CurrentVersion, "v") {
 		return
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
-
-	latestVersion, err := getLatestVersion(ctx, opts)
+	latestVersion, err := getLatestCachedVersion(opts)
 	if err != nil {
 		if opts.Debug {
 			w.Write([]byte(err.Error()))
@@ -65,10 +64,13 @@ func CheckAndWarn(opts Options, w io.Writer) {
 
 	regoArgs := []func(*rego.Rego){
 		rego.Module("update.rego", updateModule),
-		rego.Query(`data.update.needs_update`),
+		rego.ParsedQuery(query),
 		rego.ParsedInput(ast.NewObject(
 			ast.Item(ast.StringTerm("current_version"), ast.StringTerm(opts.CurrentVersion)),
 			ast.Item(ast.StringTerm("latest_version"), ast.StringTerm(latestVersion)),
+			ast.Item(ast.StringTerm("cta_url_prefix"), ast.StringTerm(opts.CTAURLPrefix)),
+			ast.Item(ast.StringTerm("release_server_host"), ast.StringTerm(opts.ReleaseServerHost)),
+			ast.Item(ast.StringTerm("release_server_path"), ast.StringTerm(opts.ReleaseServerPath)),
 		)),
 	}
 
@@ -81,45 +83,54 @@ func CheckAndWarn(opts Options, w io.Writer) {
 		return
 	}
 
-	if !rs.Allowed() {
+	result, err := resultSetToDecision(rs)
+	if err != nil {
 		if opts.Debug {
-			w.Write([]byte("Regal is up to date"))
+			w.Write([]byte(err.Error()))
+		}
+
+		return
+	}
+
+	if result.NeedsUpdate {
+		if err = saveLatestCachedVersion(opts, result.LatestVersion); err != nil && opts.Debug {
+			w.Write([]byte(err.Error()))
 		}
 
+		w.Write([]byte(result.CTA))
+
 		return
 	}
 
-	ctaURLPrefix := "https://github.com/StyraInc/regal/releases/tag/"
-	if opts.CTAURLPrefix != "" {
-		ctaURLPrefix = opts.CTAURLPrefix
+	if opts.Debug {
+		w.Write([]byte("Regal is up to date"))
 	}
+}
 
-	ctaURL := ctaURLPrefix + latestVersion
+func resultSetToDecision(rs rego.ResultSet) (decision, error) {
+	if len(rs) == 0 || rs[0].Bindings["result"] == nil {
+		return decision{}, errors.New("no result set")
+	}
 
-	tmpl := `A new version of Regal is available (%s). You are running %s.
-See %s for the latest release.
-`
+	var result decision
+	if err := encoding.JSONRoundTrip(rs[0].Bindings["result"], &result); err != nil {
+		return decision{}, fmt.Errorf("failed to decode result set: %w", err)
+	}
 
-	fmt.Fprintf(w, tmpl, latestVersion, opts.CurrentVersion, ctaURL)
+	return result, nil
 }
 
-func getLatestVersion(ctx context.Context, opts Options) (string, error) {
+func getLatestCachedVersion(opts Options) (string, error) {
 	if opts.StateDir != "" {
 		// first, attempt to get the file from previous invocations to save on remote calls
 		latestVersionFilePath := filepath.Join(opts.StateDir, "latest_version.json")
 
-		_, err := os.Stat(latestVersionFilePath)
-		if err == nil {
-			var preExistingState latestVersionFileContents
-
-			file, err := os.Open(latestVersionFilePath)
-			if err != nil {
-				return "", fmt.Errorf("failed to open file: %w", err)
-			}
+		if file, err := os.Open(latestVersionFilePath); err == nil {
+			defer file.Close()
 
-			json := encoding.JSON()
+			var preExistingState latestVersionFileContents
 
-			if err := json.NewDecoder(file).Decode(&preExistingState); err != nil {
+			if err := encoding.JSON().NewDecoder(file).Decode(&preExistingState); err != nil {
 				return "", fmt.Errorf("failed to decode existing version state file: %w", err)
 			}
 
@@ -129,60 +140,22 @@ func getLatestVersion(ctx context.Context, opts Options) (string, error) {
 		}
 	}
 
-	client := http.Client{}
+	return "", nil
+}
 
-	releaseServerHost := "https://api.github.com"
-	if opts.ReleaseServerHost != "" {
-		releaseServerHost = strings.TrimSuffix(opts.ReleaseServerHost, "/")
+func saveLatestCachedVersion(opts Options, latestVersion string) error {
+	if opts.StateDir != "" {
+		content := latestVersionFileContents{LatestVersion: latestVersion, CheckedAt: opts.CurrentTime}
 
-		if !strings.HasPrefix(releaseServerHost, "http") {
-			releaseServerHost = "https://" + releaseServerHost
+		bs, err := encoding.JSON().MarshalIndent(content, "", "  ")
+		if err != nil {
+			return fmt.Errorf("failed to marshal state file: %w", err)
 		}
-	}
-
-	releaseServerURL, err := url.Parse(releaseServerHost)
-	if err != nil {
-		return "", fmt.Errorf("failed to parse release server URL: %w", err)
-	}
-
-	releaseServerPath := "/repos/styrainc/regal/releases/latest"
-	if opts.ReleaseServerPath != "" {
-		releaseServerPath = opts.ReleaseServerPath
-	}
-
-	releaseServerURL.Path = releaseServerPath
-
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, releaseServerURL.String(), nil)
-	if err != nil {
-		return "", fmt.Errorf("failed to create request: %w", err)
-	}
 
-	resp, err := client.Do(req)
-	if err != nil {
-		return "", fmt.Errorf("failed to make request: %w", err)
-	}
-	defer resp.Body.Close()
-
-	var responseData struct {
-		TagName string `json:"tag_name"`
-	}
-
-	json := encoding.JSON()
-	if err = json.NewDecoder(resp.Body).Decode(&responseData); err != nil {
-		return "", fmt.Errorf("failed to decode response: %w", err)
-	}
-
-	stateBs, err := json.MarshalIndent(latestVersionFileContents{
-		LatestVersion: responseData.TagName,
-		CheckedAt:     opts.CurrentTime,
-	}, "", "  ")
-	if err != nil {
-		return "", fmt.Errorf("failed to marshal state file: %w", err)
-	}
-
-	if err = os.WriteFile(opts.StateDir+"/latest_version.json", stateBs, 0o600); err != nil {
-		return "", fmt.Errorf("failed to write state file: %w", err)
+		if err = os.WriteFile(opts.StateDir+"/latest_version.json", bs, 0o600); err != nil {
+			return fmt.Errorf("failed to write state file: %w", err)
+		}
 	}
 
-	return responseData.TagName, nil
+	return nil
 }

internal/update/update.rego

@@ -2,14 +2,52 @@
 # description: utility module to help determine if an update of Regal should be recommended
 package update
 
-default needs_update := false
+default check["needs_update"] := false
 
 # METADATA
 # description: true if current version is behind latest version
 # scope: document
-needs_update if {
+check["needs_update"] if {
 	current_version := trim(input.current_version, "v")
 
 	semver.is_valid(current_version)
-	semver.compare(current_version, trim(input.latest_version, "v")) == -1
+	semver.compare(current_version, trim(check.latest_version, "v")) == -1
 }
+
+# METADATA
+# description: the latest version, as determined by the release server
+# scope: document
+check["latest_version"] := input.latest_version if input.latest_version != ""
+
+check["latest_version"] := response.body.tag_name if {
+	input.latest_version == ""
+
+	response := http.send({"url": _release_server_url, "method": "GET"})
+}
+
+# METADATA
+# description: the CTA message, if an update is available
+# scope: document
+check["cta"] := sprintf(
+	"A new version of Regal is available (%s). You are running %s.\nSee %s for the latest release.\n",
+	[check.latest_version, input.current_version, concat("", [_cta_url_prefix, check.latest_version])],
+) if {
+	check.needs_update
+}
+
+_release_server_url := concat("", [_release_server_host, _release_server_path])
+
+default _release_server_host := "https://api.github.com"
+
+_release_server_host := _ensure_http(trim_suffix(input.release_server_host, "/")) if input.release_server_host != ""
+
+default _release_server_path := "/repos/styrainc/regal/releases/latest"
+
+_release_server_path := input.release_server_path if input.release_server_path != ""
+
+default _cta_url_prefix := "https://github.com/styrainc/regal/releases/tag/"
+
+_cta_url_prefix := input.cta_url_prefix if input.cta_url_prefix != ""
+
+_ensure_http(url) := url if startswith(url, "http")
+_ensure_http(url) := concat("", ["https://", url]) if not startswith(url, "http")

internal/update/update_test.go

@@ -16,6 +16,8 @@ func TestCheckAndWarn(t *testing.T) {
 
 	localReleasesServer := httptest.NewServer(
 		http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			w.Header().Set("Content-Type", "application/json")
+
 			if _, err := w.Write([]byte(`{"tag_name": "v0.2.0"}`)); err != nil {
 				t.Fatal(err)
 			}