fix: proxy web device resources transparently

internal/webdevice/proxy.go

@@ -1,6 +1,7 @@
 package webdevice
 
 import (
+	"context"
 	"errors"
 	"io"
 	"log"
@@ -18,6 +19,13 @@ var (
 	ErrInvalidProxyURL  = errors.New("invalid proxy target")
 )
 
+const (
+	webDeviceProxyConcurrency = 1
+	webDeviceProxyAttempts    = 6
+	webDeviceProxyRetryDelay  = 150 * time.Millisecond
+	webDeviceProxyQueueWait   = 8 * time.Second
+)
+
 func (s *Service) ProxyHTTP(w http.ResponseWriter, r *http.Request, targetIP, proxyPath string) error {
 	if !IsPrivateIPv4Literal(targetIP) {
 		return ErrInvalidTargetIP
@@ -36,6 +44,7 @@ func (s *Service) ProxyHTTP(w http.ResponseWriter, r *http.Request, targetIP, pr
 		proxyPath = "/"
 	}
 	rawQuery := r.URL.RawQuery
+	upgradeRequest := isUpgradeRequest(r.Header)
 
 	proxy := &httputil.ReverseProxy{
 		Transport: retryTransport{
@@ -43,8 +52,10 @@ func (s *Service) ProxyHTTP(w http.ResponseWriter, r *http.Request, targetIP, pr
 				DisableKeepAlives:     true,
 				ResponseHeaderTimeout: 5 * time.Second,
 			},
-			attempts: 3,
-			delay:    80 * time.Millisecond,
+			limiter:        s.proxyLimiter(targetIP),
+			attempts:       webDeviceProxyAttempts,
+			delay:          webDeviceProxyRetryDelay,
+			acquireTimeout: webDeviceProxyQueueWait,
 		},
 		Director: func(req *http.Request) {
 			req.URL.Scheme = targetURL.Scheme
@@ -54,7 +65,7 @@ func (s *Service) ProxyHTTP(w http.ResponseWriter, r *http.Request, targetIP, pr
 			req.URL.RawQuery = rawQuery
 			req.Host = targetURL.Host
 			req.Header = sanitizeProxyRequestHeader(req.Header, req.URL.Path)
-			req.Close = true
+			req.Close = !upgradeRequest
 		},
 		ModifyResponse: func(resp *http.Response) error {
 			proxyPrefix := "/proxy/web/" + targetIP
@@ -98,9 +109,11 @@ func (s *Service) ProxyHTTP(w http.ResponseWriter, r *http.Request, targetIP, pr
 }
 
 type retryTransport struct {
-	base     http.RoundTripper
-	attempts int
-	delay    time.Duration
+	base           http.RoundTripper
+	limiter        chan struct{}
+	attempts       int
+	delay          time.Duration
+	acquireTimeout time.Duration
 }
 
 func (t retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
@@ -119,10 +132,14 @@ func (t retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 		if attempt > 0 {
 			nextReq = req.Clone(req.Context())
 			nextReq.Header = req.Header.Clone()
-			nextReq.Close = true
+			nextReq.Close = req.Close
 		}
 
+		if err := acquireProxySlot(req.Context(), t.limiter, t.acquireTimeout); err != nil {
+			return nil, err
+		}
 		resp, err := base.RoundTrip(nextReq)
+		releaseProxySlot(t.limiter)
 		if err == nil {
 			return resp, nil
 		}
@@ -130,11 +147,43 @@ func (t retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 		if !shouldRetryProxyRequest(req, err) || attempt == attempts-1 {
 			return nil, err
 		}
-		time.Sleep(t.delay)
+		time.Sleep(t.delay * time.Duration(attempt+1))
 	}
 	return nil, lastErr
 }
 
+func acquireProxySlot(ctx context.Context, limiter chan struct{}, timeout time.Duration) error {
+	if limiter == nil {
+		return nil
+	}
+	if timeout <= 0 {
+		select {
+		case limiter <- struct{}{}:
+			return nil
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}
+
+	timer := time.NewTimer(timeout)
+	defer timer.Stop()
+	select {
+	case limiter <- struct{}{}:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.C:
+		return context.DeadlineExceeded
+	}
+}
+
+func releaseProxySlot(limiter chan struct{}) {
+	if limiter == nil {
+		return
+	}
+	<-limiter
+}
+
 func shouldRetryProxyRequest(req *http.Request, err error) bool {
 	if req == nil || err == nil {
 		return false
@@ -151,19 +200,25 @@ func shouldRetryProxyRequest(req *http.Request, err error) bool {
 }
 
 func sanitizeProxyRequestHeader(source http.Header, upstreamPath string) http.Header {
+	upgradeRequest := isUpgradeRequest(source)
 	header := source.Clone()
 	for key := range header {
-		if isProxyManagedHeader(key) {
+		if isProxyManagedHeader(key, upgradeRequest) {
 			header.Del(key)
 		}
 	}
+	header.Del("Accept-Encoding")
 
 	userAgent := strings.TrimSpace(source.Get("User-Agent"))
 	if userAgent == "" {
 		userAgent = "Mozilla/5.0"
 	}
 	header.Set("User-Agent", userAgent)
-	header.Set("Connection", "close")
+	if upgradeRequest {
+		header.Set("Connection", "Upgrade")
+	} else {
+		header.Set("Connection", "close")
+	}
 
 	if !isLoginPagePath(upstreamPath) {
 		return header
@@ -173,13 +228,12 @@ func sanitizeProxyRequestHeader(source http.Header, upstreamPath string) http.He
 	return header
 }
 
-func isProxyManagedHeader(key string) bool {
+func isProxyManagedHeader(key string, upgradeRequest bool) bool {
 	switch http.CanonicalHeaderKey(key) {
 	case "Connection",
 		"Proxy-Connection",
 		"Keep-Alive",
 		"Transfer-Encoding",
-		"Upgrade",
 		"Te",
 		"Trailer",
 		"Proxy-Authenticate",
@@ -190,11 +244,21 @@ func isProxyManagedHeader(key string) bool {
 		"X-Forwarded-Proto",
 		"X-Real-Ip":
 		return true
+	case "Upgrade":
+		return !upgradeRequest
 	default:
 		return false
 	}
 }
 
+func isUpgradeRequest(header http.Header) bool {
+	if header == nil {
+		return false
+	}
+	connection := strings.ToLower(header.Get("Connection"))
+	return header.Get("Upgrade") != "" && strings.Contains(connection, "upgrade")
+}
+
 func isLoginPagePath(path string) bool {
 	path = strings.ToLower(path)
 	return strings.HasSuffix(path, "/doc/page/login.asp") || path == "/doc/page/login.asp"