managed-portal/internal/webdevice/proxy.go

package webdevice

import (
	"errors"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"
	"time"
)

var (
	ErrInvalidTargetIP  = errors.New("invalid target ip")
	ErrTargetNotAllowed = errors.New("target ip not allowed")
	ErrInvalidProxyURL  = errors.New("invalid proxy target")
)

func (s *Service) ProxyHTTP(w http.ResponseWriter, r *http.Request, targetIP, proxyPath string) error {
	if !IsPrivateIPv4Literal(targetIP) {
		return ErrInvalidTargetIP
	}
	if !s.IsAllowed(targetIP) {
		return ErrTargetNotAllowed
	}

	rawTarget := s.ProxyTargetURL(targetIP)
	targetURL, err := url.Parse(rawTarget)
	if err != nil || targetURL.Scheme == "" || targetURL.Host == "" {
		return ErrInvalidProxyURL
	}

	if proxyPath == "" {
		proxyPath = "/"
	}
	rawQuery := r.URL.RawQuery

	proxy := &httputil.ReverseProxy{
		Transport: retryTransport{
			base: &http.Transport{
				DisableKeepAlives:     true,
				ResponseHeaderTimeout: 5 * time.Second,
			},
			attempts: 3,
			delay:    80 * time.Millisecond,
		},
		Director: func(req *http.Request) {
			req.URL.Scheme = targetURL.Scheme
			req.URL.Host = targetURL.Host
			req.URL.Path = JoinProxyTargetPath(targetURL.Path, proxyPath)
			req.URL.RawPath = ""
			req.URL.RawQuery = rawQuery
			req.Host = targetURL.Host
			req.Header = sanitizeProxyRequestHeader(req.Header, req.URL.Path)
			req.Close = true
		},
		ModifyResponse: func(resp *http.Response) error {
			proxyPrefix := "/proxy/web/" + targetIP
			if location := resp.Header.Get("Location"); location != "" {
				resp.Header.Set("Location", RewriteLocation(targetIP, targetURL, location))
			}
			if cookies := resp.Header.Values("Set-Cookie"); len(cookies) > 0 {
				resp.Header.Del("Set-Cookie")
				for _, cookie := range cookies {
					resp.Header.Add("Set-Cookie", RewriteSetCookie(cookie, proxyPrefix))
				}
			}

			contentType := resp.Header.Get("Content-Type")
			if ShouldRewriteBody(contentType) {
				body, err := io.ReadAll(resp.Body)
				if err != nil {
					return err
				}
				_ = resp.Body.Close()

				rewritten := RewriteText(string(body), proxyPrefix, targetURL, contentType)
				rewrittenBytes := []byte(rewritten)
				resp.Body = io.NopCloser(strings.NewReader(rewritten))
				resp.ContentLength = int64(len(rewrittenBytes))
				resp.Header.Set("Content-Length", strconv.Itoa(len(rewrittenBytes)))
				resp.Header.Del("Content-Encoding")
				resp.Header.Del("Content-MD5")
				resp.Header.Del("Etag")
			}
			return nil
		},
		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
			log.Printf("web device proxy failed target=%s path=%s error=%v", targetURL.String(), r.URL.RequestURI(), err)
			http.Error(w, "代理访问失败: "+err.Error(), http.StatusBadGateway)
		},
	}

	proxy.ServeHTTP(closeNotifyWriter{ResponseWriter: w}, r)
	return nil
}

type retryTransport struct {
	base     http.RoundTripper
	attempts int
	delay    time.Duration
}

func (t retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	base := t.base
	if base == nil {
		base = http.DefaultTransport
	}
	attempts := t.attempts
	if attempts < 1 {
		attempts = 1
	}

	var lastErr error
	for attempt := 0; attempt < attempts; attempt++ {
		nextReq := req
		if attempt > 0 {
			nextReq = req.Clone(req.Context())
			nextReq.Header = req.Header.Clone()
			nextReq.Close = true
		}

		resp, err := base.RoundTrip(nextReq)
		if err == nil {
			return resp, nil
		}
		lastErr = err
		if !shouldRetryProxyRequest(req, err) || attempt == attempts-1 {
			return nil, err
		}
		time.Sleep(t.delay)
	}
	return nil, lastErr
}

func shouldRetryProxyRequest(req *http.Request, err error) bool {
	if req == nil || err == nil {
		return false
	}
	switch req.Method {
	case http.MethodGet, http.MethodHead, http.MethodOptions:
	default:
		return false
	}
	message := strings.ToLower(err.Error())
	return strings.Contains(message, "eof") ||
		strings.Contains(message, "connection reset") ||
		strings.Contains(message, "broken pipe")
}

func sanitizeProxyRequestHeader(source http.Header, upstreamPath string) http.Header {
	header := make(http.Header)
	copyHeaderValue(header, source, "Accept")
	copyHeaderValue(header, source, "Content-Type")
	copyHeaderValue(header, source, "Authorization")

	userAgent := strings.TrimSpace(source.Get("User-Agent"))
	if userAgent == "" {
		userAgent = "Mozilla/5.0"
	}
	header.Set("User-Agent", userAgent)
	header.Set("Connection", "close")

	if !isLoginPagePath(upstreamPath) {
		copyHeaderValue(header, source, "Cookie")
	}
	return header
}

func copyHeaderValue(target, source http.Header, key string) {
	if value := source.Values(key); len(value) > 0 {
		target.Del(key)
		for _, item := range value {
			target.Add(key, item)
		}
	}
}

func isLoginPagePath(path string) bool {
	path = strings.ToLower(path)
	return strings.HasSuffix(path, "/doc/page/login.asp") || path == "/doc/page/login.asp"
}

type closeNotifyWriter struct {
	http.ResponseWriter
}

func (w closeNotifyWriter) CloseNotify() <-chan bool {
	ch := make(chan bool, 1)
	return ch
}

func (w closeNotifyWriter) Flush() {
	if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
		flusher.Flush()
	}
}