managed-portal/internal/webdevice/rewrite.go

package webdevice

import (
	"net/url"
	"regexp"
	"strings"
)

func JoinProxyTargetPath(basePath, requestPath string) string {
	if requestPath == "" {
		requestPath = "/"
	}
	if basePath == "" || basePath == "/" {
		return requestPath
	}
	if strings.HasSuffix(basePath, "/") && strings.HasPrefix(requestPath, "/") {
		return basePath + strings.TrimPrefix(requestPath, "/")
	}
	if !strings.HasSuffix(basePath, "/") && !strings.HasPrefix(requestPath, "/") {
		return basePath + "/" + requestPath
	}
	return basePath + requestPath
}

func RewriteLocation(targetIP string, targetURL *url.URL, location string) string {
	locationURL, err := url.Parse(location)
	if err != nil {
		return location
	}

	proxyPrefix := "/proxy/web/" + targetIP
	if locationURL.Host == "" && strings.HasPrefix(location, "/") {
		return proxyPrefix + location
	}

	if locationURL.Host != "" {
		locationHost := locationURL.Hostname()
		locationPort := locationURL.Port()
		if locationPort == "" && (locationURL.Scheme == "" || locationURL.Scheme == "http") {
			locationPort = "80"
		}

		targetHost := targetURL.Hostname()
		targetPort := targetURL.Port()
		if targetPort == "" && targetURL.Scheme == "http" {
			targetPort = "80"
		}

		if locationHost != targetHost || locationPort != targetPort {
			return location
		}

		rewrittenPath := locationURL.EscapedPath()
		if rewrittenPath == "" {
			rewrittenPath = "/"
		}
		if locationURL.RawQuery != "" {
			rewrittenPath += "?" + locationURL.RawQuery
		}
		if locationURL.Fragment != "" {
			rewrittenPath += "#" + locationURL.Fragment
		}
		return proxyPrefix + rewrittenPath
	}

	return location
}

var (
	webProxyQuotedAttrPattern = regexp.MustCompile(`(?i)\b(href|src|action|poster|data-src|data-href)\s*=\s*(['"])([^'"]*)['"]`)
	webProxyBareAttrPattern   = regexp.MustCompile(`(?i)\b(href|src|action|poster|data-src|data-href)\s*=\s*([^'">\s][^>\s]*)`)
	webProxyCSSURLPattern     = regexp.MustCompile(`(?i)url\(\s*(['"]?)([^'"\)\s]+)['"]?\s*\)`)
	webProxyQuotedURLPattern  = regexp.MustCompile(`(['"])(/[^'"<>\s\\)]*)['"]`)
)

func ShouldRewriteBody(contentType string) bool {
	contentType = strings.ToLower(contentType)
	return strings.Contains(contentType, "text/html") ||
		strings.Contains(contentType, "text/css") ||
		strings.Contains(contentType, "javascript")
}

func RewriteText(body, proxyPrefix string, targetURL *url.URL, contentType string) string {
	contentType = strings.ToLower(contentType)
	isHTML := strings.Contains(contentType, "text/html")
	isScript := strings.Contains(contentType, "javascript")

	if isHTML {
		body = webProxyQuotedAttrPattern.ReplaceAllStringFunc(body, func(match string) string {
			parts := webProxyQuotedAttrPattern.FindStringSubmatch(match)
			if len(parts) != 4 {
				return match
			}
			rewritten := rewriteURL(parts[3], proxyPrefix, targetURL)
			return strings.Replace(match, parts[2]+parts[3]+parts[2], parts[2]+rewritten+parts[2], 1)
		})

		body = webProxyBareAttrPattern.ReplaceAllStringFunc(body, func(match string) string {
			parts := webProxyBareAttrPattern.FindStringSubmatch(match)
			if len(parts) != 3 {
				return match
			}
			rewritten := rewriteURL(parts[2], proxyPrefix, targetURL)
			return strings.Replace(match, parts[2], rewritten, 1)
		})
	}

	if isHTML || isScript {
		body = webProxyQuotedURLPattern.ReplaceAllStringFunc(body, func(match string) string {
			parts := webProxyQuotedURLPattern.FindStringSubmatch(match)
			if len(parts) != 3 {
				return match
			}
			rewritten := rewriteURL(parts[2], proxyPrefix, targetURL)
			return parts[1] + rewritten + parts[1]
		})
	}

	body = webProxyCSSURLPattern.ReplaceAllStringFunc(body, func(match string) string {
		parts := webProxyCSSURLPattern.FindStringSubmatch(match)
		if len(parts) != 3 {
			return match
		}
		rewritten := rewriteURL(parts[2], proxyPrefix, targetURL)
		return "url(" + parts[1] + rewritten + parts[1] + ")"
	})

	if isHTML {
		body = injectRuntime(body, proxyPrefix)
	}

	return body
}

