reverse_proxy in go

reverse.go

package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strconv"
	"strings"
)

var (
	egress, _    = url.Parse(os.Getenv("SCOTTY_ENV_EGRESS_ENDPOINT"))
	port         = os.Getenv("SCOTTY_ENV_LISTENING_PORT")
	Warp10Header = [...]string{"X-Warp10-Elapsed", "X-Warp10-Ops", "X-Warp10-Fetched"}
)

// Transport executes a single HTTP transaction, returning
// a Response for the provided Request.
// If nil, http.DefaultTransport is used.
type transport struct {
	http.RoundTripper
}

func main() {
	fmt.Println("egress:", egress)
	fmt.Println("port:", port)

	proxy := NewWarp10ExecProxy()

	fmt.Println("Starting...\n")
	http.ListenAndServe(":"+port, proxy)
}

// NewWarp10ExecProxy is creating a reverse-proxy that is creating a new body
// NewSingleHostReverseProxy does not rewrite the Host header.
// To rewrite Host headers, use ReverseProxy directly with a custom
// Director policy.
func NewWarp10ExecProxy() *httputil.ReverseProxy {
	var proxy httputil.ReverseProxy

	proxy.Director = director
	proxy.Transport = &transport{}
	return &proxy
}

// director must be a function which modifies
// the request into a new request to be sent
// using Transport. Its response is then copied
// back to the original client unmodified.
func director(req *http.Request) {

	req.URL.Host = egress.Host
	req.Host = egress.Host // Add this line to change Host header
	req.URL.Scheme = egress.Scheme

	// Change Body
	// https://stackoverflow.com/questions/33606330/golang-rewrite-http-request-body
	warpscript := []byte("1 2 + ")
	body, _ := ioutil.ReadAll(req.Body)

	warpscript = append(warpscript, body...)
	req.Body = ioutil.NopCloser(bytes.NewBuffer(warpscript))
	req.ContentLength = int64(len(warpscript))
	req.Header.Set("Content-Length", strconv.Itoa(len(warpscript)))
}

// RoundTrip executes a single HTTP transaction, returning
// a Response for the provided Request.
func (t *transport) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	resp, err = http.DefaultTransport.RoundTrip(req)
	if err != nil {
		return nil, err
	}
	//dump, err := httputil.DumpResponse(resp, true)
	//if err != nil {
	//	log.Fatalln("error dumping request", err)
	//	return
	//}
	//fmt.Println(string(dump))

	// Getting the number of operations performed by a WarpScript
	for _, header := range Warp10Header {
		// Checking if header exists
		if len(resp.Header.Get(header)) != 0 {
			value, err := strconv.ParseInt(resp.Header.Get(header), 10, 64)
			if err != nil {
				return nil, err
			}
			resp.Header.Del(header)
			resp.Header.Add(strings.Replace(header, "Warp10", "Metrics", -1), string(value))

		}
	}

	return resp, nil

}

// Checking environnement
func init() {
	if len(egress.Host) == 0 || len(port) == 0 {
		log.Fatalln("SCOTTY_ENV_EGRESS_ENDPOINT or SCOTTY_ENV_LISTENING_PORT is not set, quitting...")
	}
}