138 lines
3.2 KiB
Go
138 lines
3.2 KiB
Go
package integration
|
|
|
|
import (
|
|
"bufio"
|
|
"net"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/pires/go-proxyproto"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/stretchr/testify/suite"
|
|
"github.com/traefik/traefik/v2/integration/try"
|
|
)
|
|
|
|
type ProxyProtocolSuite struct {
|
|
BaseSuite
|
|
whoamiIP string
|
|
}
|
|
|
|
func TestProxyProtocolSuite(t *testing.T) {
|
|
suite.Run(t, new(ProxyProtocolSuite))
|
|
}
|
|
|
|
func (s *ProxyProtocolSuite) SetupSuite() {
|
|
s.BaseSuite.SetupSuite()
|
|
|
|
s.createComposeProject("proxy-protocol")
|
|
s.composeUp()
|
|
|
|
s.whoamiIP = s.getComposeServiceIP("whoami")
|
|
}
|
|
|
|
func (s *ProxyProtocolSuite) TearDownSuite() {
|
|
s.BaseSuite.TearDownSuite()
|
|
}
|
|
|
|
func (s *ProxyProtocolSuite) TestProxyProtocolTrusted() {
|
|
file := s.adaptFile("fixtures/proxy-protocol/proxy-protocol.toml", struct {
|
|
HaproxyIP string
|
|
WhoamiIP string
|
|
}{WhoamiIP: s.whoamiIP})
|
|
|
|
s.traefikCmd(withConfigFile(file))
|
|
|
|
err := try.GetRequest("http://127.0.0.1:8000/whoami", 10*time.Second)
|
|
require.NoError(s.T(), err)
|
|
|
|
content, err := proxyProtoRequest("127.0.0.1:8000", 1)
|
|
require.NoError(s.T(), err)
|
|
assert.Contains(s.T(), content, "X-Forwarded-For: 1.2.3.4")
|
|
|
|
content, err = proxyProtoRequest("127.0.0.1:8000", 2)
|
|
require.NoError(s.T(), err)
|
|
assert.Contains(s.T(), content, "X-Forwarded-For: 1.2.3.4")
|
|
}
|
|
|
|
func (s *ProxyProtocolSuite) TestProxyProtocolNotTrusted() {
|
|
file := s.adaptFile("fixtures/proxy-protocol/proxy-protocol.toml", struct {
|
|
HaproxyIP string
|
|
WhoamiIP string
|
|
}{WhoamiIP: s.whoamiIP})
|
|
|
|
s.traefikCmd(withConfigFile(file))
|
|
|
|
err := try.GetRequest("http://127.0.0.1:9000/whoami", 10*time.Second)
|
|
require.NoError(s.T(), err)
|
|
|
|
content, err := proxyProtoRequest("127.0.0.1:9000", 1)
|
|
require.NoError(s.T(), err)
|
|
assert.Contains(s.T(), content, "X-Forwarded-For: 127.0.0.1")
|
|
|
|
content, err = proxyProtoRequest("127.0.0.1:9000", 2)
|
|
require.NoError(s.T(), err)
|
|
assert.Contains(s.T(), content, "X-Forwarded-For: 127.0.0.1")
|
|
}
|
|
|
|
func proxyProtoRequest(address string, version byte) (string, error) {
|
|
// Open a TCP connection to the server
|
|
conn, err := net.Dial("tcp", address)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer conn.Close()
|
|
|
|
// Create a Proxy Protocol header with v1
|
|
proxyHeader := &proxyproto.Header{
|
|
Version: version,
|
|
Command: proxyproto.PROXY,
|
|
TransportProtocol: proxyproto.TCPv4,
|
|
DestinationAddr: &net.TCPAddr{
|
|
IP: net.ParseIP("127.0.0.1"),
|
|
Port: 8000,
|
|
},
|
|
SourceAddr: &net.TCPAddr{
|
|
IP: net.ParseIP("1.2.3.4"),
|
|
Port: 62541,
|
|
},
|
|
}
|
|
|
|
// After the connection was created write the proxy headers first
|
|
_, err = proxyHeader.WriteTo(conn)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
// Create an HTTP request
|
|
request := "GET /whoami HTTP/1.1\r\n" +
|
|
"Host: 127.0.0.1\r\n" +
|
|
"Connection: close\r\n" +
|
|
"\r\n"
|
|
|
|
// Write the HTTP request to the TCP connection
|
|
writer := bufio.NewWriter(conn)
|
|
_, err = writer.WriteString(request)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
// Flush the buffer to ensure the request is sent
|
|
err = writer.Flush()
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
// Read the response from the server
|
|
var content string
|
|
scanner := bufio.NewScanner(conn)
|
|
for scanner.Scan() {
|
|
content += scanner.Text() + "\n"
|
|
}
|
|
|
|
if scanner.Err() != nil {
|
|
return "", err
|
|
}
|
|
|
|
return content, nil
|
|
}
|