Golang实现反向代理的示例代码

// 设置要转发的地址
target, err := url.Parse("http://www.domain.com")
if err != nil {
    panic(err)
}
// 实例化 ReverseProxy 包
proxy := httputil.NewSingleHostReverseProxy(target)
//http.HandleFunc("/", proxy.ServeHTTP)
// 启动服务
log.Fatal(http.ListenAndServe(":8082", proxy))

type ReverseProxy struct {
	// 对请求内容进行修改 (对象是业务传入req的一个副本)
	Director func(*http.Request)
	// 连接池复用连接，用于执行请求, 默认为http.DefaultTransport 
	Transport http.RoundTripper
	// 定时刷新内容到客户端的时间间隔(流式/无内容此参数忽略)
	FlushInterval time.Duration
	// 默认为std.err，用于记录内部错误日志
	ErrorLog *log.Logger
	// 用于执行 copyBuffer 复制响应体时，利用的bytes内存池化
	BufferPool BufferPool
	// 如果配置后, 可修改目标代理的响应结果(响应头和内容)
    // 如果此方法返回error, 将调用 ErrorHandler 方法
	ModifyResponse func(*http.Response) error
    // 配置后代理执行过程中, 发生错误均会回调此方法
    // 默认逻辑不响应任务内容, 状态码返回502
	ErrorHandler func(http.ResponseWriter, *http.Request, error)
}

// 实例化 ReverseProxy 包
proxy := httputil.NewSingleHostReverseProxy(target)

// 实例化 ReverseProxy 对象
// 初始化 Director 对象, 将请求地址转换为代理目标地址.
// 对请求header头进行处理
func NewSingleHostReverseProxy(target *url.URL) *ReverseProxy {
	targetQuery := target.RawQuery
	director := func(req *http.Request) {
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		req.URL.Path, req.URL.RawPath = joinURLPath(target, req.URL)
		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
		if _, ok := req.Header["User-Agent"]; !ok {
			// explicitly disable User-Agent so it's not set to default value
			req.Header.Set("User-Agent", "")
		}
	}
	return &ReverseProxy{Director: director}
}

type Handler interface {
	ServeHTTP(ResponseWriter, *Request)
}

