go-proxy/proxy.go

122 lines
2.7 KiB
Go

package proxy
import (
"fmt"
"net/http"
"net/http/httputil"
"net/url"
"sync"
"forge.cadoles.com/Cadoles/go-proxy/util"
"github.com/pkg/errors"
)
type Proxy struct {
reversers sync.Map
handler http.Handler
proxyResponseTransformer ProxyResponseTransformer
proxyRequestTransformer ProxyRequestTransformer
}
// ServeHTTP implements http.Handler
func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
p.handler.ServeHTTP(w, r)
}
func (p *Proxy) proxyRequest(w http.ResponseWriter, r *http.Request) {
var reverser *httputil.ReverseProxy
key := fmt.Sprintf("%s://%s", r.URL.Scheme, r.URL.Host)
createAndStore := func() {
target := &url.URL{
Scheme: r.URL.Scheme,
Host: r.URL.Host,
}
reverser = httputil.NewSingleHostReverseProxy(target)
originalDirector := reverser.Director
if p.proxyRequestTransformer != nil {
reverser.Director = func(r *http.Request) {
originalDirector(r)
p.proxyRequestTransformer.TransformRequest(r)
}
}
if p.proxyResponseTransformer != nil {
reverser.ModifyResponse = func(r *http.Response) error {
if err := p.proxyResponseTransformer.TransformResponse(r); err != nil {
return errors.WithStack(err)
}
return nil
}
}
p.reversers.Store(key, reverser)
}
raw, exists := p.reversers.Load(key)
if !exists {
createAndStore()
}
reverser, ok := raw.(*httputil.ReverseProxy)
if !ok {
createAndStore()
}
reverser.ServeHTTP(w, r)
}
func New(funcs ...OptionFunc) *Proxy {
opts := defaultOptions()
for _, fn := range funcs {
fn(opts)
}
proxy := &Proxy{}
handler := http.HandlerFunc(proxy.proxyRequest)
proxy.handler = createMiddlewareChain(handler, opts.Middlewares)
proxy.proxyRequestTransformer = createProxyRequestChain(&defaultProxyRequestTransformer{}, opts.ProxyRequestMiddlewares)
proxy.proxyResponseTransformer = createProxyResponseChain(&defaultProxyResponseTransformer{}, opts.ProxyResponseMiddlewares)
return proxy
}
var _ http.Handler = &Proxy{}
func createMiddlewareChain(handler http.Handler, middlewares []Middleware) http.Handler {
util.Reverse(middlewares)
for _, m := range middlewares {
handler = m(handler)
}
return handler
}
func createProxyResponseChain(transformer ProxyResponseTransformer, middlewares []ProxyResponseMiddleware) ProxyResponseTransformer {
util.Reverse(middlewares)
for _, m := range middlewares {
transformer = m(transformer)
}
return transformer
}
func createProxyRequestChain(transformer ProxyRequestTransformer, middlewares []ProxyRequestMiddleware) ProxyRequestTransformer {
util.Reverse(middlewares)
for _, m := range middlewares {
transformer = m(transformer)
}
return transformer
}