EdgeNode/internal/nodes/http_request_reverse_proxy.go

package nodes

import (
	"context"
	"errors"
	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/shared"
	"github.com/TeaOSLab/EdgeNode/internal/utils"
	"github.com/iwind/TeaGo/logs"
	"io"
	"net/url"
	"strconv"
	"strings"
)

// 处理反向代理
func (this *HTTPRequest) doReverseProxy() {
	if this.reverseProxy == nil {
		return
	}

	// 对URL的处理
	stripPrefix := this.reverseProxy.StripPrefix
	requestURI := this.reverseProxy.RequestURI
	requestURIHasVariables := this.reverseProxy.RequestURIHasVariables()
	requestHost := this.reverseProxy.RequestHost
	requestHostHasVariables := this.reverseProxy.RequestHostHasVariables()

	// 源站
	requestCall := shared.NewRequestCall()
	origin := this.reverseProxy.NextOrigin(requestCall)
	if origin == nil {
		err := errors.New(this.requestPath() + ": no available backends for reverse proxy")
		logs.Error(err)
		this.write500(err)
		return
	}
	this.origin = origin // 设置全局变量是为了日志等处理
	if len(origin.StripPrefix) > 0 {
		stripPrefix = origin.StripPrefix
	}
	if len(origin.RequestURI) > 0 {
		requestURI = origin.RequestURI
		requestURIHasVariables = origin.RequestURIHasVariables()
	}
	if len(origin.RequestHost) > 0 {
		requestHost = origin.RequestHost
		requestHostHasVariables = origin.RequestHostHasVariables()
	}

	// 处理Scheme
	if origin.Addr == nil {
		err := errors.New(this.requestPath() + ": origin '" + strconv.FormatInt(origin.Id, 10) + "' does not has a address")
		logs.Error(err)
		this.write500(err)
		return
	}
	this.RawReq.URL.Scheme = origin.Addr.Protocol.Primary().Scheme()

	// StripPrefix
	if len(stripPrefix) > 0 {
		if stripPrefix[0] != '/' {
			stripPrefix = "/" + stripPrefix
		}
		this.uri = strings.TrimPrefix(this.uri, stripPrefix)
		if len(this.uri) == 0 || this.uri[0] != '/' {
			this.uri = "/" + this.uri
		}
	}

	// RequestURI
	if len(requestURI) > 0 {
		if requestURIHasVariables {
			this.uri = this.Format(requestURI)
		} else {
			this.uri = requestURI
		}
		if len(this.uri) == 0 || this.uri[0] != '/' {
			this.uri = "/" + this.uri
		}

		// 处理RequestURI中的问号
		questionMark := strings.LastIndex(this.uri, "?")
		if questionMark > 0 {
			path := this.uri[:questionMark]
			if strings.Contains(path, "?") {
				this.uri = path + "&" + this.uri[questionMark+1:]
			}
		}

		// 去除多个/
		this.uri = utils.CleanPath(this.uri)
	}

	// RequestHost
	if len(requestHost) > 0 {
		if requestHostHasVariables {
			this.RawReq.Host = this.Format(requestHost)
		} else {
			this.RawReq.Host = this.reverseProxy.RequestHost
		}
		this.RawReq.URL.Host = this.RawReq.Host
	} else {
		this.RawReq.URL.Host = this.Host
	}

	// 重组请求URL
	questionMark := strings.Index(this.uri, "?")
	if questionMark > -1 {
		this.RawReq.URL.Path = this.uri[:questionMark]
		this.RawReq.URL.RawQuery = this.uri[questionMark+1:]
	} else {
		this.RawReq.URL.Path = this.uri
		this.RawReq.URL.RawQuery = ""
	}
	this.RawReq.RequestURI = ""

	// 处理Header
	this.setForwardHeaders(this.RawReq.Header)
	this.processRequestHeaders(this.RawReq.Header)

	// 判断是否为Websocket请求
	if this.RawReq.Header.Get("Upgrade") == "websocket" {
		this.doWebsocket()
		return
	}

	// 获取请求客户端
	client, addr, err := SharedHTTPClientPool.Client(this, origin)
	if err != nil {
		logs.Error(err)
		this.write500(err)
		return
	}

	this.originAddr = addr

	// 开始请求
	resp, err := client.Do(this.RawReq)
	if err != nil {
		// 客户端取消请求，则不提示
		httpErr, ok := err.(*url.Error)
		if !ok || httpErr.Err != context.Canceled {
			// TODO 如果超过最大失败次数，则下线

			this.write500(err)
			logs.Println("[proxy]'" + this.RawReq.URL.String() + "': " + err.Error())
		} else {
			// 是否为客户端方面的错误
			isClientError := false
			if ok {
				if httpErr.Err == context.Canceled {
					isClientError = true
					this.addError(errors.New(httpErr.Op + " " + httpErr.URL + ": client closed the connection"))
					this.writer.WriteHeader(499) // 仿照nginx
				}
			}

			if !isClientError {
				this.write500(err)
			}
		}
		if resp != nil && resp.Body != nil {
			_ = resp.Body.Close()
		}
		return
	}

	// WAF对出站进行检查
	// TODO

	// TODO 清除源站错误次数

	// 特殊页面
	// TODO

	// 设置Charset
	// TODO 这里应该可以设置文本类型的列表，以及是否强制覆盖所有文本类型的字符集
	if this.web.Charset != nil && this.web.Charset.IsOn && len(this.web.Charset.Charset) > 0 {
		contentTypes, ok := resp.Header["Content-Type"]
		if ok && len(contentTypes) > 0 {
			contentType := contentTypes[0]
			if _, found := textMimeMap[contentType]; found {
				resp.Header["Content-Type"][0] = contentType + "; charset=" + this.web.Charset.Charset
			}
		}
	}

	// 响应Header
	this.writer.AddHeaders(resp.Header)
	this.processResponseHeaders(resp.StatusCode)

	// 是否需要刷新
	shouldAutoFlush := this.reverseProxy.AutoFlush || this.RawReq.Header.Get("Accept") == "text/event-stream"

	// 准备
	this.writer.Prepare(resp.ContentLength)

	// 设置响应代码
	this.writer.WriteHeader(resp.StatusCode)

	// 输出到客户端
	pool := this.bytePool(resp.ContentLength)
	buf := pool.Get()
	if shouldAutoFlush {
		for {
			n, readErr := resp.Body.Read(buf)
			if n > 0 {
				_, err = this.writer.Write(buf[:n])
				this.writer.Flush()
				if err != nil {
					break
				}
			}
			if readErr != nil {
				err = readErr
				break
			}
		}
	} else {
		_, err = io.CopyBuffer(this.writer, resp.Body, buf)
	}
	pool.Put(buf)

	err1 := resp.Body.Close()
	if err1 != nil {
		logs.Error(err1)
	}

	if err != nil && err != io.EOF {
		logs.Error(err)
		this.addError(err)
	}
}