// 服务请求处理方法
func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// 检测是否设置http.Transport对象
    // 如果未设置则使用默认对象
    transport := p.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}
    // 检测请求是否被终止
    // 终止请求或是正常结束请求等 notifyChan 都会收到请求结束通知, 之后进行cancel
	ctx := req.Context()
	if cn, ok := rw.(http.CloseNotifier); ok {
		var cancel context.CancelFunc
		ctx, cancel = context.WithCancel(ctx)
		defer cancel()
		notifyChan := cn.CloseNotify()
		go func() {
			select {
			case <-notifyChan:
				cancel()
			case <-ctx.Done():
			}
		}()
	}
    // 对外部传入的http.Request对象进行克隆
    // outreq 是给代理服务器传入的请求对象
	outreq := req.Clone(ctx)
	if req.ContentLength == 0 {
        // 主要修复 ReverseProxy 与 http.Transport 重试不兼容性问题
        // 如果请求方法为 GET、HEAD、OPTIONS、TRACE, 同时body为nil情况下, 将会发生重试
        // 避免因为复制传入的request创建传入代理的请求内容, 导致无法发生重试.
        // https://github.com/golang/go/issues/16036
		outreq.Body = nil
	}
	if outreq.Body != nil {
		// 避免因panic问题导致请求未正确关闭, 其他协程继续从中读取
        // https://github.com/golang/go/issues/46866
		defer outreq.Body.Close()
	}
	if outreq.Header == nil {
        // Issue 33142: historical behavior was to always allocate
		outreq.Header = make(http.Header)
	}
    // 调用实现的 Director 方法修改请求代理的request对象
	p.Director(outreq)
	if outreq.Form != nil {
		outreq.URL.RawQuery = cleanQueryParams(outreq.URL.RawQuery)
	}
	outreq.Close = false
    // 升级http协议，HTTP Upgrade
    // 判断header Connection 中是否有Upgrade
	reqUpType := upgradeType(outreq.Header)
    // 根据《网络交换的 ASCII 格式》规范, 升级协议中是否包含禁止使用的字符
    // https://datatracker.ietf.org/doc/html/rfc20#section-4.2
	if !ascii.IsPrint(reqUpType) {
        // 调用 ReverseProxy 对象的 ErrorHandler 方法
		p.getErrorHandler()(
            rw, 
        	req, 
         	fmt.Errorf("client tried to switch to invalid protocol %q", reqUpType))
		return
	}
    // 请求下游移除Connetion头
    // https://datatracker.ietf.org/doc/html/rfc7230#section-6.1
	removeConnectionHeaders(outreq.Header)
	// 请求下游根据RFC规范移除协议头
	for _, h := range hopHeaders {
		outreq.Header.Del(h)
	}
	// Transfer-Encoding: chunked 分块传输编码
	if httpguts.HeaderValuesContainsToken(req.Header["Te"], "trailers") {
		outreq.Header.Set("Te", "trailers")
	}
	// 请求下游指定协议升级, 例如 websockeet
	if reqUpType != "" {
		outreq.Header.Set("Connection", "Upgrade")
		outreq.Header.Set("Upgrade", reqUpType)
	}
    // 添加 X-Forwarded-For 头
    // 最开始的是离服务端最远的设备 IP，然后是每一级代理设备的 IP
    // 类似于 X-Forwarded-For: client, proxy1, proxy2
	if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
		prior, ok := outreq.Header["X-Forwarded-For"]
        // 如果header头 X-Forwarded-For 设置为nil, 则不再 X-Forwarded-For
        // 这个参数下面我们将详细说明
		omit := ok && prior == nil
		if len(prior) > 0 {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		if !omit {
			outreq.Header.Set("X-Forwarded-For", clientIP)
		}
	}
    // 使用transport对象中维护的链接池, 向下游发起请求
	res, err := transport.RoundTrip(outreq)
	if err != nil {
		p.getErrorHandler()(rw, outreq, err)
		return
	}
    // 处理下游响应的升级协议请求
	// Deal with 101 Switching Protocols responses: (WebSocket, h2c, etc)
	if res.StatusCode == http.StatusSwitchingProtocols {
		if !p.modifyResponse(rw, res, outreq) {
			return
		}
		p.handleUpgradeResponse(rw, outreq, res)
		return
	}
    // 根据协议规范删除响应 Connection 头
	removeConnectionHeaders(res.Header)
    // 下游响应根据RFC规范移除协议头
	for _, h := range hopHeaders {
		res.Header.Del(h)
	}
    // 如有设置 modifyResponse, 则修改响应内容
    // 调用 ReverseProxy 对象 modifyResponse 方法
	if !p.modifyResponse(rw, res, outreq) {
		return
	}
    // 拷贝响应Header到上游response对象
	copyHeader(rw.Header(), res.Header)
	// 分块传输部分协议 header 头设置, 已跳过
    // 写入响应码到上游response对象
	rw.WriteHeader(res.StatusCode)
    // 拷贝结果到上游
    // flushInterval将响应定时刷新到缓冲区
	err = p.copyResponse(rw, res.Body, p.flushInterval(res))
	if err != nil {
		defer res.Body.Close()
		// ... 调用errorHandler
        panic(http.ErrAbortHandler)
	}
    // 关闭响应body
	res.Body.Close()
    // chunked 分块传输编码调用flush刷新到客户端
	if len(res.Trailer) > 0 {
		// Force chunking if we saw a response trailer.
		// This prevents net/http from calculating the length for short
		// bodies and adding a Content-Length.
		if fl, ok := rw.(http.Flusher); ok {
			fl.Flush()
		}
	}
    // 以下为分块传输编码相关header设置
	if len(res.Trailer) == announcedTrailers {
		copyHeader(rw.Header(), res.Trailer)
		return
	}
	for k, vv := range res.Trailer {
		k = http.TrailerPrefix + k
		for _, v := range vv {
			rw.Header().Add(k, v)
		}
	}
}

