e2e: cover actor UDP and raw TCP egress

demos/egress/main.go

@@ -13,8 +13,10 @@
 // limitations under the License.
 
 // Command egress is a small HTTP service for demonstrating per-Actor egress
-// policy. It accepts a URL, fetches it, and returns the upstream response, and
-// on a second endpoint it makes gRPC calls and returns what came back.
+// policy. It accepts a URL, fetches it, and returns the upstream response. On
+// /grpc it makes gRPC calls and returns what came back, and on /tcp it opens a
+// raw TCP connection, so that egress can be exercised with something other than
+// HTTP.
 package main
 
 import (
@@ -34,6 +36,7 @@ import (
 	"os"
 	"strconv"
 	"sync"
+	"syscall"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -210,6 +213,7 @@ func newHandler(client *http.Client) http.Handler {
 	mux.HandleFunc("/mux", handleMux)
 	mux.HandleFunc("/websocket", handleWebSocket)
 	mux.HandleFunc("/grpc", handleGRPC)
+	mux.HandleFunc("/tcp", handleTCPProbe)
 	return mux
 }
 
@@ -862,6 +866,146 @@ func writeGRPCFailure(w http.ResponseWriter, rpc string, err error) {
 	})
 }
 
+// tcpProbeRequest asks for one raw TCP exchange.
+type tcpProbeRequest struct {
+	Address   string `json:"address"`
+	Send      string `json:"send,omitempty"`
+	ReadBytes int    `json:"readBytes,omitempty"`
+	Timeout   string `json:"timeout,omitempty"`
+	// SkipBanner writes Send as soon as the connection is up instead of first
+	// waiting for the peer to speak, so a path that classifies traffic by the
+	// client's first bytes sees them right away.
+	SkipBanner bool `json:"skipBanner,omitempty"`
+}
+
+type tcpProbeResponse struct {
+	// Banner is whatever the peer sent before being spoken to.
+	Banner   string `json:"banner,omitempty"`
+	Received string `json:"received,omitempty"`
+	// Closed reports that the peer closed or reset the connection, as opposed
+	// to a read that only ran out its timeout.
+	Closed bool   `json:"closed,omitempty"`
+	Error  string `json:"error,omitempty"`
+}
+
+// handleTCPProbe opens a TCP connection and, unless SkipBanner is set, reads
+// before it writes.
+func handleTCPProbe(w http.ResponseWriter, r *http.Request) {
+	const defaultProbeTimeout = 5 * time.Second
+	// Enough for an SSH identification string or a test banner.
+	const defaultProbeReadBytes = 512
+	const maxProbeReadBytes = 8 << 10 // 8 KiB
+
+	if r.Method != http.MethodPost {
+		w.Header().Set("Allow", http.MethodPost)
+		writeJSON(w, http.StatusMethodNotAllowed, tcpProbeResponse{Error: "method must be POST"})
+		return
+	}
+
+	var input tcpProbeRequest
+	decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, maxRequestBody))
+	if err := decoder.Decode(&input); err != nil {
+		writeJSON(w, http.StatusBadRequest, tcpProbeResponse{Error: fmt.Sprintf("invalid JSON payload: %v", err)})
+		return
+	}
+	if _, _, err := net.SplitHostPort(input.Address); err != nil {
+		writeJSON(w, http.StatusBadRequest, tcpProbeResponse{Error: fmt.Sprintf("address must be host:port: %v", err)})
+		return
+	}
+
+	timeout := defaultProbeTimeout
+	if input.Timeout != "" {
+		parsed, err := time.ParseDuration(input.Timeout)
+		if err != nil {
+			writeJSON(w, http.StatusBadRequest, tcpProbeResponse{Error: fmt.Sprintf("invalid timeout: %v", err)})
+			return
+		}
+		timeout = parsed
+	}
+	readBytes := input.ReadBytes
+	if readBytes <= 0 {
+		readBytes = defaultProbeReadBytes
+	}
+	readBytes = min(readBytes, maxProbeReadBytes)
+
+	dialer := net.Dialer{Timeout: timeout}
+	connection, err := dialer.DialContext(r.Context(), "tcp", input.Address)
+	if err != nil {
+		writeJSON(w, http.StatusBadGateway, tcpProbeResponse{Error: fmt.Sprintf("dialing %s: %v", input.Address, err)})
+		return
+	}
+	defer connection.Close()
+
+	// Read before writing anything at all, so an empty banner really does mean
+	// the peer stayed silent.
+	var response tcpProbeResponse
+	if !input.SkipBanner {
+		banner, closed, err := readWithTimeout(connection, readBytes, timeout)
+		if err != nil {
+			writeJSON(w, http.StatusBadGateway, tcpProbeResponse{Error: fmt.Sprintf("reading banner from %s: %v", input.Address, err)})
+			return
+		}
+		response.Banner = string(banner)
+		if closed {
+			response.Closed = true
+			writeJSON(w, http.StatusOK, response)
+			return
+		}
+	}
+	if input.Send == "" {
+		writeJSON(w, http.StatusOK, response)
+		return
+	}
+
+	if err := connection.SetWriteDeadline(time.Now().Add(timeout)); err != nil {
+		writeJSON(w, http.StatusBadGateway, tcpProbeResponse{Error: fmt.Sprintf("setting write deadline: %v", err)})
+		return
+	}
+	if _, err := io.WriteString(connection, input.Send); err != nil {
+		if peerClosed(err) {
+			response.Closed = true
+			writeJSON(w, http.StatusOK, response)
+			return
+		}
+		writeJSON(w, http.StatusBadGateway, tcpProbeResponse{Error: fmt.Sprintf("writing to %s: %v", input.Address, err)})
+		return
+	}
+	received, closed, err := readWithTimeout(connection, readBytes, timeout)
+	if err != nil {
+		writeJSON(w, http.StatusBadGateway, tcpProbeResponse{Error: fmt.Sprintf("reading reply from %s: %v", input.Address, err)})
+		return
+	}
+	response.Received = string(received)
+	response.Closed = closed
+	writeJSON(w, http.StatusOK, response)
+}
+
+// readWithTimeout returns the bytes of a single read, capped at limit, and
+// whether the peer closed or reset the connection. A peer that says nothing
+// within timeout yields no bytes, closed false and no error, since silence is a
+// legitimate answer to "does this peer speak first?".
+func readWithTimeout(connection net.Conn, limit int, timeout time.Duration) ([]byte, bool, error) {
+	if err := connection.SetReadDeadline(time.Now().Add(timeout)); err != nil {
+		return nil, false, fmt.Errorf("setting read deadline: %w", err)
+	}
+	buffer := make([]byte, limit)
+	n, err := connection.Read(buffer)
+	switch {
+	case err == nil, errors.Is(err, os.ErrDeadlineExceeded):
+		return buffer[:n], false, nil
+	case peerClosed(err):
+		return buffer[:n], true, nil
+	default:
+		return nil, false, err
+	}
+}
+
+// peerClosed reports whether err means the peer closed or reset the
+// connection.
+func peerClosed(err error) bool {
+	return errors.Is(err, io.EOF) || errors.Is(err, syscall.ECONNRESET) || errors.Is(err, syscall.EPIPE)
+}
+
 func validateURL(raw string) error {
 	parsed, err := url.Parse(raw)
 	if err != nil {

demos/egress/main_test.go

@@ -850,3 +850,203 @@ func TestSlowHeadersInvalidRequests(t *testing.T) {
 		}
 	})
 }
