feat: set User-Agent HTTP header on requests (#249)

semaphoreci · Aug 9, 2024 · 24f7f04 · 24f7f04
1 parent 58209f7
commit 24f7f04
Show file tree

Hide file tree

Showing 17 changed files with 116 additions and 39 deletions.
diff --git a/main.go b/main.go
@@ -29,6 +29,7 @@ import (
 )
 
 var VERSION = "dev"
+var HTTPUserAgent = fmt.Sprintf("SemaphoreAgent/%s", VERSION)
 
 func main() {
 	logfile := OpenLogfile()
@@ -220,6 +221,7 @@ func RunListener(httpClient *http.Client, logfile io.Writer) {
 		FailOnPreJobHookError:            viper.GetBool(config.FailOnPreJobHookError),
 		SourcePreJobHook:                 viper.GetBool(config.SourcePreJobHook),
 		AgentVersion:                     VERSION,
+		UserAgent:                        HTTPUserAgent,
 		ExitOnShutdown:                   true,
 		KubernetesExecutor:               viper.GetBool(config.KubernetesExecutor),
 		KubernetesPodSpec:                viper.GetString(config.KubernetesPodSpec),
@@ -398,6 +400,7 @@ func RunServer(httpClient *http.Client, logfile io.Writer) {
 	server.NewServer(server.ServerConfig{
 		Host:                  *host,
 		Port:                  *port,
+		UserAgent:             HTTPUserAgent,
 		TLSCertPath:           *tlsCertPath,
 		TLSKeyPath:            *tlsKeyPath,
 		Version:               VERSION,

diff --git a/pkg/eventlogger/default.go b/pkg/eventlogger/default.go
@@ -13,12 +13,22 @@ import (
 const LoggerMethodPull = "pull"
 const LoggerMethodPush = "push"
 
-func CreateLogger(request *api.JobRequest, refreshTokenFn func() (string, error)) (*Logger, error) {
-	switch request.Logger.Method {
+type LoggerOptions struct {
+	Request        *api.JobRequest
+	RefreshTokenFn func() (string, error)
+	UserAgent      string
+}
+
+func CreateLogger(options LoggerOptions) (*Logger, error) {
+	if options.Request == nil {
+		return nil, fmt.Errorf("request is required")
+	}
+
+	switch options.Request.Logger.Method {
 	case LoggerMethodPull:
-		return Default(request)
+		return Default(options.Request)
 	case LoggerMethodPush:
-		return DefaultHTTP(request, refreshTokenFn)
+		return DefaultHTTP(options)
 	default:
 		return nil, fmt.Errorf("unknown logger type")
 	}
@@ -50,19 +60,21 @@ func Default(request *api.JobRequest) (*Logger, error) {
 	return logger, nil
 }
 
-func DefaultHTTP(request *api.JobRequest, refreshTokenFn func() (string, error)) (*Logger, error) {
+func DefaultHTTP(options LoggerOptions) (*Logger, error) {
+	request := options.Request
 	if request.Logger.URL == "" {
 		return nil, errors.New("HTTP logger needs a URL")
 	}
 
-	if refreshTokenFn == nil {
+	if options.RefreshTokenFn == nil {
 		return nil, errors.New("HTTP logger needs a refresh token function")
 	}
 
 	backend, err := NewHTTPBackend(HTTPBackendConfig{
 		URL:                   request.Logger.URL,
 		Token:                 request.Logger.Token,
-		RefreshTokenFn:        refreshTokenFn,
+		RefreshTokenFn:        options.RefreshTokenFn,
+		UserAgent:             options.UserAgent,
 		LinesPerRequest:       MaxLinesPerRequest,
 		FlushTimeoutInSeconds: DefaultFlushTimeoutInSeconds,
 	})

diff --git a/pkg/eventlogger/httpbackend.go b/pkg/eventlogger/httpbackend.go
@@ -34,6 +34,7 @@ type HTTPBackend struct {
 
 type HTTPBackendConfig struct {
 	URL                   string
+	UserAgent             string
 	Token                 string
 	LinesPerRequest       int
 	FlushTimeoutInSeconds int
@@ -183,6 +184,7 @@ func (l *HTTPBackend) newRequest() error {
 
 	request.Header.Set("Content-Type", "text/plain")
 	request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", l.config.Token))
+	request.Header.Set("User-Agent", l.config.UserAgent)
 	response, err := l.client.Do(request)
 	if err != nil {
 		return err

diff --git a/pkg/eventlogger/httpbackend_test.go b/pkg/eventlogger/httpbackend_test.go
@@ -1,6 +1,7 @@
 package eventlogger
 
 import (
+	"fmt"
 	"testing"
 	"time"
 
@@ -93,6 +94,7 @@ func Test__LogsArePushedToHTTPEndpoint(t *testing.T) {
 		RefreshTokenFn:        func() (string, error) { return "", nil },
 		LinesPerRequest:       20,
 		FlushTimeoutInSeconds: DefaultFlushTimeoutInSeconds,
+		UserAgent:             fmt.Sprintf("SemaphoreAgent/%s", testsupport.AgentVersionExpected),
 	})
 
 	assert.Nil(t, err)
@@ -132,6 +134,7 @@ func Test__RequestsAreCappedAtLinesPerRequest(t *testing.T) {
 		RefreshTokenFn:        func() (string, error) { return "", nil },
 		LinesPerRequest:       2,
 		FlushTimeoutInSeconds: DefaultFlushTimeoutInSeconds,
+		UserAgent:             fmt.Sprintf("SemaphoreAgent/%s", testsupport.AgentVersionExpected),
 	})
 
 	assert.Nil(t, err)
@@ -183,6 +186,7 @@ func Test__FlushingGivesUpAfterTimeout(t *testing.T) {
 		RefreshTokenFn:        func() (string, error) { return "", nil },
 		LinesPerRequest:       2,
 		FlushTimeoutInSeconds: 10,
+		UserAgent:             fmt.Sprintf("SemaphoreAgent/%s", testsupport.AgentVersionExpected),
 	})
 
 	assert.Nil(t, err)
@@ -217,6 +221,7 @@ func Test__ExecutesOnCloseCallback(t *testing.T) {
 		RefreshTokenFn:        func() (string, error) { return "", nil },
 		LinesPerRequest:       10,
 		FlushTimeoutInSeconds: 10,
+		UserAgent:             fmt.Sprintf("SemaphoreAgent/%s", testsupport.AgentVersionExpected),
 	})
 
 	assert.Nil(t, err)
@@ -246,6 +251,7 @@ func Test__TokenIsRefreshed(t *testing.T) {
 		Token:                 testsupport.ExpiredLogToken,
 		LinesPerRequest:       20,
 		FlushTimeoutInSeconds: DefaultFlushTimeoutInSeconds,
+		UserAgent:             fmt.Sprintf("SemaphoreAgent/%s", testsupport.AgentVersionExpected),
 		RefreshTokenFn: func() (string, error) {
 			tokenWasRefreshed = true
 			return "some-new-and-shiny-valid-token", nil

diff --git a/pkg/jobs/job.go b/pkg/jobs/job.go
@@ -37,6 +37,7 @@ type Job struct {
 	Stopped        bool
 	Finished       bool
 	UploadJobLogs  string
+	UserAgent      string
 }
 
 type JobOptions struct {
@@ -54,6 +55,7 @@ type JobOptions struct {
 	KubernetesImageValidator         *kubernetes.ImageValidator
 	UploadJobLogs                    string
 	RefreshTokenFn                   func() (string, error)
+	UserAgent                        string
 }
 
 func NewJob(request *api.JobRequest, client *http.Client) (*Job, error) {
@@ -90,7 +92,12 @@ func NewJobWithOptions(options *JobOptions) (*Job, error) {
 	if options.Logger != nil {
 		job.Logger = options.Logger
 	} else {
-		l, err := eventlogger.CreateLogger(options.Request, options.RefreshTokenFn)
+		l, err := eventlogger.CreateLogger(eventlogger.LoggerOptions{
+			Request:        options.Request,
+			RefreshTokenFn: options.RefreshTokenFn,
+			UserAgent:      options.UserAgent,
+		})
+
 		if err != nil {
 			return nil, err
 		}
@@ -643,6 +650,7 @@ func (job *Job) SendCallback(url string, payload string) error {
 		request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", job.Request.Callbacks.Token))
 	}
 
+	request.Header.Set("User-Agent", job.UserAgent)
 	response, err := job.Client.Do(request)
 	if err != nil {
 		return err

diff --git a/pkg/listener/job_processor.go b/pkg/listener/job_processor.go
@@ -26,6 +26,7 @@ func StartJobProcessor(httpClient *http.Client, apiClient *selfhostedapi.API, co
 	p := &JobProcessor{
 		HTTPClient:                       httpClient,
 		APIClient:                        apiClient,
+		UserAgent:                        config.UserAgent,
 		LastSuccessfulSync:               time.Now(),
 		forceSyncCh:                      make(chan bool),
 		State:                            selfhostedapi.AgentStateWaitingForJobs,
@@ -84,6 +85,7 @@ type JobProcessor struct {
 	FileInjections                   []config.FileInjection
 	FailOnMissingFiles               bool
 	UploadJobLogs                    string
+	UserAgent                        string
 	FailOnPreJobHookError            bool
 	SourcePreJobHook                 bool
 	ExitOnShutdown                   bool
@@ -211,6 +213,7 @@ func (p *JobProcessor) RunJob(jobID string) {
 		KubernetesLabels:                 p.KubernetesLabels,
 		KubernetesImageValidator:         p.KubernetesImageValidator,
 		UploadJobLogs:                    p.UploadJobLogs,
+		UserAgent:                        p.UserAgent,
 		RefreshTokenFn: func() (string, error) {
 			return p.APIClient.RefreshToken()
 		},

diff --git a/pkg/listener/listener.go b/pkg/listener/listener.go
@@ -44,6 +44,7 @@ type Config struct {
 	SourcePreJobHook                 bool
 	ExitOnShutdown                   bool
 	AgentVersion                     string
+	UserAgent                        string
 	AgentName                        string
 	KubernetesExecutor               bool
 	KubernetesPodSpec                string
@@ -55,7 +56,7 @@ type Config struct {
 func Start(httpClient *http.Client, config Config) (*Listener, error) {
 	listener := &Listener{
 		Config: config,
-		Client: selfhostedapi.New(httpClient, config.Scheme, config.Endpoint, config.Token),
+		Client: selfhostedapi.New(httpClient, config.Scheme, config.Endpoint, config.Token, config.UserAgent),
 	}
 
 	listener.DisplayHelloMessage()