From 1b9f55a002266bdcd36dbc4e6ee43b1920f85507 Mon Sep 17 00:00:00 2001 From: Devin Carr Date: Fri, 30 Jun 2023 11:27:37 -0700 Subject: [PATCH] TUN-7550: Add pprof endpoint to management service --- component-tests/test_management.py | 93 ++++++++++++++++++++++++++++++ component-tests/test_tunnel.py | 32 ---------- connection/quic.go | 8 +++ management/service.go | 5 ++ 4 files changed, 106 insertions(+), 32 deletions(-) create mode 100644 component-tests/test_management.py diff --git a/component-tests/test_management.py b/component-tests/test_management.py new file mode 100644 index 00000000..3a4708d9 --- /dev/null +++ b/component-tests/test_management.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python +import requests +from conftest import CfdModes +from constants import METRICS_PORT, MAX_RETRIES, BACKOFF_SECS +from retrying import retry +from cli import CloudflaredCli +from util import LOGGER, write_config, start_cloudflared, wait_tunnel_ready, send_requests +import platform + +""" +Each test in TestManagement will: +1. Acquire a management token from Cloudflare public API +2. Make a request against the management service for the running tunnel +""" +class TestManagement: + """ + test_get_host_details does the following: + 1. It gets a management token from Tunnelstore using cloudflared tail token + 2. It gets the connector_id after starting a cloudflare tunnel + 3. It sends a request to the management host with the connector_id and management token + 4. Asserts that the response has a hostname and ip. + """ + def test_get_host_details(self, tmp_path, component_tests_config): + # TUN-7377 : wait_tunnel_ready does not work properly in windows. + # Skipping this test for windows for now and will address it as part of tun-7377 + if platform.system() == "Windows": + return + config = component_tests_config(cfd_mode=CfdModes.NAMED, run_proxy_dns=False, provide_ingress=False) + LOGGER.debug(config) + headers = {} + headers["Content-Type"] = "application/json" + config_path = write_config(tmp_path, config.full_config) + with start_cloudflared(tmp_path, config, cfd_pre_args=["tunnel", "--ha-connections", "1", "--label" , "test"], cfd_args=["run", "--hello-world"], new_process=True): + wait_tunnel_ready(tunnel_url=config.get_url(), + require_min_connections=1) + cfd_cli = CloudflaredCli(config, config_path, LOGGER) + connector_id = cfd_cli.get_connector_id(config)[0] + url = cfd_cli.get_management_url("host_details", config, config_path) + resp = send_request(url, headers=headers) + + # Assert response json. + assert resp.status_code == 200, "Expected cloudflared to return 200 for host details" + assert resp.json()["hostname"] == "custom:test", "Expected cloudflared to return hostname" + assert resp.json()["ip"] != "", "Expected cloudflared to return ip" + assert resp.json()["connector_id"] == connector_id, "Expected cloudflared to return connector_id" + + """ + test_get_metrics will verify that the /metrics endpoint returns the prometheus metrics dump + """ + def test_get_metrics(self, tmp_path, component_tests_config): + # TUN-7377 : wait_tunnel_ready does not work properly in windows. + # Skipping this test for windows for now and will address it as part of tun-7377 + if platform.system() == "Windows": + return + config = component_tests_config(cfd_mode=CfdModes.NAMED, run_proxy_dns=False, provide_ingress=False) + LOGGER.debug(config) + config_path = write_config(tmp_path, config.full_config) + with start_cloudflared(tmp_path, config, cfd_pre_args=["tunnel", "--ha-connections", "1"], new_process=True): + wait_tunnel_ready(require_min_connections=1) + cfd_cli = CloudflaredCli(config, config_path, LOGGER) + url = cfd_cli.get_management_url("metrics", config, config_path) + resp = send_request(url) + + # Assert response. + assert resp.status_code == 200, "Expected cloudflared to return 200 for /metrics" + assert "# HELP build_info Build and version information" in resp.text, "Expected /metrics to have with the build_info details" + + """ + test_get_pprof_heap will verify that the /debug/pprof/heap endpoint returns a pprof/heap dump response + """ + def test_get_pprof_heap(self, tmp_path, component_tests_config): + # TUN-7377 : wait_tunnel_ready does not work properly in windows. + # Skipping this test for windows for now and will address it as part of tun-7377 + if platform.system() == "Windows": + return + config = component_tests_config(cfd_mode=CfdModes.NAMED, run_proxy_dns=False, provide_ingress=False) + LOGGER.debug(config) + config_path = write_config(tmp_path, config.full_config) + with start_cloudflared(tmp_path, config, cfd_pre_args=["tunnel", "--ha-connections", "1"], new_process=True): + wait_tunnel_ready(require_min_connections=1) + cfd_cli = CloudflaredCli(config, config_path, LOGGER) + url = cfd_cli.get_management_url("debug/pprof/heap", config, config_path) + resp = send_request(url) + + # Assert response. + assert resp.status_code == 200, "Expected cloudflared to return 200 for /debug/pprof/heap" + assert resp.headers["Content-Type"] == "application/octet-stream", "Expected /debug/pprof/heap to have return a binary response" + + +@retry(stop_max_attempt_number=MAX_RETRIES, wait_fixed=BACKOFF_SECS * 1000) +def send_request(url, headers={}): + with requests.Session() as s: + return s.get(url, timeout=BACKOFF_SECS, headers=headers) diff --git a/component-tests/test_tunnel.py b/component-tests/test_tunnel.py index 8266157c..93a39c51 100644 --- a/component-tests/test_tunnel.py +++ b/component-tests/test_tunnel.py @@ -16,38 +16,6 @@ class TestTunnel: with start_cloudflared(tmp_path, config, cfd_pre_args=["tunnel", "--ha-connections", "1"], cfd_args=["run", "--hello-world"], new_process=True): wait_tunnel_ready(tunnel_url=config.get_url(), require_min_connections=1) - - """ - test_get_host_details does the following: - 1. It gets a management token from Tunnelstore using cloudflared tail token - 2. It gets the connector_id after starting a cloudflare tunnel - 3. It sends a request to the management host with the connector_id and management token - 4. Asserts that the response has a hostname and ip. - """ - def test_get_host_details(self, tmp_path, component_tests_config): - # TUN-7377 : wait_tunnel_ready does not work properly in windows. - # Skipping this test for windows for now and will address it as part of tun-7377 - if platform.system() == "Windows": - return - config = component_tests_config(cfd_mode=CfdModes.NAMED, run_proxy_dns=False, provide_ingress=False) - LOGGER.debug(config) - headers = {} - headers["Content-Type"] = "application/json" - config_path = write_config(tmp_path, config.full_config) - with start_cloudflared(tmp_path, config, cfd_pre_args=["tunnel", "--ha-connections", "1", "--label" , "test"], cfd_args=["run", "--hello-world"], new_process=True): - wait_tunnel_ready(tunnel_url=config.get_url(), - require_min_connections=1) - cfd_cli = CloudflaredCli(config, config_path, LOGGER) - connector_id = cfd_cli.get_connector_id(config)[0] - url = cfd_cli.get_management_url("host_details", config, config_path) - resp = send_request(url, headers=headers) - - # Assert response json. - assert resp.status_code == 200, "Expected cloudflared to return 200 for host details" - assert resp.json()["hostname"] == "custom:test", "Expected cloudflared to return hostname" - assert resp.json()["ip"] != "", "Expected cloudflared to return ip" - assert resp.json()["connector_id"] == connector_id, "Expected cloudflared to return connector_id" - def test_tunnel_url(self, tmp_path, component_tests_config): config = component_tests_config(cfd_mode=CfdModes.NAMED, run_proxy_dns=False, provide_ingress=False) diff --git a/connection/quic.go b/connection/quic.go index c1ccf2f0..df7188e3 100644 --- a/connection/quic.go +++ b/connection/quic.go @@ -449,6 +449,14 @@ func (hrw *httpResponseAdapter) WriteRespHeaders(status int, header http.Header) return hrw.WriteConnectResponseData(nil, metadata...) } +func (hrw *httpResponseAdapter) Write(p []byte) (int, error) { + // Make sure to send WriteHeader response if not called yet + if !hrw.connectResponseSent { + hrw.WriteRespHeaders(http.StatusOK, hrw.headers) + } + return hrw.RequestServerStream.Write(p) +} + func (hrw *httpResponseAdapter) Header() http.Header { return hrw.headers } diff --git a/management/service.go b/management/service.go index 94925abe..120ab189 100644 --- a/management/service.go +++ b/management/service.go @@ -5,6 +5,7 @@ import ( "fmt" "net" "net/http" + "net/http/pprof" "os" "sync" "time" @@ -74,6 +75,10 @@ func New(managementHostname string, r.Head("/ping", ping) r.Get("/logs", s.logs) r.Get("/metrics", s.metricsHandler.ServeHTTP) + + // Supports only heap and goroutine + r.Get("/debug/pprof/{profile:heap|goroutine}", pprof.Index) + r.Route("/host_details", func(r chi.Router) { // CORS middleware required to allow dash to access management.argotunnel.com requests r.Use(cors.Handler(cors.Options{