go 实现端口代理

首页 > GoLand学习作者：皮皮华 2022年6月1日 10:33 星期三热度：976° 字号：小中大评论：0 条

时间：2022-6-1 10:33 热度：976° 评论：0 条

package main

import (
	"bytes"
	"compress/gzip"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"runtime"
)

/**
启动
-port=8080 -target=http://127.0.0.1:9411
*/

var (
	targetURL  *url.URL
	targetAddr *string //要代理的服务的http地址

)

/*
编译跨平台的只需要修改GOOS、GOARCH、CGO_ENABLED三个环境变量即可
GOOS：目标平台的操作系统(darwin、freebsd、linux、windows)
GOARCH：目标平台的体系架构32位还是64位(386、amd64、arm)
SET CGO_ENABLED=0
SET GOOS=linux
SET GOARCH=amd64
go build
*/
//参数识别
func Daemonize(args ...string) {
	sysType := runtime.GOOS
	if sysType == "linux" {
		var arg []string
		if len(args) > 1 {
			arg = args[1:]
		}
		fmt.Println("run in backMode")
		cmd := exec.Command(args[0], arg...)
		cmd.Env = os.Environ()
		cmd.Start()
	} else if sysType == "windows" {

	}
}
func printMsg(r *http.Request) error {

	//gzip压缩格式
	if r.Header.Get("Content-Encoding") == "gzip" {
		fmt.Println("Content-Encoding: " + r.Header.Get("Content-Encoding"))
		reader, e := gzip.NewReader(r.Body)
		if e != nil {
			fmt.Printf("create gzip reader err: %s \n", e.Error())
			return e
		} else {
			//ioutil.ReadAll() 读取之后会关闭流
			//这里读出之后再新建一个流然后重新赋值
			bodyBytes, e := ioutil.ReadAll(reader)
			if e != nil {
				fmt.Printf("read request body err: %s \n", e.Error())
				return e
			} else {
				if len(bodyBytes) > 0 {
					fmt.Println(string(bodyBytes))
				}
			}

			//需要将数据再以gzip格式压缩，否则和header中的content-length对不上
			var zBuf bytes.Buffer
			zw := gzip.NewWriter(&zBuf)
			if _, err := zw.Write(bodyBytes); err != nil {
				return err
			}
			zw.Close()

			r.Body = ioutil.NopCloser(&zBuf)
		}
	} else {
		//其它压缩格式
		bodyBytes, e := ioutil.ReadAll(r.Body)
		r.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes))
		if e != nil {
			fmt.Printf("read request body err: %s \n", e.Error())
			return e
		} else {
			if len(bodyBytes) > 0 {
				fmt.Println(string(bodyBytes))
			}
		}
	}
	return nil
}

func proxy(w http.ResponseWriter, r *http.Request) {
	if e := printMsg(r); e != nil {
		_, _ = fmt.Fprint(w, e.Error())
		return
	}

	o := new(http.Request)
	*o = *r

	o.Host = targetURL.Host
	o.URL.Scheme = targetURL.Scheme
	o.URL.Host = targetURL.Host
	o.URL.Path = r.URL.Path
	o.URL.RawQuery = r.URL.RawQuery
	o.Proto = r.Proto
	o.ProtoMajor = 1
	o.ProtoMinor = 1
	o.Close = false

	res, err := http.DefaultTransport.RoundTrip(o)
	if err != nil {
		fmt.Printf("http: proxy error: %v\n", err)
		w.WriteHeader(http.StatusInternalServerError)
		_, _ = w.Write([]byte(err.Error()))
		return
	}

	hdr := w.Header()
	for k, vv := range res.Header {
		for _, v := range vv {
			hdr.Add(k, v)
		}
	}
	for _, c := range res.Cookies() {
		w.Header().Add("Set-Cookie", c.Raw)
	}

	w.WriteHeader(res.StatusCode)
	if res.Body != nil {
		_, _ = io.Copy(w, res.Body)
	}
}

func main() {

	args := os.Args
	daemon := false
	for k, v := range args {
		if v == "-d" {
			daemon = true
			args[k] = ""
		} else if v == "restart" {
			daemon = true
			args[k] = ""
		}
	}
	if daemon {
		Daemonize(args...)
		return
	}
	//代理自身的端口号
	port := flag.Int("port", 8098, "Http proxy liston port")
	//后端服务的地址
	targetAddr = flag.String("target", "https://api-testing.phonepe.com", "The server address. Format: http://ip:port")
	flag.Parse()

	addr := fmt.Sprintf(":%d", *port)
	fmt.Println("proxy addr: " + addr)
	fmt.Println("target addr: " + *targetAddr)

	if *port <= 0 {
		panic("proxy port is invalid!")
	}

	if len(*targetAddr) == 0 {
		panic("target addr is empty! Format: http://ip:port")
	}

	var e error
	targetURL, e = url.Parse(*targetAddr)
	if e != nil {
		panic(e.Error())
	}

	http.HandleFunc("/", proxy)
	err := http.ListenAndServe(addr, nil)
	if err != nil {
		panic(err.Error())
	}
}