granular control over WYSIWYG

.vscode/settings.json

@@ -13,7 +13,9 @@
     "Dawid",
     "eqfield",
     "Expirable",
+    "expiremap",
     "extldflags",
+    "geolocation",
     "ginzap",
     "Ginzap",
     "goccy",
@@ -40,6 +42,7 @@
     "multitemplate",
     "netgo",
     "nosniff",
+    "nursik",
     "NXDOMAIN",
     "osusergo",
     "pquerna",

cmd/kagi-proxy/main.go

@@ -4,10 +4,11 @@ import (
 	"flag"
 	"fmt"
 	"html"
-	"math"
 	"net/http"
+	"net/url"
 	"os"
 	"slices"
+	"strings"
 	"time"
 
 	"github.com/gin-contrib/cors"
@@ -65,7 +66,14 @@ func main() {
 			*proxyHost:                "kagi.com",
 			"translate." + *proxyHost: "translate.kagi.com",
 			"assets." + *proxyHost:    "assets.kagi.com",
-			"status" + *proxyHost:     "status.kagi.com",
+			"status." + *proxyHost:    "status.kagi.com",
+		}),
+		common.WithProxyGuardRules(common.Ruleset{
+			common.Rule{Path: "/settings", PathType: common.Exact, Query: url.Values{"p": {"billing"}}},
+			common.Rule{Path: "/settings", PathType: common.Exact, Query: url.Values{"p": {"gift"}}},
+			common.Rule{Path: "/settings", PathType: common.Exact, Query: url.Values{"p": {"user_details"}}},
+			common.Rule{Path: "/settings", PathType: common.Exact, Query: url.Values{"p": {"api"}, "generate": {"1"}}},
+			common.Rule{Path: "/api/user_token", PathType: common.Prefix},
 		}),
 		common.WithProxyRedirectLoginURL("/signin"),
 	)
