Skip to content

Commit

Permalink
Add /tracer_flare/v1 endpoint with basic form-field validation.
Browse files Browse the repository at this point in the history
  • Loading branch information
mcculls committed Nov 6, 2023
1 parent 07e1de9 commit 059c9a5
Show file tree
Hide file tree
Showing 4 changed files with 76 additions and 0 deletions.
29 changes: 29 additions & 0 deletions ddapm_test_agent/agent.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import argparse
from asyncio import StreamReader
import atexit
import base64
from collections import OrderedDict
Expand All @@ -22,6 +23,7 @@
from urllib.parse import urlunparse

from aiohttp import ClientSession
from aiohttp import MultipartReader
from aiohttp import web
from aiohttp.web import Request
from aiohttp.web import middleware
Expand Down Expand Up @@ -193,6 +195,7 @@ def __init__(self):
"/v0.7/config",
"/telemetry/proxy/api/v2/apmtelemetry",
"/v0.1/pipeline_stats",
"/tracer_flare/v1",
]

async def traces(self) -> TraceMap:
Expand Down Expand Up @@ -512,6 +515,29 @@ async def handle_v2_apmtelemetry(self, request: Request) -> web.Response:
# TODO: Snapshots
return web.HTTPOk()

async def handle_v1_tracer_flare(self, request: Request) -> web.Response:
# reconstruct stream from previously cached bytes
stream = StreamReader()
stream.feed_data(self._request_data(request))
stream.feed_eof()

# validate we have all the required fields
remainingFields = {"source", "case_id", "email", "hostname", "flare_file"}
async for part in MultipartReader(request.headers, stream):
if part.name is not None:
remainingFields.discard(part.name)
if part.name == "flare_file":
request["_flare_file"] = await part.read() # zipfile
else:
request["_flare_" + part.name] = await part.text()

if len(remainingFields) == 0:
return web.HTTPOk()
else:
msg = f"Flare request is missing {','.join(remainingFields)}"
log.error(msg)
return web.HTTPBadRequest(text=msg)

async def handle_put_tested_integrations(self, request: Request) -> web.Response:
# we need to store the request manually since this is not a real DD agent endpoint
await self._store_request(request)
Expand Down Expand Up @@ -570,6 +596,7 @@ async def handle_info(self, request: Request) -> web.Response:
"/v0.6/stats",
"/telemetry/proxy/",
"/v0.7/config",
"/tracer_flare/v1",
],
"feature_flags": [],
"config": {},
Expand Down Expand Up @@ -767,6 +794,7 @@ async def handle_session_requests(self, request: Request) -> web.Response:
self.handle_v2_apmtelemetry,
self.handle_v1_profiling,
self.handle_v07_remoteconfig,
self.handle_v1_tracer_flare,
):
continue
resp.append(
Expand Down Expand Up @@ -989,6 +1017,7 @@ def make_app(
web.post("/v0.7/config", agent.handle_v07_remoteconfig),
web.post("/telemetry/proxy/api/v2/apmtelemetry", agent.handle_v2_apmtelemetry),
web.post("/profiling/v1/input", agent.handle_v1_profiling),
web.post("/tracer_flare/v1", agent.handle_v1_tracer_flare),
web.get("/info", agent.handle_info),
web.get("/test/session/start", agent.handle_session_start),
web.get("/test/session/clear", agent.handle_session_clear),
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
---
features:
- |
Add ``/tracer_flare/v1`` endpoint with basic form-field validation.
1 change: 1 addition & 0 deletions tests/test_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ async def test_info(agent):
"/v0.6/stats",
"/telemetry/proxy/",
"/v0.7/config",
"/tracer_flare/v1",
],
"feature_flags": [],
"config": {},
Expand Down
42 changes: 42 additions & 0 deletions tests/test_tracerflare.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
from aiohttp import FormData


async def test_tracerflare(agent):
form = FormData()
form.add_field("source", "tracer_test")
form.add_field("case_id", "12345")
form.add_field("email", "its.me@datadoghq.com")
form.add_field("hostname", "my.hostname")
form.add_field(
"flare_file",
b"PK\x05\x06\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",
filename="test-flare.zip",
content_type="application/octet-stream",
)
resp = await agent.post("/tracer_flare/v1", data=form)
assert resp.status == 200


async def test_tracerflare_missing_case_id(agent):
form = FormData()
form.add_field("source", "tracer_test")
form.add_field("email", "its.me@datadoghq.com")
form.add_field("hostname", "my.hostname")
form.add_field(
"flare_file",
b"PK\x05\x06\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",
filename="test-flare.zip",
content_type="application/octet-stream",
)
resp = await agent.post("/tracer_flare/v1", data=form)
assert resp.status == 400


async def test_tracerflare_missing_file(agent):
form = FormData()
form.add_field("source", "tracer_test")
form.add_field("case_id", "12345")
form.add_field("email", "its.me@datadoghq.com")
form.add_field("hostname", "my.hostname")
resp = await agent.post("/tracer_flare/v1", data=form)
assert resp.status == 400

0 comments on commit 059c9a5

Please sign in to comment.