From 26953f54484977c51d7d97ab5291bec82e05d8ae Mon Sep 17 00:00:00 2001
From: Blake Gentry <bgentry@users.noreply.github.com>
Date: Sat, 7 Dec 2024 16:17:01 -0600
Subject: [PATCH] add basic http auth, adopt apimiddleware system (#241)

Co-authored-by: TArch64 <tarasturchenko1551@gmail.com>
---
 CHANGELOG.md                   |  4 ++++
 cmd/riverui/auth_middleware.go | 31 +++++++++++++++++++++++++++
 cmd/riverui/main.go            | 39 +++++++++++++++++++++-------------
 docs/README.md                 |  5 +++++
 4 files changed, 64 insertions(+), 15 deletions(-)
 create mode 100644 cmd/riverui/auth_middleware.go

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b17c74f..b2423f9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add support for basic auth to the riverui executable. Thanks [Taras Turchenko](https://github.com/TArch64)! [PR #241](https://github.com/riverqueue/riverui/pull/241).
+
 ## [0.6.0] - 2024-11-26
 
 ### Added
diff --git a/cmd/riverui/auth_middleware.go b/cmd/riverui/auth_middleware.go
new file mode 100644
index 0000000..acc6d72
--- /dev/null
+++ b/cmd/riverui/auth_middleware.go
@@ -0,0 +1,31 @@
+package main
+
+import (
+	"crypto/subtle"
+	"net/http"
+)
+
+type authMiddleware struct {
+	username string
+	password string
+}
+
+func (m *authMiddleware) Middleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
+		if isReqAuthorized(req, m.username, m.password) {
+			next.ServeHTTP(res, req)
+			return
+		}
+
+		res.Header().Set("WWW-Authenticate", "Basic realm=\"riverui\"")
+		http.Error(res, "Unauthorized", http.StatusUnauthorized)
+	})
+}
+
+func isReqAuthorized(req *http.Request, username, password string) bool {
+	reqUsername, reqPassword, ok := req.BasicAuth()
+
+	return ok &&
+		subtle.ConstantTimeCompare([]byte(reqUsername), []byte(username)) == 1 &&
+		subtle.ConstantTimeCompare([]byte(reqPassword), []byte(password)) == 1
+}
diff --git a/cmd/riverui/main.go b/cmd/riverui/main.go
index 9d90e03..7c4d06a 100644
--- a/cmd/riverui/main.go
+++ b/cmd/riverui/main.go
@@ -21,6 +21,7 @@ import (
 	"github.com/riverqueue/river/riverdriver/riverpgxv5"
 
 	"riverqueue.com/riverui"
+	"riverqueue.com/riverui/internal/apimiddleware"
 )
 
 func main() {
@@ -48,11 +49,13 @@ func initAndServe(ctx context.Context) int {
 	pathPrefix = riverui.NormalizePathPrefix(pathPrefix)
 
 	var (
-		corsOrigins = strings.Split(os.Getenv("CORS_ORIGINS"), ",")
-		dbURL       = mustEnv("DATABASE_URL")
-		host        = os.Getenv("RIVER_HOST") // may be left empty to bind to all local interfaces
-		otelEnabled = os.Getenv("OTEL_ENABLED") == "true"
-		port        = cmp.Or(os.Getenv("PORT"), "8080")
+		basicAuthUsername = os.Getenv("RIVER_BASIC_AUTH_USER")
+		basicAuthPassword = os.Getenv("RIVER_BASIC_AUTH_PASS")
+		corsOrigins       = strings.Split(os.Getenv("CORS_ORIGINS"), ",")
+		dbURL             = mustEnv("DATABASE_URL")
+		host              = os.Getenv("RIVER_HOST") // may be left empty to bind to all local interfaces
+		otelEnabled       = os.Getenv("OTEL_ENABLED") == "true"
+		port              = cmp.Or(os.Getenv("PORT"), "8080")
 	)
 
 	dbPool, err := getDBPool(ctx, dbURL)
@@ -62,11 +65,6 @@ func initAndServe(ctx context.Context) int {
 	}
 	defer dbPool.Close()
 
-	corsHandler := cors.New(cors.Options{
-		AllowedMethods: []string{"GET", "HEAD", "POST", "PUT"},
-		AllowedOrigins: corsOrigins,
-	})
-
 	client, err := river.NewClient(riverpgxv5.New(dbPool), &river.Config{})
 	if err != nil {
 		logger.ErrorContext(ctx, "error creating river client", slog.String("error", err.Error()))
@@ -88,21 +86,32 @@ func initAndServe(ctx context.Context) int {
 		return 1
 	}
 
-	if err := server.Start(ctx); err != nil {
+	if err = server.Start(ctx); err != nil {
 		logger.ErrorContext(ctx, "error starting UI server", slog.String("error", err.Error()))
 		return 1
 	}
 
-	logHandler := sloghttp.Recovery(server)
-	config := sloghttp.Config{
+	corsHandler := cors.New(cors.Options{
+		AllowedMethods: []string{"GET", "HEAD", "POST", "PUT"},
+		AllowedOrigins: corsOrigins,
+	})
+	logHandler := sloghttp.NewWithConfig(logger, sloghttp.Config{
 		WithSpanID:  otelEnabled,
 		WithTraceID: otelEnabled,
+	})
+
+	middlewareStack := apimiddleware.NewMiddlewareStack(
+		apimiddleware.MiddlewareFunc(sloghttp.Recovery),
+		apimiddleware.MiddlewareFunc(corsHandler.Handler),
+		apimiddleware.MiddlewareFunc(logHandler),
+	)
+	if basicAuthUsername != "" && basicAuthPassword != "" {
+		middlewareStack.Use(&authMiddleware{username: basicAuthUsername, password: basicAuthPassword})
 	}
-	wrappedHandler := sloghttp.NewWithConfig(logger, config)(corsHandler.Handler(logHandler))
 
 	srv := &http.Server{
 		Addr:              host + ":" + port,
-		Handler:           wrappedHandler,
+		Handler:           middlewareStack.Mount(server),
 		ReadHeaderTimeout: 5 * time.Second,
 	}
 
diff --git a/docs/README.md b/docs/README.md
index 190804d..ea7e10b 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -55,6 +55,11 @@ The `riverui` command utilizes the `RIVER_LOG_LEVEL` environment variable to con
 * `warn`
 * `error`
 
+### Basic HTTP Authentication
+
+The `riverui` supports basic HTTP authentication to protect access to the UI.
+To enable it, set the `RIVER_BASIC_AUTH_USER` and `RIVER_BASIC_AUTH_PASS` environment variables.
+
 ## Development
 
 See [developing River UI](./development.md).