proxy_protocol_test.go

package risefront

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"testing"
	"time"
)

func TestProxyProtocol(t *testing.T) {
	// Track the client IP seen by the child process
	clientIPChan := make(chan string, 1)
	portChan := make(chan string, 1)

	cfg := Config{
		Addresses: []string{"127.0.0.1:0"},
		Run: func(listeners []net.Listener) error {
			// Send the actual port to the test
			select {
			case portChan <- listeners[0].Addr().String():
			default:
			}

			mux := http.NewServeMux()
			mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
				// Get the remote address which should be the real client IP
				// thanks to PROXY protocol
				clientIP := r.RemoteAddr
				select {
				case clientIPChan <- clientIP:
				default:
				}
				fmt.Fprintf(w, "OK")
			})

			server := &http.Server{
				Handler: mux,
			}

			go func() {
				if err := server.Serve(listeners[0]); err != nil && err != http.ErrServerClosed {
					t.Logf("Server error: %v", err)
				}
			}()

			// Wait for test to complete
			time.Sleep(2 * time.Second)
			return server.Shutdown(context.Background())
		},
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// Start the parent process
	errChan := make(chan error, 1)
	go func() {
		if err := New(ctx, cfg); err != nil && err != context.DeadlineExceeded && err != context.Canceled {
			errChan <- err
		}
	}()

	// Get the actual listening address
	var addr string
	select {
	case addr = <-portChan:
		t.Logf("Server listening on: %s", addr)
	case err := <-errChan:
		t.Fatalf("Parent process error: %v", err)
	case <-time.After(2 * time.Second):
		t.Fatal("Timeout waiting for server to start")
	}

	// Make a request to the parent
	resp, err := http.Get("http://" + addr + "/")
	if err != nil {
		t.Fatalf("Failed to connect to server: %v", err)
	}
	defer resp.Body.Close()

	// Wait for client IP to be received
	select {
	case receivedClientIP := <-clientIPChan:
		t.Logf("Child process received client IP: %s", receivedClientIP)

		// The client IP should not be from the unix socket
		if receivedClientIP == "" {
			t.Error("Client IP was empty")
		}

		// Check that it's not a unix socket address
		if receivedClientIP == "@" || receivedClientIP == "unix" {
			t.Errorf("Client IP appears to be from unix socket, not preserved: %s", receivedClientIP)
		}

	case <-time.After(1 * time.Second):
		t.Error("Timeout waiting for client IP")
	}
}

func TestProxyProtocolPreservesClientIP(t *testing.T) {
	// This test verifies that client IP is preserved through the PROXY protocol
	// when connections are forwarded from parent to child process

	receivedIPs := make(chan string, 10)

	cfg := Config{
		Addresses: []string{"127.0.0.1:9999"},
		Run: func(listeners []net.Listener) error {
			// Create a simple TCP server that reads the connection's remote address
			for _, ln := range listeners {
				go func(listener net.Listener) {
					for {
						conn, err := listener.Accept()
						if err != nil {
							return
						}

						// The RemoteAddr should be the original client IP
						// preserved through PROXY protocol
						remoteAddr := conn.RemoteAddr().String()
						select {
						case receivedIPs <- remoteAddr:
						default:
						}

						conn.Write([]byte("OK\n"))
						conn.Close()
					}
				}(ln)
			}

			time.Sleep(1 * time.Second)
			return nil
		},
	}

	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()

	// Start the risefront parent/child system
	go func() {
		if err := New(ctx, cfg); err != nil && err != context.Canceled {
			t.Logf("New() error (might be expected): %v", err)
		}
	}()

	// Give system time to start
	time.Sleep(500 * time.Millisecond)

	// Connect as a client
	conn, err := net.Dial("tcp", "127.0.0.1:9999")
	if err != nil {
		t.Logf("Note: Connection failed, might be expected in test environment: %v", err)
		return
	}
	defer conn.Close()

	originalClientAddr := conn.LocalAddr().String()
	t.Logf("Original client address: %s", originalClientAddr)

	// Read response
	buf := make([]byte, 10)
	conn.Read(buf)

	// Check what IP the child process saw
	select {
	case receivedIP := <-receivedIPs:
		t.Logf("Child process saw client IP as: %s", receivedIP)

		// Verify it's not a unix socket address
		if receivedIP == "" || receivedIP == "@" {
			t.Errorf("Client IP was not preserved, got: %s", receivedIP)
		}

		// The IP should contain a real TCP address
		if _, _, err := net.SplitHostPort(receivedIP); err != nil {
			t.Errorf("Received IP is not a valid TCP address: %s", receivedIP)
		}

	case <-time.After(1 * time.Second):
		t.Error("No client IP received by child process")
	}
}