Wormhole/internal/proxy/forward_pipeline.go

package proxy

import (
	"bytes"
	"io"
	"net/http"

	"github.com/wormhole/wormhole/internal/forwardplugin"
	"github.com/wormhole/wormhole/internal/store"
)

func pluginEntriesFromStore(plugins store.ForwardPlugins) []forwardplugin.PluginEntry {
	if len(plugins) == 0 {
		return nil
	}
	out := make([]forwardplugin.PluginEntry, len(plugins))
	for i, p := range plugins {
		out[i] = forwardplugin.PluginEntry{
			Type:    p.Type,
			Enabled: p.Enabled,
			Config:  p.Config,
		}
	}
	return out
}

func runForwardPluginPipeline(
	w http.ResponseWriter,
	clientReq *http.Request,
	df *store.DomainForward,
	upstreamReq *http.Request,
	transport http.RoundTripper,
) error {
	chain, err := forwardplugin.BuildChain(pluginEntriesFromStore(df.Plugins))
	if err != nil {
		return err
	}

	forwardplugin.BindForwardID(chain, df.ID)

	ctx := &forwardplugin.Context{
		ForwardID: df.ID,
		ClientReq: clientReq,
	}

	outReq := upstreamReq.Clone(clientReq.Context())
	if chain.NeedsBody() {
		outReq.Header.Del("Accept-Encoding")
	}

	short, err := chain.HandleRequest(ctx, outReq)
	if err != nil {
		return err
	}

	var resp *forwardplugin.BufferedResponse
	if short && ctx.CachedResp != nil {
		resp = ctx.CachedResp.Clone()
	} else {
		resp, err = fetchUpstreamResponse(outReq, transport, chain.NeedsBody())
		if err != nil {
			return err
		}
		chain.StoreRawResponses(ctx, resp)
	}

	if err := chain.HandleResponse(ctx, resp); err != nil {
		return err
	}
	writeBufferedResponse(w, resp)
	return nil
}

func fetchUpstreamResponse(req *http.Request, transport http.RoundTripper, _ bool) (*forwardplugin.BufferedResponse, error) {
	client := &http.Client{Transport: transport}
	upstreamResp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer upstreamResp.Body.Close()

	body, err := io.ReadAll(upstreamResp.Body)
	if err != nil {
		return nil, err
	}
	return &forwardplugin.BufferedResponse{
		StatusCode: upstreamResp.StatusCode,
		Header:     upstreamResp.Header.Clone(),
		Body:       body,
	}, nil
}

func writeBufferedResponse(w http.ResponseWriter, resp *forwardplugin.BufferedResponse) {
	for k, vals := range resp.Header {
		for _, v := range vals {
			w.Header().Add(k, v)
		}
	}
	if resp.StatusCode == 0 {
		resp.StatusCode = http.StatusOK
	}
	w.WriteHeader(resp.StatusCode)
	if len(resp.Body) > 0 {
		_, _ = io.Copy(w, bytes.NewReader(resp.Body))
	}
}

func hasEnabledPlugins(plugins store.ForwardPlugins) bool {
	for _, p := range plugins {
		if p.Enabled {
			return true
		}
	}
	return false
}