[minor] add retry on rollout restart (#20)

.golangci.yml

@@ -1,4 +1,7 @@
 linters:
   enable:
     - gofmt
-
+issues:
+  exclude-dirs:
+    - ../../go
+    - ../../../../opt

README.md

@@ -44,6 +44,7 @@ This service requires two environment variables.
 ```
 - `ROLLOUT_CMD` (default: `/bin/bash`) - the command to execute a rollout
 - `ROLLOUT_ARGS` (default: `/rollout.sh` ) - the args to pass to `ROLLOUT_CMD`
+- `ROLLOUT_LOCK_FILE` (optional) - a lock file that is set during rollouts. When the file is on disk additional rollouts can not happen, and if the file is present on startup initiates a rollout
 
 ## Dynamic environment variables for ROLLOUT_CMD
 

go.mod

@@ -6,7 +6,6 @@ require (
 	github.com/golang-jwt/jwt/v5 v5.2.1
 	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510
 	github.com/lestrrat-go/jwx v1.2.30
-	github.com/lestrrat-go/jwx/v2 v2.1.3
 	github.com/stretchr/testify v1.10.0
 )
 

go.sum

@@ -19,7 +19,6 @@ github.com/lestrrat-go/iter v1.0.2 h1:gMXo1q4c2pHmC3dn8LzRhJfP1ceCbgSiT9lUydIzlt
 github.com/lestrrat-go/iter v1.0.2/go.mod h1:Momfcq3AnRlRjI5b5O8/G5/BvpzrhoFTZcn06fEOPt4=
 github.com/lestrrat-go/jwx v1.2.30 h1:VKIFrmjYn0z2J51iLPadqoHIVLzvWNa1kCsTqNDHYPA=
 github.com/lestrrat-go/jwx v1.2.30/go.mod h1:vMxrwFhunGZ3qddmfmEm2+uced8MSI6QFWGTKygjSzQ=
-github.com/lestrrat-go/jwx/v2 v2.1.3/go.mod h1:q6uFgbgZfEmQrfJfrCo90QcQOcXFMfbI/fO0NqRtvZo=
 github.com/lestrrat-go/option v1.0.0/go.mod h1:5ZHFbivi4xwXxhxY9XHDe2FHo6/Z7WWmtT7T5nBBp3I=
 github.com/lestrrat-go/option v1.0.1 h1:oAzP2fvZGQKWkvHa1/SAcFolBEca1oN+mQ7eooNBEYU=
 github.com/lestrrat-go/option v1.0.1/go.mod h1:5ZHFbivi4xwXxhxY9XHDe2FHo6/Z7WWmtT7T5nBBp3I=

main.go

@@ -12,7 +12,7 @@ import (
 	"reflect"
 	"regexp"
 
-	"github.com/golang-jwt/jwt/v5"
+	jwt "github.com/golang-jwt/jwt/v5"
 	"github.com/google/shlex"
 	"github.com/lestrrat-go/jwx/jwk"
 )
@@ -33,6 +33,8 @@ func init() {
 }
 
 func main() {
+	var err error
+	rolloutLockFile := os.Getenv("ROLLOUT_LOCK_FILE")
 	if os.Getenv("JWKS_URI") == "" {
 		slog.Error("JWKS_URI is required. e.g. JWKS_URI=https://gitlab.com/oauth/discovery/keys")
 		os.Exit(1)
@@ -42,13 +44,23 @@ func main() {
 		os.Exit(1)
 	}
 
+	if lockExists(rolloutLockFile, false) {
+		slog.Info("Rollout triggered on startup.")
+		err = rollout()
+		if err != nil {
+			slog.Error("Error running rollout", "err", err)
+			os.Exit(1)
+		}
+
+	}
+
 	http.HandleFunc("/healthcheck", func(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusOK)
 		fmt.Fprintln(w, "OK")
 	})
 	http.HandleFunc("/", Rollout)
 	slog.Info("Server is running on :8080")
-	err := http.ListenAndServe(":8080", nil)
+	err = http.ListenAndServe(":8080", nil)
 	if err != nil {
 		slog.Error("Unable to start service")
 		os.Exit(1)
@@ -66,8 +78,6 @@ func Rollout(w http.ResponseWriter, r *http.Request) {
 	}
 	// Assuming "Bearer " prefix
 	tokenString := a[7:]
-
-	// Parse and verify the token
 	token, err := jwt.Parse(tokenString, ParseToken)
 	if err != nil {
 		slog.Info("Failed to verify token for", "forwarded-ip", realIp, "lasthop-ip", lastIP, "err", err.Error())
@@ -111,23 +121,21 @@ func Rollout(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	name := os.Getenv("ROLLOUT_CMD")
-	if name == "" {
-		name = "/bin/bash"
+	rolloutLockFile := os.Getenv("ROLLOUT_LOCK_FILE")
+	if rolloutLockFile != "" && lockExists(rolloutLockFile, true) {
+		slog.Error("Lock file exists. Not rolling out")
+		http.Error(w, "Internal error", http.StatusInternalServerError)
+		return
 	}
-	cmd := exec.Command(name, getRolloutCmdArgs()...)
 
-	var stdOut, stdErr bytes.Buffer
-	cmd.Stdout = &stdOut
-	cmd.Stderr = &stdErr
-	cmd.Env = os.Environ()
-	if err := cmd.Run(); err != nil {
-		slog.Error("Error running", "command", cmd.String(), "stdout", stdOut.String(), "stderr", stdErr.String())
+	err = rollout()
+	if err != nil {
+		slog.Error("Error running", "err", err)
 		http.Error(w, "Script execution failed", http.StatusInternalServerError)
 		return
 	}
+	slog.Info("Rollout complete for", "forwarded-ip", r.Header.Get("X-Forwarded-For"), "lasthop-ip", r.RemoteAddr)
 
-	slog.Info("Rollout complete for", "forwarded-ip", realIp, "lasthop-ip", lastIP)
 	fmt.Fprintln(w, "Rollout complete")
 }
 
@@ -253,3 +261,55 @@ func setEnvFromStruct(data interface{}) error {
 	}
 	return nil
 }
+
+func lockExists(filePath string, createOnNotExists bool) bool {
+	if filePath == "" {
+		return false
+	}
+
+	_, err := os.ReadFile(filePath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			if !createOnNotExists {
+				return false
+			}
+			err := os.WriteFile(filePath, []byte("rollout"), 0644)
+			if err != nil {
+				slog.Error("Failed to create rollout lock file", "error", err)
+			}
+			return false
+		}
+		slog.Error("Error reading rollout lock file", "error", err)
+		return false
+	}
+
+	return true
+}
+
+func rollout() error {
+	name := os.Getenv("ROLLOUT_CMD")
+	if name == "" {
+		name = "/bin/bash"
+	}
+	cmd := exec.Command(name, getRolloutCmdArgs()...)
+
+	var stdOut, stdErr bytes.Buffer
+	cmd.Stdout = &stdOut
+	cmd.Stderr = &stdErr
+	cmd.Env = os.Environ()
+	if err := cmd.Run(); err != nil {
+		return fmt.Errorf("command: %s had stdout:%s stderr:%s", cmd.String(), stdOut.String(), stdErr.String())
+	}
+
+	rolloutLockFile := os.Getenv("ROLLOUT_LOCK_FILE")
+	if rolloutLockFile == "" {
+		return nil
+	}
+
+	err := os.Remove(rolloutLockFile)
+	if err != nil {
+		return fmt.Errorf("failed to remove rollout lock file: %v", err)
+	}
+
+	return nil
+}

main_test.go

@@ -16,7 +16,7 @@ import (
 	"testing"
 	"time"
 
-	"github.com/golang-jwt/jwt/v5"
+	jwt "github.com/golang-jwt/jwt/v5"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -480,3 +480,108 @@ func RemoveFileIfExists(filePath string) error {
 
 	return nil
 }
+
+func TestLockFile(t *testing.T) {
+	lockFile := "/tmp/rollout.lock"
+	testFile := "/tmp/rollout-test.txt"
+
+	// have our test rollout cmd just touch a file
+	os.Setenv("ROLLOUT_CMD", "touch")
+	os.Setenv("ROLLOUT_ARGS", testFile)
+	os.Setenv("ROLLOUT_LOCK_FILE", lockFile)
+
+	// make sure the test and lock file doesn't exist
+	err := RemoveFileIfExists(testFile)
+	if err != nil {
+		slog.Error("Unable to cleanup test file", "err", err)
+		os.Exit(1)
+	}
+	err = RemoveFileIfExists(lockFile)
+	if err != nil {
+		slog.Error("Unable to cleanup lock file", "err", err)
+		os.Exit(1)
+	}
+
+	// create the lock file
+	lockExists(lockFile, true)
+	s := createMockJwksServer()
+	defer s.Close()
+
+	// get a valid token
+	exp := time.Now().Add(time.Hour * 1).Unix()
+	jwtToken, err := CreateSignedJWT(kid, aud, claim, exp, privateKey)
+	if err != nil {
+		t.Fatalf("Unable to create a JWT with our test key: %v", err)
+	}
+
+	tests := []Test{
+		{
+			name:           "Do not roll out when locked",
+			authHeader:     "Bearer " + jwtToken,
+			expectedStatus: http.StatusInternalServerError,
+			expectedBody:   "Internal error\n",
+		},
+		{
+			name:           "Rollout OK when not locked",
+			authHeader:     "Bearer " + jwtToken,
+			expectedStatus: http.StatusOK,
+			expectedBody:   "Rollout complete\n",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			recorder := httptest.NewRecorder()
+			request := createRequest(tt.authHeader, "GET", nil)
+			if tt.name == "No custom claim" {
+				os.Setenv("CUSTOM_CLAIMS", "")
+			} else {
+				os.Setenv("CUSTOM_CLAIMS", `{"foo": "bar"}`)
+			}
+			if tt.cmdArgs != "" {
+				os.Setenv("ROLLOUT_ARGS", tt.cmdArgs)
+			}
+
+			Rollout(recorder, request)
+
+			assert.Equal(t, tt.expectedStatus, recorder.Code)
+			assert.Equal(t, tt.expectedBody, recorder.Body.String())
+
+			// on the first test, remove the lock file
+			// so the second test should pass OK
+			if tt.name == "Do not roll out when locked" {
+				// remove the file
+				err = RemoveFileIfExists(lockFile)
+				if err != nil {
+					slog.Error("Unable to cleanup lock file", "err", err)
+					os.Exit(1)
+				}
+			}
+		})
+	}
+
+	// make sure the lock file was removed
+	_, err = os.Stat(lockFile)
+	if err == nil {
+		t.Errorf("The successful test did not cleanup the lock file %s", lockFile)
+	}
+
+	// make sure the rollout command actually ran the command
+	// which creates the file
+	_, err = os.Stat(testFile)
+	if err != nil && os.IsNotExist(err) {
+		t.Errorf("The successful test did not create the expected file %s", testFile)
+	}
+
+	testFiles := []string{
+		testFile,
+		lockFile,
+	}
+	for _, f := range testFiles {
+		// cleanup
+		err = RemoveFileIfExists(f)
+		if err != nil {
+			slog.Error("Unable to cleanup test file", "file", f, "err", err)
+			os.Exit(1)
+		}
+	}
+}