perf: update check origin

.github/workflows/golangci-lint.yml

@@ -26,7 +26,7 @@ jobs:
         with:
           go-version: stable
 
-      - uses: golangci/golangci-lint-action@v6
+      - uses: golangci/golangci-lint-action@v9
         with:
           # Optional: version of golangci-lint to use in form of v1.2 or v1.2.3 or `latest` to use the latest version
           version: latest

.golangci.yml

@@ -1,3 +1,5 @@
+version: "2"
+
 run:
   timeout: 5m
   modules-download-mode: readonly

pkg/config/config.go

@@ -68,6 +68,8 @@ type Config struct {
 	DisableKeyboardRecord bool `mapstructure:"DISABLE_KEYBOARD_RECORD"`
 
 	DriveScope string `mapstructure:"LION_DRIVE_SCOPE"` // user or session
+	// DOMAINS=* "demo.example.com:443,172.17.200.191:80"
+	DOMAINS string `mapstructure:"DOMAINS"`
 }
 
 func (c *Config) UpdateRedisPassword(val string) {

pkg/tunnel/server.go

@@ -7,6 +7,8 @@
 	"io"
 	"net"
 	"net/http"
+	"net/netip"
+	"net/url"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -32,13 +34,77 @@
 	defaultBufferSize = 1024
 )
 
+func normalizeHost(h string) string {
+	h = strings.TrimSpace(strings.ToLower(h))
+	host, _, err := net.SplitHostPort(h)
+	if err == nil {
+		return strings.Trim(host, "[]")
+	}
+	return strings.Trim(h, "[]")
+}
+
+func hostMatches(allowedHost, originHost string) bool {
+	allowedName := normalizeHost(allowedHost)
+	originName := normalizeHost(originHost)
+	if allowedName == "" || originName == "" {
+		return false
+	}
+	return allowedName == originName
+}
+
+func isInternalHost(host string) bool {
+	host = normalizeHost(host)
+	if host == "localhost" {
+		return true
+	}
+
+	addr, err := netip.ParseAddr(host)
+	if err != nil {
+		return false
+	}
+	return addr.IsLoopback() || addr.IsPrivate() || addr.IsLinkLocalUnicast()
+}
+
+func domainsAllowOrigin(originHost, domains string) bool {
+	for _, domain := range strings.Split(domains, ",") {
+		domain = strings.TrimSpace(domain)
+		if domain == "" {
+			continue
+		}
+		if domain == "*" || hostMatches(domain, originHost) {
+			return true
+		}
+	}
+	return false
+}
+
+func checkOrigin(r *http.Request) bool {
+	origin := r.Header.Get("Origin")
+	// 允许非浏览器，客户端访问
+	if len(origin) == 0 {
+		return true
+	}
+	u, err := url.Parse(origin)
+	if err != nil {
+		return false
+	}
+	originHost := normalizeHost(u.Host)
+	reqHost := normalizeHost(r.Host)
+	if hostMatches(originHost, reqHost) {
+		return true
+	}
+	if isInternalHost(originHost) {
+		return true
+	}
+
+	return domainsAllowOrigin(originHost, config.GlobalConfig.DOMAINS)
+}
+
 var upGrader = websocket.Upgrader{
 	ReadBufferSize:  defaultBufferSize,
 	WriteBufferSize: defaultBufferSize,
 	Subprotocols:    []string{"guacamole"},
-	CheckOrigin: func(r *http.Request) bool {
-		return true
-	},
+	CheckOrigin:     checkOrigin,
 }
 
 type GuacamoleTunnelServer struct {
@@ -81,7 +147,7 @@
 		ctx.AbortWithStatus(http.StatusBadRequest)
 		return
 	}
 	defer ws.Close()
 
 	tokenId, ok := ctx.GetQuery("TOKEN_ID")
 	if !ok {
@@ -206,7 +272,7 @@
 		g.RecordLifecycleLog(sessionId, model.AssetConnectFinished, reason)
 		return
 	}
 	defer tunnel.Close()
 	g.RecordLifecycleLog(sessionId, model.AssetConnectSuccess, model.EmptyLifecycleLog)
 
 	logger.Infof("Session[%s] use resolution (%d*%d)",
@@ -403,7 +469,7 @@
 		ctx.AbortWithStatus(http.StatusBadRequest)
 		return
 	}
 	defer ws.Close()
 	userItem, ok := ctx.Get(config.GinCtxUserKey)
 	if !ok {
 		_ = ws.WriteMessage(websocket.TextMessage, []byte(ErrAuthUser.String()))
@@ -439,7 +505,7 @@
 		_ = ws.WriteMessage(websocket.TextMessage, []byte(ErrNoSession.String()))
 		return
 	}
 	defer tunnelCon.Close()
 
 	ints := guacd.NewInstruction(INTERNALDATAOPCODE, tunnelCon.UUID())
 	_ = ws.WriteMessage(websocket.TextMessage, []byte(ints.String()))