From a1cfdbc1d61c02997b4a39c8e675d3d961d5dae7 Mon Sep 17 00:00:00 2001
From: dingjun <fangdingjun@gmail.com>
Date: Mon, 15 Apr 2019 11:56:21 +0800
Subject: [PATCH] add support to send proxy protocol to dst

---
 config.sample.yaml | 10 ++++++++--
 main.go            | 44 +++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/config.sample.yaml b/config.sample.yaml
index f0aa56a..739b68d 100644
--- a/config.sample.yaml
+++ b/config.sample.yaml
@@ -18,10 +18,16 @@ listen:
 # exact match or wildcard match
 forward_rules:
     # www.example.com forward to 127.0.0.1:8443
-    www.example.com: 127.0.0.1:8443
+    # send proxy protocol v2 header
+    www.example.com: 127.0.0.1:8443 proxy-v2
 
     # b.example.com forward to 127.0.0.1:8541
-    b.example.com: 127.0.0.1:8541
+    # send proxy protocol v1 header
+    b.example.com: 127.0.0.1:8541 proxy-v1
+    
+    # b.example.com forward to 127.0.0.1:8541
+    # no proxy protocol header send
+    b.example.com: 127.0.0.1:8542
 
     # forward by SNI, wildcard 
     # ex: SNI a.example.com from port 9999 forward to a.example.com:443
diff --git a/main.go b/main.go
index 7acef1b..5f188db 100644
--- a/main.go
+++ b/main.go
@@ -2,10 +2,14 @@ package main
 
 import (
 	"fmt"
-	"github.com/go-yaml/yaml"
-	"github.com/golang/glog"
 	"io"
 	"io/ioutil"
+	"strings"
+
+	"github.com/go-yaml/yaml"
+	"github.com/golang/glog"
+	proxyproto "github.com/pires/go-proxyproto"
+
 	//"crypto/tls"
 	"flag"
 	//"log"
@@ -56,7 +60,37 @@ func getSNIServerName(buf []byte) string {
 }
 
 func forward(c net.Conn, data []byte, dst string) {
-	c1, err := net.Dial("tcp", dst)
+	addr := dst
+	proxyProto := 0
+
+	ss := strings.Fields(dst)
+
+	var hdr proxyproto.Header
+
+	if len(ss) > 1 {
+		addr = ss[0]
+		raddr := c.RemoteAddr().(*net.TCPAddr)
+		hdr = proxyproto.Header{
+			Version:            1,
+			Command:            proxyproto.PROXY,
+			TransportProtocol:  proxyproto.TCPv4,
+			SourceAddress:      raddr.IP.To4(),
+			DestinationAddress: net.IP{0, 0, 0, 0},
+			SourcePort:         uint16(raddr.Port),
+			DestinationPort:    0,
+		}
+
+		switch strings.ToLower(ss[1]) {
+		case "proxy-v1":
+			proxyProto = 1
+			hdr.Version = 1
+		case "proxy-v2":
+			proxyProto = 2
+			hdr.Version = 2
+		}
+	}
+
+	c1, err := net.Dial("tcp", addr)
 	if err != nil {
 		glog.Error(err)
 		return
@@ -64,6 +98,10 @@ func forward(c net.Conn, data []byte, dst string) {
 
 	defer c1.Close()
 
+	if proxyProto != 0 {
+		hdr.WriteTo(c1)
+	}
+
 	if _, err = c1.Write(data); err != nil {
 		glog.Error(err)
 		return