type ServeHTTP struct {
    // 代理链接地址
	targetUrl         string
    // net/http 内置的 ReverseProxy 对象
	reverseProxy      *httputil.ReverseProxy
    // 代理错误处理
	proxyErrorHandler ProxyErrorHandler
    // 日志对象
	logger            log.Logger
}

// NewServeHTTP 初始化代理对象
func NewServeHTTP(targetUrl string, logger log.Logger) *ServeHTTP {
	target, err := url.Parse(targetUrl)
	if err != nil {
		panic(err)
	}
    // 重新设置 Director 复制请求处理
	proxy := &httputil.ReverseProxy{Director: func(req *http.Request) {
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		req.Host = target.Host
		if _, ok := req.Header["User-Agent"]; !ok {
			req.Header.Set("User-Agent", "")
		}
		if req.Header.Get("Content-Length") == "0" {
			req.Header.Del("Content-Length")
		}
		req.Header["X-Forwarded-For"] = nil
		for _, name := range removeRequestHeaders {
			req.Header.Del(name)
		}
	}}
	serveHttp := &ServeHTTP{
		targetUrl:         targetUrl,
		logger:            logger,
		reverseProxy:      proxy,
		proxyErrorHandler: DefaultProxyErrorHandler,
	}
    // 设置trasport处理对象(主要调配链接池大小和超时时间)
	serveHttp.reverseProxy.Transport = HttpTransportDefault()
    // 定义错误处理
	serveHttp.reverseProxy.ErrorHandler = serveHttp.getErrorHandler(logger)
    // 定义响应处理
	serveHttp.reverseProxy.ModifyResponse = serveHttp.getResponseHandler(logger)
	return serveHttp
}
// SetProxyErrorFunc 设置错误处理函数
func (s *ServeHTTP) SetProxyErrorFunc(handler ProxyErrorHandler) *ServeHTTP {
	s.proxyErrorHandler = handler
	return s
}

netHttp.Transport{
		Proxy: proxyURL,
		DialContext: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		ForceAttemptHTTP2:     true,
		MaxIdleConns:          200,
		MaxIdleConnsPerHost:   200,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}

// ServeHTTP 服务转发
func (s *ServeHTTP) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	var (
		reqBody []byte
		err     error
        // 生成traceId
		traceId = s.getTraceId(request)
	)
    // 前置获取请求头, 放入context中
    // 调用结束后请求 body 将会被关闭, 后面将无法再获取
	if request.Body != nil {
		reqBody, err = io.ReadAll(request.Body)
		if err == nil {
			request.Body = io.NopCloser(bytes.NewBuffer(reqBody))
		}
	}
    // header 设置 traceId和超时时间传递
	request.Header.Set(utils.TraceKey, traceId)
	request.Header.Set(utils.Timeoutkey, cast.ToString(s.getTimeout(request)))
	// 计算获取超时时间, 发起转发请求
	ctx, cancel := context.WithTimeout(
        request.Context(), 
         time.Duration(s.getTimeout(request))*time.Millisecond,
    )
	defer cancel()
    // 设置请求体
	ctx = context.WithValue(ctx, ctxReqBody, string(reqBody))
    // 设置请求时间, 用于响应结束后计算请求耗时
	ctx = context.WithValue(ctx, ctxReqTime, time.Now())
    // context 设置traceId, 用于链路日志打印
	ctx = context.WithValue(ctx, utils.TraceKey, traceId)
	request = request.WithContext(ctx)
    // 调用 reverseProxy ServeHTTP, 处理转发逻辑
	s.reverseProxy.ServeHTTP(writer, request)
}

// getTraceId 获取traceId
// header头中不存在则生成
func (s *ServeHTTP) getTraceId(request *http.Request) string {
	traceId := request.Header.Get(utils.TraceKey)
	if traceId != "" {
		return traceId
	}
	return uuid.NewV4().String()
}
// getTimeout 获取超时时间
// header中不存在timeoutKey, 返回默认超时时间
// header头存在, 则判断是否大于默认超时时间, 大于则使用默认超时时间
// 否则返回header设置的超时时间
func (s *ServeHTTP) getTimeout(request *http.Request) uint32 {
	timeout := request.Header.Get(utils.Timeoutkey)
	if timeout == "" {
		return DefaultTimeoutMs
	}
	headerTimeoutMs := cast.ToUint32(timeout)
	if headerTimeoutMs > DefaultTimeoutMs {
		return DefaultTimeoutMs
	}
	return cast.ToUint32(timeout)
}