@@ -79,30 +87,50 @@ func main() {
 		hashKey, blockKey := common.MakeKeyPair([]byte(*proxySessionSecret))
 		store := cookie.NewStore(hashKey, blockKey)
 		store.Options(sessions.Options{
-			Domain:   *proxyHost, // Required to support wildcard subdomains
-			Path:     "/",
-			MaxAge:   math.MaxInt32, // 0: session cookie until the browser is closed, -1: delete the cookie, math.MaxInt32: 68 years
+			Domain: *proxyHost, // Required to support wildcard subdomains
+			Path:   "/",
+			// 0: session cookie until the browser is closed, -1: delete the cookie, math.MaxInt32: 68 years
+			MaxAge:   int((24 * 7 * time.Hour).Seconds()),
 			Secure:   true,
 			HttpOnly: true,
 			SameSite: http.SameSiteLaxMode,
 		})
 		e.Use(sessions.Sessions("proxy_session", store))
 
 		// Setup CORS
-		config := cors.DefaultConfig()
+		config := cors.Config{}
 		config.AllowCredentials = true
 		config.AllowBrowserExtensions = true
-		config.AllowWebSockets = true
-		config.AddAllowHeaders("X-Requested-With")
-		config.AllowOrigins = append(config.AllowOrigins,
-			"http://localhost:"+fmt.Sprint(*port),
-			"https://kagi.com",
-			"https://"+*proxyHost,
-			"https://*.kagi.com",
-			"https://*."+*proxyHost,
-		)
+		config.AddAllowHeaders("X-Requested-With", "Content-Type", "Authorization", "Origin", "Accept")
+		config.AllowOriginFunc = func(origin string) bool {
+			parsed, err := url.Parse(origin)
+			if err != nil {
+				return false
+			}
+
+			if parsed.Hostname() == "localhost" {
+				return true
+			}
+
+			hostname := parsed.Hostname()
+			for targetDomain, proxyDomain := range common.ConfigProxyTargetHosts() {
+				switch {
+				case hostname == targetDomain,
+					hostname == proxyDomain,
+					strings.HasSuffix(hostname, "."+targetDomain),
+					strings.HasSuffix(hostname, "."+proxyDomain):
+
+					return true
+				}
+			}
+
+			return false
+		}
 		config.AddExposeHeaders("Location", "Content-Disposition")
-		config.AllowWildcard = true
+		config.AllowMethods = []string{http.MethodGet, http.MethodPost, http.MethodPut, http.MethodPatch, http.MethodDelete, http.MethodOptions}
+		if err := config.Validate(); err != nil {
+			common.Logger().Fatal("Invalid CORS configuration", zap.Error(err))
+		}
 		e.Use(cors.New(config))
 
 		// Use the ginzap middleware to log requests.
@@ -143,14 +171,20 @@ func main() {
 	// Overwrite the logout route.
 	router.GET("/logout", api.HandleLogout())
 
-	// Overwrite the settings route.
-	router.GET("/settings", api.HandleUnauthorized())
-
 	// Add a proxy route for anything else, do not require authentication for the favicon.
-	router.NoRoute(api.BasicAuth([]string{"/favicon.ico"}), api.ProxyPass())
+	router.NoRoute(api.BasicAuth([]string{"/favicon.ico"}), api.ProxyGuard(), api.ProxyPass())
+
+	server := &http.Server{
+		Addr:    fmt.Sprintf(":%d", *port),
+		Handler: router,
+	}
+	server.RegisterOnShutdown(func() {
+		common.Logger().Info("Server shutdown")
+		api.SessionMap.Close()
+	})
 
 	// Start the server.
-	if err := router.Run(fmt.Sprintf(":%d", *port)); err != nil {
+	if err := server.ListenAndServe(); err != nil {
 		common.Logger().Fatal("Unexpected server error", zap.Error(err))
 	}
 }

go.mod

@@ -8,7 +8,9 @@ require (
 	github.com/gin-contrib/sessions v1.0.2
 	github.com/gin-contrib/zap v1.1.4
 	github.com/gin-gonic/gin v1.10.0
+	github.com/google/uuid v1.6.0
 	github.com/klauspost/compress v1.17.7
+	github.com/nursik/go-expire-map v1.2.0
 	github.com/pquerna/otp v1.4.0
 	github.com/utrack/gin-csrf v0.0.0-20190424104817-40fb8d2c8fca
 	go.uber.org/zap v1.27.0

go.sum

@@ -52,6 +52,8 @@ github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/gofuzz v1.2.0 h1:xRy4A+RhZaiKjJ1bPfwQ8sedCA+YS2YcCHW6ec7JMi0=
 github.com/google/gofuzz v1.2.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/gorilla/context v1.1.1/go.mod h1:kBGZzfjB9CEq2AlWe17Uuf7NDRt0dE0s8S51q0aT7Yg=
 github.com/gorilla/context v1.1.2 h1:WRkNAv2uoa03QNIc1A6u4O7DAGMUVoopZhkiXWA2V1o=
 github.com/gorilla/context v1.1.2/go.mod h1:KDPwT9i/MeWHiLl90fuTgrt4/wPcv75vFAZLaOOcbxM=
@@ -91,6 +93,8 @@ github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJ
 github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
 github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
 github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
+github.com/nursik/go-expire-map v1.2.0 h1:3yl3sVLSnfw4vbo6OVZWP1tFND4CBm4hrYCrrAT5jkU=
+github.com/nursik/go-expire-map v1.2.0/go.mod h1:Mrqzxpk2G81At+TAlfImiNrKECj0gfVw7URv3p/eLzU=
 github.com/pelletier/go-toml/v2 v2.2.3 h1:YmeHyLY8mFWbdkNWwpr+qIL2bEqT0o95WSdkNHvL12M=
 github.com/pelletier/go-toml/v2 v2.2.3/go.mod h1:MfCQTFTvCcUyyvvwm1+G6H/jORL20Xlb6rzQu9GuUkc=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=

pkg/api/auth.go

@@ -11,6 +11,8 @@ import (
 
 	"github.com/gin-contrib/sessions"
 	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
+	expiremap "github.com/nursik/go-expire-map"
 	"github.com/pquerna/otp"
 	"github.com/pquerna/otp/totp"
 	"github.com/sarumaj/kagi-proxy/pkg/common"
@@ -19,6 +21,20 @@ import (
 	"golang.org/x/crypto/bcrypt"
 )
 
+// SessionMap is a map that stores the user session.
+// It uses the session ID as the key and the session information as the value.
+var SessionMap = expiremap.New()
+
+// SessionInfo is a struct that stores the user session information.
+type SessionInfo struct {
+	// SessionId is the session ID.
+	SessionId string
+	// CreatedAt is the session creation time.
+	CreatedAt time.Time
+	// Rules are volatile rules that are evaluated at runtime for given session.
+	Rules common.Ruleset
+}
+
 // BasicAuth is a middleware that checks if the user is authenticated.
 // It skips authentication for the paths in exceptPaths.
 // It seeks the proxy_token query parameter to authenticate the user.
@@ -47,6 +63,9 @@ func BasicAuth(exceptPaths []string) gin.HandlerFunc {
 
 				// Establish or overwrite the user session
 				session.Set("user", common.ConfigProxyUser())
+				sessionId, _ := uuid.NewRandom()
+				session.Set("session_id", sessionId.String())
+				session.Set("created_at", time.Now().Unix())
 				if !sessionSave(session, ctx) {
 					return
 				}
@@ -277,6 +296,9 @@ func HandleLogin() gin.HandlerFunc {
 		// Success: Redirect the user to the root page or the location he attempted to access before page
 		if validUsername && validPassword && validOTP {
 			session.Set("user", request.Username)
+			sessionId, _ := uuid.NewRandom()
+			session.Set("session_id", sessionId.String())
+			session.Set("created_at", time.Now().Unix())
 			redirectURL := session.Get("redirect_url")
 			session.Delete("redirect_url")
 			if !sessionSave(session, ctx) {
@@ -320,9 +342,18 @@ func HandleLogout() gin.HandlerFunc {
 	}
 }
 
-// HandleUnauthorized is a handler that displays an unauthorized page.
-func HandleUnauthorized() gin.HandlerFunc {
+// ProxyGuard is a middleware that checks if the request is allowed.
+// It uses the proxy guard rules to evaluate the request.
+func ProxyGuard() gin.HandlerFunc {
 	return func(ctx *gin.Context) {
+		session := sessions.Default(ctx)
+		sessionId := common.QuickGet[string](session, "session_id")
+		sessionInfo := common.QuickGet[SessionInfo](SessionMap, sessionId)
+		if append(common.ConfigProxyGuardRules(), sessionInfo.Rules...).Evaluate(ctx.Request) == common.Allow {
+			ctx.Next()
+			return
+		}
+
 		nonce, _ := common.GetNonce()
 		SetContentSecurityHeaders(ctx.Writer, nonce)
 		ctx.HTML(http.StatusForbidden, "error.html", gin.H{
@@ -331,6 +362,7 @@ func HandleUnauthorized() gin.HandlerFunc {
 			"error": nil,
 			"nonce": nonce,
 		})
+		ctx.Abort()
 	}
 }
 

pkg/api/fs.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"bytes"
 	"embed"
 	"encoding/json"
 	"io/fs"
@@ -13,16 +14,18 @@ import (
 
 var (
 	funcsMap = map[string]any{
-		"json": func(v any) string {
-			out, _ := json.Marshal(v)
-			return string(out)
+		"json": func(v any) htmlTemplate.JS {
+			out, err := json.Marshal(v)
+			if err != nil {
+				return htmlTemplate.JS("null")
+			}
+			return htmlTemplate.JS(bytes.ReplaceAll(out, []byte{'\\'}, []byte(`\x5c`)))
 		},
 		"codeString": func(code int) string {
 			txt := http.StatusText(code)
 			if len(txt) > 0 {
 				return txt
 			}
-
 			return "Unknown Status Code"
 		},
 	}