feat: add ability to set custom headers

cmd/root.go

@@ -133,7 +133,7 @@ var rootCmd = &cobra.Command{
 		hooks := hooks.NewHooks(auth, providers)
 
 		// Create handlers
-		handlers := handlers.NewHandlers(serverConfig, auth, hooks, providers)
+		handlers := handlers.NewHandlers(serverConfig, auth, hooks, providers, docker)
 
 		// Create API
 		api := api.NewAPI(apiConfig, handlers)

docker-compose.dev.yml

@@ -14,7 +14,6 @@ services:
     labels:
       traefik.enable: true
       traefik.http.routers.nginx.rule: Host(`whoami.example.com`)
-      traefik.http.services.nginx.loadbalancer.server.port: 80
       traefik.http.routers.nginx.middlewares: tinyauth
 
   tinyauth-frontend:
@@ -29,7 +28,6 @@ services:
     labels:
       traefik.enable: true
       traefik.http.routers.tinyauth.rule: Host(`tinyauth.example.com`)
-      traefik.http.services.tinyauth.loadbalancer.server.port: 5173
 
   tinyauth-backend:
     container_name: tinyauth-backend
@@ -41,6 +39,7 @@ services:
       - ./internal:/tinyauth/internal
       - ./cmd:/tinyauth/cmd
       - ./main.go:/tinyauth/main.go
+      - /var/run/docker.sock:/var/run/docker.sock
     ports:
       - 3000:3000
     labels:

docker-compose.example.yml

@@ -14,7 +14,6 @@ services:
     labels:
       traefik.enable: true
       traefik.http.routers.nginx.rule: Host(`whoami.example.com`)
-      traefik.http.services.nginx.loadbalancer.server.port: 80
       traefik.http.routers.nginx.middlewares: tinyauth
 
   tinyauth:
@@ -27,5 +26,4 @@ services:
     labels:
       traefik.enable: true
       traefik.http.routers.tinyauth.rule: Host(`tinyauth.example.com`)
-      traefik.http.services.tinyauth.loadbalancer.server.port: 3000
       traefik.http.middlewares.tinyauth.forwardauth.address: http://tinyauth:3000/api/auth/traefik

internal/api/api_test.go

@@ -78,7 +78,7 @@ func getAPI(t *testing.T) *api.API {
 	hooks := hooks.NewHooks(auth, providers)
 
 	// Create handlers service
-	handlers := handlers.NewHandlers(handlersConfig, auth, hooks, providers)
+	handlers := handlers.NewHandlers(handlersConfig, auth, hooks, providers, docker)
 
 	// Create API
 	api := api.NewAPI(apiConfig, handlers)

internal/auth/auth.go

@@ -159,41 +159,35 @@ func (auth *Auth) ResourceAllowed(c *gin.Context, context types.UserContext) (bo
 	// Get app id
 	appId := strings.Split(host, ".")[0]
 
-	// Check if resource is allowed
-	allowed, err := auth.Docker.ContainerAction(appId, func(labels types.TinyauthLabels) (bool, error) {
-		// If the container has an oauth whitelist, check if the user is in it
-		if context.OAuth {
-			if len(labels.OAuthWhitelist) == 0 {
-				return true, nil
-			}
-			log.Debug().Msg("Checking OAuth whitelist")
-			if slices.Contains(labels.OAuthWhitelist, context.Username) {
-				return true, nil
-			}
-			return false, nil
-		}
-
-		// If the container has users, check if the user is in it
-		if len(labels.Users) != 0 {
-			log.Debug().Msg("Checking users")
-			if slices.Contains(labels.Users, context.Username) {
-				return true, nil
-			}
-			return false, nil
-		}
-
-		// Allowed
-		return true, nil
-	})
+	// Get the container labels
+	labels, err := auth.Docker.GetLabels(appId)
 
 	// If there is an error, return false
 	if err != nil {
-		log.Error().Err(err).Msg("Error checking if resource is allowed")
 		return false, err
 	}
 
-	// Return if the resource is allowed
-	return allowed, nil
+	// Check if oauth is allowed
+	if context.OAuth {
+		if len(labels.OAuthWhitelist) == 0 {
+			return true, nil
+		}
+		log.Debug().Msg("Checking OAuth whitelist")
+		if slices.Contains(labels.OAuthWhitelist, context.Username) {
+			return true, nil
+		}
+	}
+
+	// Check if user is allowed
+	if len(labels.Users) != 0 {
+		log.Debug().Msg("Checking users")
+		if slices.Contains(labels.Users, context.Username) {
+			return true, nil
+		}
+	}
+
+	// Not allowed
+	return false, nil
 }
 
 func (auth *Auth) AuthEnabled(c *gin.Context) (bool, error) {
@@ -204,40 +198,37 @@ func (auth *Auth) AuthEnabled(c *gin.Context) (bool, error) {
 	// Get app id
 	appId := strings.Split(host, ".")[0]
 
-	// Check if auth is enabled
-	enabled, err := auth.Docker.ContainerAction(appId, func(labels types.TinyauthLabels) (bool, error) {
-		// Check if the allowed label is empty
-		if labels.Allowed == "" {
-			// Auth enabled
-			return true, nil
-		}
-
-		// Compile regex
-		regex, err := regexp.Compile(labels.Allowed)
-
-		// If there is an error, invalid regex, auth enabled
-		if err != nil {
-			log.Warn().Err(err).Msg("Invalid regex")
-			return true, err
-		}
-
-		// Check if the uri matches the regex
-		if regex.MatchString(uri) {
-			// Auth disabled
-			return false, nil
-		}
-
-		// Auth enabled
-		return true, nil
-	})
+	// Get the container labels
+	labels, err := auth.Docker.GetLabels(appId)
 
 	// If there is an error, auth enabled
 	if err != nil {
-		log.Error().Err(err).Msg("Error checking if auth is enabled")
 		return true, err
 	}
 
-	return enabled, nil
+	// Check if the allowed label is empty
+	if labels.Allowed == "" {
+		// Auth enabled
+		return true, nil
+	}
+
+	// Compile regex
+	regex, err := regexp.Compile(labels.Allowed)
+
+	// If there is an error, invalid regex, auth enabled
+	if err != nil {
+		log.Warn().Err(err).Msg("Invalid regex")
+		return true, err
+	}
+
+	// Check if the uri matches the regex
+	if regex.MatchString(uri) {
+		// Auth disabled
+		return false, nil
+	}
+
+	// Auth enabled
+	return true, nil
 }
 
 func (auth *Auth) GetBasicAuth(c *gin.Context) *types.User {

internal/constants/constants.go

@@ -5,4 +5,5 @@ var TinyauthLabels = []string{
 	"tinyauth.oauth.whitelist",
 	"tinyauth.users",
 	"tinyauth.allowed",
+	"tinyauth.headers",
 }

internal/docker/docker.go

@@ -3,7 +3,7 @@ package docker
 import (
 	"context"
 	"strings"
-	appTypes "tinyauth/internal/types"
+	"tinyauth/internal/types"
 	"tinyauth/internal/utils"
 
 	apiTypes "github.com/docker/docker/api/types"
@@ -70,14 +70,14 @@ func (docker *Docker) DockerConnected() bool {
 	return err == nil
 }
 
-func (docker *Docker) ContainerAction(appId string, runCheck func(labels appTypes.TinyauthLabels) (bool, error)) (bool, error) {
+func (docker *Docker) GetLabels(appId string) (types.TinyauthLabels, error) {
 	// Check if we have access to the Docker API
 	isConnected := docker.DockerConnected()
 
-	// If we don't have access, it is assumed that the check passed
+	// If we don't have access, return an empty struct
 	if !isConnected {
-		log.Debug().Msg("Docker not connected, passing check")
-		return true, nil
+		log.Debug().Msg("Docker not connected, returning empty labels")
+		return types.TinyauthLabels{}, nil
 	}
 
 	// Get the containers
@@ -85,7 +85,7 @@ func (docker *Docker) ContainerAction(appId string, runCheck func(labels appType
 
 	// If there is an error, return false
 	if err != nil {
-		return false, err
+		return types.TinyauthLabels{}, err
 	}
 
 	log.Debug().Msg("Got containers")
@@ -97,11 +97,11 @@ func (docker *Docker) ContainerAction(appId string, runCheck func(labels appType
 
 		// If there is an error, return false
 		if err != nil {
-			return false, err
+			return types.TinyauthLabels{}, err
 		}
 
 		// Get the container name (for some reason it is /name)
-		containerName := strings.Split(inspect.Name, "/")[1]
+		containerName := strings.TrimPrefix(inspect.Name, "/")
 
 		// There is a container with the same name as the app ID
 		if containerName == appId {
@@ -112,14 +112,14 @@ func (docker *Docker) ContainerAction(appId string, runCheck func(labels appType
 
 			log.Debug().Msg("Got labels")
 
-			// Run the check
-			return runCheck(labels)
+			// Return labels
+			return labels, nil
 		}
 
 	}
 
-	log.Debug().Msg("No matching container found, passing check")
+	log.Debug().Msg("No matching container found, returning empty labels")
 
-	// If no matching container is found, pass check
-	return true, nil
+	// If no matching container is found, return empty labels
+	return types.TinyauthLabels{}, nil
 }

internal/handlers/handlers.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"strings"
 	"tinyauth/internal/auth"
+	"tinyauth/internal/docker"
 	"tinyauth/internal/hooks"
 	"tinyauth/internal/providers"
 	"tinyauth/internal/types"
@@ -16,12 +17,13 @@ import (
 	"github.com/rs/zerolog/log"
 )
 
-func NewHandlers(config types.HandlersConfig, auth *auth.Auth, hooks *hooks.Hooks, providers *providers.Providers) *Handlers {
+func NewHandlers(config types.HandlersConfig, auth *auth.Auth, hooks *hooks.Hooks, providers *providers.Providers, docker *docker.Docker) *Handlers {
 	return &Handlers{
 		Config:    config,
 		Auth:      auth,
 		Hooks:     hooks,
 		Providers: providers,
+		Docker:    docker,
 	}
 }
 
@@ -30,6 +32,7 @@ type Handlers struct {
 	Auth      *auth.Auth
 	Hooks     *hooks.Hooks
 	Providers *providers.Providers
+	Docker    *docker.Docker
 }
 
 func (h *Handlers) AuthHandler(c *gin.Context) {
@@ -60,12 +63,39 @@ func (h *Handlers) AuthHandler(c *gin.Context) {
 
 	log.Debug().Interface("proxy", proxy.Proxy).Msg("Got proxy")
 
+	// Get headers
+	uri := c.Request.Header.Get("X-Forwarded-Uri")
+	proto := c.Request.Header.Get("X-Forwarded-Proto")
+	host := c.Request.Header.Get("X-Forwarded-Host")
+
 	// Check if auth is enabled
 	authEnabled, err := h.Auth.AuthEnabled(c)
 
-	// Handle error
+	// Check if there was an error
 	if err != nil {
-		log.Error().Err(err).Msg("Failed to check if auth is enabled")
+		log.Error().Err(err).Msg("Failed to check if app is allowed")
+
+		if proxy.Proxy == "nginx" || !isBrowser {
+			c.JSON(500, gin.H{
+				"status":  500,
+				"message": "Internal Server Error",
+			})
+			return
+		}
+
+		c.Redirect(http.StatusPermanentRedirect, fmt.Sprintf("%s/error", h.Config.AppURL))
+		return
+	}
+
+	// Get the app id
+	appId := strings.Split(host, ".")[0]
+
+	// Get the container labels
+	labels, err := h.Docker.GetLabels(appId)
+
+	// Check if there was an error
+	if err != nil {
+		log.Error().Err(err).Msg("Failed to check if app is allowed")
 
 		if proxy.Proxy == "nginx" || !isBrowser {
 			c.JSON(500, gin.H{
@@ -81,6 +111,10 @@ func (h *Handlers) AuthHandler(c *gin.Context) {
 
 	// If auth is not enabled, return 200
 	if !authEnabled {
+		for key, value := range labels.Headers {
+			log.Debug().Str("key", key).Str("value", value).Msg("Setting header")
+			c.Header(key, value)
+		}
 		c.JSON(200, gin.H{
 			"status":  200,
 			"message": "Authenticated",
@@ -91,11 +125,6 @@ func (h *Handlers) AuthHandler(c *gin.Context) {
 	// Get user context
 	userContext := h.Hooks.UseUserContext(c)
 
-	// Get headers
-	uri := c.Request.Header.Get("X-Forwarded-Uri")
-	proto := c.Request.Header.Get("X-Forwarded-Proto")
-	host := c.Request.Header.Get("X-Forwarded-Host")
-
 	// Check if user is logged in
 	if userContext.IsLoggedIn {
 		log.Debug().Msg("Authenticated")
@@ -157,6 +186,12 @@ func (h *Handlers) AuthHandler(c *gin.Context) {
 		// Set the user header
 		c.Header("Remote-User", userContext.Username)
 
+		// Set the rest of the headers
+		for key, value := range labels.Headers {
+			log.Debug().Str("key", key).Str("value", value).Msg("Setting header")
+			c.Header(key, value)
+		}
+
 		// The user is allowed to access the app
 		c.JSON(200, gin.H{
 			"status":  200,

internal/types/types.go

@@ -124,6 +124,7 @@ type TinyauthLabels struct {
 	OAuthWhitelist []string
 	Users          []string
 	Allowed        string
+	Headers        map[string]string
 }
 
 // TailscaleQuery is the query parameters for the tailscale endpoint

internal/utils/utils.go

@@ -193,6 +193,16 @@ func GetTinyauthLabels(labels map[string]string) types.TinyauthLabels {
 				tinyauthLabels.Users = strings.Split(value, ",")
 			case "tinyauth.allowed":
 				tinyauthLabels.Allowed = value
+			case "tinyauth.headers":
+				tinyauthLabels.Headers = make(map[string]string)
+				headers := strings.Split(value, ",")
+				for _, header := range headers {
+					headerSplit := strings.Split(header, "=")
+					if len(headerSplit) != 2 {
+						continue
+					}
+					tinyauthLabels.Headers[headerSplit[0]] = headerSplit[1]
+				}
 			}
 		}
 	}