func injectRuntime(body, proxyPrefix string) string {
	if strings.Contains(body, "data-web-proxy-runtime") {
		return body
	}

	script := `<script data-web-proxy-runtime>(function(){var p="` + proxyPrefix + `";var d=["/ISAPI","/SDK","/PSIA","/doc","/webSocket"];function q(x){if(x.indexOf(p+"/")===0||x.indexOf("/proxy/web/")===0){return x}for(var i=0;i<d.length;i++){if(x===d[i]||x.indexOf(d[i]+"/")===0||x.indexOf(d[i]+"?")===0){return p+x}}return x}function r(u){if(typeof u!=="string"){return u}if(u.charAt(0)==="/"&&u.indexOf("//")!==0){return q(u)}try{var a=new URL(u,window.location.href);if(a.origin===window.location.origin){var x=a.pathname+a.search+a.hash;var y=q(x);if(y!==x){return y}}}catch(e){}return u}if(window.XMLHttpRequest){var o=XMLHttpRequest.prototype.open;XMLHttpRequest.prototype.open=function(m,u){arguments[1]=r(u);return o.apply(this,arguments)}}if(window.fetch){var f=window.fetch;window.fetch=function(i,n){if(typeof i==="string"){i=r(i)}else if(i&&i.url){i=new Request(r(i.url),i)}return f.call(this,i,n)}}function a(e){if(!e||!e.getAttribute){return}["src","href","action","data-src","data-href"].forEach(function(k){var v=e.getAttribute(k);if(v){var nv=r(v);if(nv!==v){e.setAttribute(k,nv)}}})}if(window.MutationObserver){new MutationObserver(function(ms){ms.forEach(function(m){if(m.type==="attributes"){a(m.target)}else{Array.prototype.forEach.call(m.addedNodes,function(n){a(n);if(n&&n.querySelectorAll){Array.prototype.forEach.call(n.querySelectorAll("[src],[href],[action],[data-src],[data-href]"),a)}})}})}).observe(document.documentElement,{childList:true,subtree:true,attributes:true,attributeFilter:["src","href","action","data-src","data-href"]})}})();</script>`

	lower := strings.ToLower(body)
	if idx := strings.Index(lower, "</head>"); idx >= 0 {
		return body[:idx] + script + body[idx:]
	}
	if idx := strings.Index(lower, "<body"); idx >= 0 {
		if end := strings.Index(body[idx:], ">"); end >= 0 {
			insertAt := idx + end + 1
			return body[:insertAt] + script + body[insertAt:]
		}
	}
	return script + body
}

func rewriteURL(rawURL, proxyPrefix string, targetURL *url.URL) string {
	rawURL = strings.TrimSpace(rawURL)
	if rawURL == "" ||
		strings.HasPrefix(rawURL, "#") ||
		strings.HasPrefix(rawURL, "//") ||
		strings.HasPrefix(rawURL, proxyPrefix+"/") ||
		strings.HasPrefix(rawURL, "/proxy/web/") {
		return rawURL
	}

	lower := strings.ToLower(rawURL)
	for _, prefix := range []string{"data:", "blob:", "mailto:", "tel:", "javascript:"} {
		if strings.HasPrefix(lower, prefix) {
			return rawURL
		}
	}

	if strings.HasPrefix(rawURL, "/") {
		return proxyPrefix + rawURL
	}

	parsed, err := url.Parse(rawURL)
	if err != nil || parsed.Host == "" {
		return rawURL
	}
	if targetURL == nil || !sameUpstreamHost(parsed, targetURL) {
		return rawURL
	}

	rewrittenPath := parsed.EscapedPath()
	if rewrittenPath == "" {
		rewrittenPath = "/"
	}
	if parsed.RawQuery != "" {
		rewrittenPath += "?" + parsed.RawQuery
	}
	if parsed.Fragment != "" {
		rewrittenPath += "#" + parsed.Fragment
	}
	return proxyPrefix + rewrittenPath
}

func sameUpstreamHost(left, right *url.URL) bool {
	leftPort := left.Port()
	if leftPort == "" && (left.Scheme == "" || left.Scheme == "http") {
		leftPort = "80"
	}
	if leftPort == "" && left.Scheme == "https" {
		leftPort = "443"
	}

	rightPort := right.Port()
	if rightPort == "" && (right.Scheme == "" || right.Scheme == "http") {
		rightPort = "80"
	}
	if rightPort == "" && right.Scheme == "https" {
		rightPort = "443"
	}

	return strings.EqualFold(left.Hostname(), right.Hostname()) && leftPort == rightPort
}

func RewriteSetCookie(cookie, proxyPrefix string) string {
	parts := strings.Split(cookie, ";")
	if len(parts) == 0 {
		return cookie
	}

	rewritten := []string{strings.TrimSpace(parts[0])}
	hasPath := false
	for _, part := range parts[1:] {
		attr := strings.TrimSpace(part)
		if attr == "" {
			continue
		}
		lower := strings.ToLower(attr)
		switch {
		case strings.HasPrefix(lower, "domain="):
			continue
		case strings.HasPrefix(lower, "path="):
			hasPath = true
			rewritten = append(rewritten, "Path="+proxyPrefix+"/")
		default:
			rewritten = append(rewritten, attr)
		}
	}
	if !hasPath {
		rewritten = append(rewritten, "Path="+proxyPrefix+"/")
	}
	return strings.Join(rewritten, "; ")
}