From 13a8a58d75ee056d67c37f44a633c4cf252c4de2 Mon Sep 17 00:00:00 2001
From: CrazyMax <crazy-max@users.noreply.github.com>
Date: Thu, 27 Apr 2023 17:39:28 +0200
Subject: [PATCH] allow setting user agent products

Signed-off-by: CrazyMax <crazy-max@users.noreply.github.com>
---
 version/ua.go      | 49 +++++++++++++++++++++++++++++++++++++++++++++
 version/ua_test.go | 50 ++++++++++++++++++++++++++++++++++++++++++++++
 version/version.go | 28 --------------------------
 3 files changed, 99 insertions(+), 28 deletions(-)
 create mode 100644 version/ua.go
 create mode 100644 version/ua_test.go

diff --git a/version/ua.go b/version/ua.go
new file mode 100644
index 000000000000..01cfe67cd047
--- /dev/null
+++ b/version/ua.go
@@ -0,0 +1,49 @@
+package version
+
+import (
+	"fmt"
+	"regexp"
+	"strings"
+	"sync"
+)
+
+var (
+	reRelease *regexp.Regexp
+	reDev     *regexp.Regexp
+	reOnce    sync.Once
+	uapCbs    map[string]func() string
+)
+
+func UserAgent() string {
+	uaVersion := defaultVersion
+
+	reOnce.Do(func() {
+		reRelease = regexp.MustCompile(`^(v[0-9]+\.[0-9]+)\.[0-9]+$`)
+		reDev = regexp.MustCompile(`^(v[0-9]+\.[0-9]+)\.[0-9]+`)
+	})
+
+	if matches := reRelease.FindAllStringSubmatch(Version, 1); len(matches) > 0 {
+		uaVersion = matches[0][1]
+	} else if matches := reDev.FindAllStringSubmatch(Version, 1); len(matches) > 0 {
+		uaVersion = matches[0][1] + "-dev"
+	}
+
+	res := &strings.Builder{}
+	fmt.Fprintf(res, "buildkit/%s", uaVersion)
+	for pname, pver := range uapCbs {
+		fmt.Fprintf(res, " %s/%s", pname, pver())
+	}
+
+	return res.String()
+}
+
+// SetUserAgentProduct sets a callback to get the version of a product to be
+// included in the User-Agent header. The callback is called every time the
+// User-Agent header is generated. Caller must ensure that the callback is
+// cached if it is expensive to compute.
+func SetUserAgentProduct(name string, cb func() (version string)) {
+	if uapCbs == nil {
+		uapCbs = make(map[string]func() string)
+	}
+	uapCbs[name] = cb
+}
diff --git a/version/ua_test.go b/version/ua_test.go
new file mode 100644
index 000000000000..ce0c15c67c93
--- /dev/null
+++ b/version/ua_test.go
@@ -0,0 +1,50 @@
+package version
+
+import "testing"
+
+func TestUserAgent(t *testing.T) {
+	cases := []struct {
+		name     string
+		version  string
+		products map[string]string
+		want     string
+	}{
+		{
+			name:    "dev",
+			version: defaultVersion,
+			want:    "buildkit/v0.0-dev",
+		},
+		{
+			name:    "unknown",
+			version: "0.0.0+unknown",
+			want:    "buildkit/v0.0.0+unknown",
+		},
+		{
+			name:    "release",
+			version: "v0.11.6",
+			want:    "buildkit/v0.11",
+		},
+		{
+			name:    "product",
+			version: "v0.11.6",
+			products: map[string]string{
+				"moby": "v24.0",
+			},
+			want: "buildkit/v0.11 moby/v24.0",
+		},
+	}
+	for _, tt := range cases {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			Version = tt.version
+			for pname, pver := range tt.products {
+				SetUserAgentProduct(pname, func() string {
+					return pver
+				})
+			}
+			if g, w := UserAgent(), tt.want; g != w {
+				t.Fatalf("got: %q\nwant: %q", g, w)
+			}
+		})
+	}
+}
diff --git a/version/version.go b/version/version.go
index 4dfe4618ba46..9cddea63c022 100644
--- a/version/version.go
+++ b/version/version.go
@@ -17,11 +17,6 @@
 
 package version
 
-import (
-	"regexp"
-	"sync"
-)
-
 const (
 	defaultVersion = "v0.0.0+unknown"
 )
@@ -37,26 +32,3 @@ var (
 	// the program at linking time.
 	Revision = ""
 )
-
-var (
-	reRelease *regexp.Regexp
-	reDev     *regexp.Regexp
-	reOnce    sync.Once
-)
-
-func UserAgent() string {
-	uaVersion := defaultVersion
-
-	reOnce.Do(func() {
-		reRelease = regexp.MustCompile(`^(v[0-9]+\.[0-9]+)\.[0-9]+$`)
-		reDev = regexp.MustCompile(`^(v[0-9]+\.[0-9]+)\.[0-9]+`)
-	})
-
-	if matches := reRelease.FindAllStringSubmatch(Version, 1); len(matches) > 0 {
-		uaVersion = matches[0][1]
-	} else if matches := reDev.FindAllStringSubmatch(Version, 1); len(matches) > 0 {
-		uaVersion = matches[0][1] + "-dev"
-	}
-
-	return "buildkit/" + uaVersion
-}