+
+// TestTCPProbeServerSpeaksFirst covers the ordering the probe exists for: the
+// peer's greeting is reported without the probe having written anything, and
+// the reply to Send comes back separately.
+func TestTCPProbeServerSpeaksFirst(t *testing.T) {
+	const banner = "TESTBANNER/1.0\r\n"
+	address := startTestPeer(t, func(connection net.Conn) {
+		if _, err := io.WriteString(connection, banner); err != nil {
+			return
+		}
+		buffer := make([]byte, 64)
+		n, err := connection.Read(buffer)
+		if err != nil {
+			return
+		}
+		_, _ = connection.Write(buffer[:n])
+	})
+
+	got := probe(t, tcpProbeRequest{Address: address, Send: "ping", Timeout: "5s"}, http.StatusOK)
+	if got.Banner != banner {
+		t.Errorf("banner = %q, want %q", got.Banner, banner)
+	}
+	if got.Received != "ping" {
+		t.Errorf("received = %q, want %q", got.Received, "ping")
+	}
+}
+
+// TestTCPProbeSilentPeer verifies that a client-speaks-first peer yields an empty banner and a 200.
+func TestTCPProbeSilentPeer(t *testing.T) {
+	address := startTestPeer(t, func(connection net.Conn) {
+		buffer := make([]byte, 64)
+		n, err := connection.Read(buffer)
+		if err != nil {
+			return
+		}
+		_, _ = connection.Write(buffer[:n])
+	})
+
+	got := probe(t, tcpProbeRequest{Address: address, Send: "ping", Timeout: "250ms"}, http.StatusOK)
+	if got.Banner != "" {
+		t.Errorf("banner = %q, want empty for a peer that does not speak first", got.Banner)
+	}
+	if got.Received != "ping" {
+		t.Errorf("received = %q, want %q", got.Received, "ping")
+	}
+}
+
+// TestTCPProbeSkipBannerWritesFirst verifies that SkipBanner sends before
+// reading: the peer only answers once it has read, and the probe's timeout is
+// far shorter than any banner wait would need, so the bytes must already be
+// on the wire when the peer reads.
+func TestTCPProbeSkipBannerWritesFirst(t *testing.T) {
+	firstRead := make(chan time.Duration, 1)
+	address := startTestPeer(t, func(connection net.Conn) {
+		accepted := time.Now()
+		buffer := make([]byte, 64)
+		n, err := connection.Read(buffer)
+		if err != nil {
+			return
+		}
+		firstRead <- time.Since(accepted)
+		_, _ = connection.Write(buffer[:n])
+	})
+
+	got := probe(t, tcpProbeRequest{Address: address, Send: "ping", SkipBanner: true, Timeout: "2s"}, http.StatusOK)
+	if got.Banner != "" {
+		t.Errorf("banner = %q, want empty when the banner read is skipped", got.Banner)
+	}
+	if got.Received != "ping" {
+		t.Errorf("received = %q, want %q", got.Received, "ping")
+	}
+	// Reading first would have held the write back for the whole 2s timeout.
+	if elapsed := <-firstRead; elapsed >= time.Second {
+		t.Errorf("peer waited %v for the first bytes, want them sent without a banner read", elapsed)
+	}
+}
+
+func TestTCPProbeReadBytesCapsTheBanner(t *testing.T) {
+	address := startTestPeer(t, func(connection net.Conn) {
+		_, _ = io.WriteString(connection, "0123456789")
+	})
+
+	got := probe(t, tcpProbeRequest{Address: address, ReadBytes: 4, Timeout: "5s"}, http.StatusOK)
+	if got.Banner != "0123" {
+		t.Errorf("banner = %q, want %q", got.Banner, "0123")
+	}
+}
+
+func TestTCPProbeDialFailure(t *testing.T) {
+	// Port 0 is not connectable, so this fails without depending on which
+	// ports happen to be free.
+	got := probe(t, tcpProbeRequest{Address: "127.0.0.1:0", Timeout: "2s"}, http.StatusBadGateway)
+	if got.Error == "" {
+		t.Error("error = empty, want a dial failure")
+	}
+}
+
+func TestTCPProbeInvalidRequests(t *testing.T) {
+	tests := []struct {
+		name   string
+		method string
+		body   string
+		status int
+	}{
+		{name: "method", method: http.MethodGet, body: `{}`, status: http.StatusMethodNotAllowed},
+		{name: "malformed JSON", method: http.MethodPost, body: `{`, status: http.StatusBadRequest},
+		{name: "address without port", method: http.MethodPost, body: `{"address":"example.com"}`, status: http.StatusBadRequest},
+		{name: "invalid timeout", method: http.MethodPost, body: `{"address":"127.0.0.1:9","timeout":"soon"}`, status: http.StatusBadRequest},
+	}
+
+	handler := newHandler(http.DefaultClient)
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			recorder := httptest.NewRecorder()
+			request := httptest.NewRequest(test.method, "/tcp", strings.NewReader(test.body))
+			handler.ServeHTTP(recorder, request)
+			if recorder.Code != test.status {
+				t.Errorf("status = %d, want %d; body = %s", recorder.Code, test.status, recorder.Body.String())
+			}
+		})
+	}
+}
+
+// TestTCPProbeReportsClosed separates a peer that closes the connection from
+// one that only stays quiet, which an empty read alone cannot.
+func TestTCPProbeReportsClosed(t *testing.T) {
+	tests := []struct {
+		name       string
+		serve      func(net.Conn)
+		skipBanner bool
+		wantClosed bool
+	}{
+		{name: "closes before speaking", serve: func(c net.Conn) { _ = c.Close() }, wantClosed: true},
+		{name: "closes after reading", skipBanner: true, wantClosed: true, serve: func(c net.Conn) {
+			_, _ = c.Read(make([]byte, 64))
+			_ = c.Close()
+		}},
+		{name: "stays silent", skipBanner: true, wantClosed: false, serve: func(c net.Conn) {
+			_, _ = c.Read(make([]byte, 64))
+			time.Sleep(time.Second)
+		}},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			address := startTestPeer(t, tt.serve)
+			got := probe(t, tcpProbeRequest{Address: address, Send: "ping\r\n", SkipBanner: tt.skipBanner, Timeout: "300ms"}, http.StatusOK)
+			if got.Closed != tt.wantClosed {
+				t.Errorf("closed = %v, want %v (response %+v)", got.Closed, tt.wantClosed, got)
+			}
+			if got.Banner != "" || got.Received != "" {
+				t.Errorf("response = %+v, want no bytes", got)
+			}
+		})
+	}
+}
+
+func probe(t *testing.T, input tcpProbeRequest, wantStatus int) tcpProbeResponse {
+	t.Helper()
+	payload, err := json.Marshal(input)
+	if err != nil {
+		t.Fatal(err)
+	}
+	recorder := httptest.NewRecorder()
+	request := httptest.NewRequest(http.MethodPost, "/tcp", strings.NewReader(string(payload)))
+	newHandler(http.DefaultClient).ServeHTTP(recorder, request)
+
+	if recorder.Code != wantStatus {
+		t.Fatalf("status = %d, want %d; body = %s", recorder.Code, wantStatus, recorder.Body.String())
+	}
+	var got tcpProbeResponse
+	if err := json.NewDecoder(recorder.Body).Decode(&got); err != nil {
+		t.Fatalf("decoding response: %v", err)
+	}
+	return got
+}
+
+// startTestPeer listens on loopback and hands each connection to serve. It
+// returns the address to probe.
+func startTestPeer(t *testing.T, serve func(net.Conn)) string {
+	t.Helper()
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listening: %v", err)
+	}
+	t.Cleanup(func() { listener.Close() })
+
+	go func() {
+		for {
+			connection, err := listener.Accept()
+			if err != nil {
+				return
+			}
+			go func() {
+				defer connection.Close()
+				serve(connection)
+			}()
+		}
+	}()
+	return listener.Addr().String()
+}