2 files changed, 31 insertions, 1 deletions
diff --git a/internal/auth/middleware.go b/internal/auth/middleware.go
index 78f3b53..04e7111 100644
--- a/internal/auth/middleware.go
+++ b/internal/auth/middleware.go
@@ -6,6 +6,7 @@ import (
 	"crypto/subtle"
 	"encoding/base64"
 	"net/http"
+	"net/url"
 	"strings"
 
 	"github.com/alexedwards/scs/v2"
@@ -120,6 +121,31 @@ func GetCSRFTokenFromContext(ctx context.Context) string {
 	return token
 }
 
+// OriginCheck returns middleware that validates the Origin (or Referer) header
+// against allowedOrigin for state-changing requests. Use for reverse-proxy
+// routes where injecting a CSRF token into the proxied UI is not practical.
+// When allowedOrigin is empty the check is skipped (dev/no-TLS setups).
+func OriginCheck(allowedOrigin string) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if allowedOrigin != "" &&
+				(r.Method == "POST" || r.Method == "PUT" || r.Method == "DELETE" || r.Method == "PATCH") {
+				origin := r.Header.Get("Origin")
+				if origin == "" {
+					if ref, err := url.Parse(r.Header.Get("Referer")); err == nil && ref.Host != "" {
+						origin = ref.Scheme + "://" + ref.Host
+					}
+				}
+				if origin != allowedOrigin {
+					http.Error(w, "Forbidden", http.StatusForbidden)
+					return
+				}
+			}
+			next.ServeHTTP(w, r)
+		})
+	}
+}
+
 func generateToken() (string, error) {
 	b := make([]byte, 32)
 	_, err := rand.Read(b)
diff --git a/internal/handlers/claudomator_proxy.go b/internal/handlers/claudomator_proxy.go
index bfbbabc..f3f0cd2 100644
--- a/internal/handlers/claudomator_proxy.go
+++ b/internal/handlers/claudomator_proxy.go
@@ -13,7 +13,7 @@ import (
 // targetURL, stripping the "/claudomator" prefix from the path. WebSocket
 // upgrade requests are handled via raw TCP hijacking to support long-lived
 // connections.
-func NewClaudomatorProxy(targetURL string) http.Handler {
+func NewClaudomatorProxy(targetURL, allowedOrigin string) http.Handler {
 	target, err := url.Parse(targetURL)
 	if err != nil {
 		panic("claudomator: invalid target URL: " + err.Error())
@@ -50,6 +50,10 @@ func NewClaudomatorProxy(targetURL string) http.Handler {
 
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if strings.EqualFold(r.Header.Get("Upgrade"), "websocket") {
+			if allowedOrigin != "" && r.Header.Get("Origin") != allowedOrigin {
+				http.Error(w, "forbidden", http.StatusForbidden)
+				return
+			}
 			proxyWebSocket(w, r, target)
 			return
 		}