// getErrorHandler 记录错误记录
func (s *ServeHTTP) getErrorHandler(logger log.Logger) ErrorHandler {
	return func(writer http.ResponseWriter, request *http.Request, e error) {
		var (
			reqBody []byte
			err     error
		)
		if request.Body != nil {
			reqBody, err = io.ReadAll(request.Body)
			if err == nil {
				request.Body = io.NopCloser(bytes.NewBuffer(reqBody))
			}
		}
        // 初始化时确认proxyErrorHandler具体处理方法
        // 调用 proxyErrorHandler,处理响应部分
		s.proxyErrorHandler(writer, e)
        // 获取必要信息, 记录错误日志
		scheme := s.getSchemeDataByRequest(request)
		_ = log.WithContext(request.Context(), logger).Log(log.LevelError,
			"x_module", "proxy/server/error",
			"x_component", scheme.kind,
			"x_error", e,
			"x_header", request.Header,
			"x_action", scheme.operation,
			"x_param", string(reqBody),
			"x_trace_id", request.Context().Value(utils.TraceKey),
		)
	}
}
// 具体代理业务错误处理
// 包含默认错误响应和具体代理业务错误响应. 
// 以下为某个业务响应
func XXXProxyErrorHandler(writer http.ResponseWriter, err error) {
	resp := HttpXXXResponse{
		ErrCode: 1,
		ErrMsg:  err.Error(),
		Data:    struct{}{},
	}
	writer.Header().Set("Content-Type", "application/json; charset=utf-8")
	writer.Header().Set("Connection", "keep-alive")
	writer.Header().Set("Cache-Control", "no-cache")
    // 设置状态码为200
	writer.WriteHeader(http.StatusOK)
    // 将响应值序列化
	respByte, _ := json.Marshal(resp)
    // 将response数据写入writer, 刷新到Flush
    // 关于Flush部分, 一般是不需要主动刷新的, 请求结束后会自动Flush
	_, _ = fmt.Fprintf(writer, string(respByte))
	if f, ok := writer.(http.Flusher); ok {
		f.Flush()
	}
}

// getResponseHandler 获取响应数据
func (s *ServeHTTP) getResponseHandler(logger log.Logger) func(response *http.Response) error {
	return func(response *http.Response) error {
		var (
			duration float64
			logLevel = log.LevelInfo
			header   http.Header
		)
        // 获取请求体
		reqBody := response.Request.Context().Value(ctxReqBody)
        // 获取开始请求时间, 计算请求耗时
		startTime := response.Request.Context().Value(ctxReqBody)
		if startTime != nil {
			_, ok := startTime.(time.Time)
			if ok {
				duration = time.Since(startTime.(time.Time)).Seconds()
			}
		}
        // 获取响应数据
        // 如果响应码非200, 调整日志等级
		scheme := s.getSchemeDataByResponse(response)
		if response.StatusCode != http.StatusOK {
			logLevel = log.LevelError
			header = scheme.header
		}
        // 记录日志
		_ = log.WithContext(response.Request.Context(), logger).Log(logLevel,
			"x_module", "proxy/server/resp",
			"x_component", "http",
			"x_code", scheme.code,
			"x_header", header,
			"x_action", scheme.operation,
			"x_params", reqBody,
			"x_response", scheme.responseData,
			"x_duration", duration,
			"x_trace_id", response.Request.Context().Value(utils.TraceKey),
		)
		// 设置响应头
		response.Header.Set("Content-Type", "application/json; charset=utf-8")
		return nil
	}
}

urlStr := "https://" + targetHost
proxy := utilsProxy.NewServeHTTP(urlStr, logger).SetProxyErrorFunc(utilsProxy.XXXProxyErrorHandler)
log.Fatal(http.ListenAndServe(":8082", proxy))