e2e: cover egress for raw TCP and SSH traffic, in both speaking orders

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 (
@@ -150,6 +152,7 @@ func newHandler(client *http.Client) http.Handler {
 		writeJSON(w, response.StatusCode, fetchResponse{StatusCode: response.StatusCode, Body: string(body)})
 	})
 	mux.HandleFunc("/grpc", handleGRPC)
+	mux.HandleFunc("/tcp", handleTCPProbe)
 	return mux
 }
 
@@ -290,6 +293,114 @@ 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"`
+}
+
+type tcpProbeResponse struct {
+	// Banner is whatever the peer sent before being spoken to.
+	Banner   string `json:"banner,omitempty"`
+	Received string `json:"received,omitempty"`
+	Error    string `json:"error,omitempty"`
+}
+
+// handleTCPProbe opens a TCP connection and 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.
+	banner, 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 := tcpProbeResponse{Banner: string(banner)}
+	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 {
+		writeJSON(w, http.StatusBadGateway, tcpProbeResponse{Error: fmt.Sprintf("writing to %s: %v", input.Address, err)})
+		return
+	}
+	received, 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)
+	writeJSON(w, http.StatusOK, response)
+}
+
+// readWithTimeout returns the bytes of a single read, capped at limit. A peer
+// that says nothing within timeout yields no bytes 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, error) {
+	if err := connection.SetReadDeadline(time.Now().Add(timeout)); err != nil {
+		return nil, fmt.Errorf("setting read deadline: %w", err)
+	}
+	buffer := make([]byte, limit)
+	n, err := connection.Read(buffer)
+	if err != nil && !errors.Is(err, os.ErrDeadlineExceeded) && !errors.Is(err, io.EOF) {
+		return nil, err
+	}
+	return buffer[:n], nil
+}
+
 func validateURL(raw string) error {
 	parsed, err := url.Parse(raw)
 	if err != nil {

demos/egress/main_test.go

@@ -110,6 +110,143 @@ func TestOutboundFailure(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")
+	}
+}
+
+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 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())
+			}
+		})
+	}
+}
+
+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")
+	}
+}
+
+// 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()
+}
+
+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
+}
+
 type roundTripFunc func(*http.Request) (*http.Response, error)
 
 func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) {

internal/e2e/fixtures/testserver/main.go

@@ -19,6 +19,7 @@
 //
 //	testserver grpc  --listen=:50051   a cleartext HTTP/2 gRPC echo origin
 //	testserver http  --listen=:8080    a plain HTTP origin serving /healthz
+//	testserver tcpecho --listen=:2222  a raw TCP echo origin, greeting optional
 //	testserver egressprobe --listen=:8080  a client that drives the egress gateway
 //	testserver websocket --listen=:8080  a websocket server that responds to PINGs
 //
@@ -43,7 +44,7 @@ func main() {
 		Use:   "testserver",
 		Short: "Multi-mode helper server for the egress e2e suites.",
 	}
-	root.AddCommand(newGRPCCmd(), newHTTPCmd(), newEgressProbeCmd(), newWebsocketCmd())
+	root.AddCommand(newGRPCCmd(), newHTTPCmd(), newTCPEchoCmd(), newEgressProbeCmd(), newWebsocketCmd())
 	if err := root.Execute(); err != nil {
 		fmt.Fprintln(os.Stderr, err)
 		os.Exit(1)