From 8de41696308e53d921b4fff2c644b745362f355a Mon Sep 17 00:00:00 2001
From: Aleksey Kashapov <nanodirijabl@gmail.com>
Date: Thu, 5 Feb 2026 15:37:45 +0300
Subject: [PATCH] Adds OTEL middleware for swag's cowboy server

- Moves OTEL's text map extractor to middleware
- Implements `X-Amzn-Trace-Id` header's parsing as remote OTEL's span context
---
 apps/capi/src/capi_handler.erl         | 11 +---
 apps/capi/src/capi_otel_middleware.erl | 87 ++++++++++++++++++++++++++
 apps/capi/src/capi_swagger_server.erl  |  3 +-
 3 files changed, 90 insertions(+), 11 deletions(-)
 create mode 100644 apps/capi/src/capi_otel_middleware.erl

diff --git a/apps/capi/src/capi_handler.erl b/apps/capi/src/capi_handler.erl
index f4e56ee..0418a0f 100644
--- a/apps/capi/src/capi_handler.erl
+++ b/apps/capi/src/capi_handler.erl
@@ -66,8 +66,7 @@
 
 -spec authorize_api_key(operation_id(), swag_server:api_key(), request_context(), handler_opts()) ->
     Result :: false | {true, capi_auth:preauth_context()}.
-authorize_api_key(OperationID, ApiKey, Context, _HandlerOpts) ->
-    ok = set_otel_context(Context),
+authorize_api_key(OperationID, ApiKey, _Context, _HandlerOpts) ->
     %% Since we require the request id field to create a woody context for our trip to token_keeper
     %% it seems it is no longer possible to perform any authorization in this method.
     %% To gain this ability back be would need to rewrite the swagger generator to perform its
@@ -296,14 +295,6 @@ process_general_error(Class, Reason, Stacktrace, Req, SwagContext) ->
     ),
     {error, server_error(500)}.
 
-set_otel_context(#{cowboy_req := Req}) ->
-    Headers = cowboy_req:headers(Req),
-    %% Implicitly puts OTEL context into process dictionary.
-    %% Since cowboy does not reuse process for other requests, we don't care
-    %% about cleaning it up.
-    _OtelCtx = otel_propagator_text_map:extract(maps:to_list(Headers)),
-    ok.
-
 -spec set_context_meta(processing_context()) -> ok.
 set_context_meta(Context) ->
     AuthContext = capi_handler_utils:get_auth_context(Context),
diff --git a/apps/capi/src/capi_otel_middleware.erl b/apps/capi/src/capi_otel_middleware.erl
new file mode 100644
index 0000000..a83adbd
--- /dev/null
+++ b/apps/capi/src/capi_otel_middleware.erl
@@ -0,0 +1,87 @@
+-module(capi_otel_middleware).
+
+-behaviour(cowboy_middleware).
+
+-export([execute/2]).
+
+-behaviour(otel_propagator_text_map).
+
+-export([fields/1]).
+-export([inject/4]).
+-export([extract/5]).
+
+-spec execute(Req, Env) -> {ok, Req, Env} | {stop, Req} when Req :: cowboy_req:req(), Env :: cowboy_middleware:env().
+execute(#{headers := Headers} = Req, Env) ->
+    Propagator =
+        {otel_propagator_text_map_composite, [
+            %% First try getting OTEL context from x-ray trace header
+            ?MODULE,
+            %% But if trace context headers are present, then get current span context from there
+            otel_propagator:builtin_to_module(tracecontext)
+        ]},
+    OtelCtx = otel_propagator_text_map:extract_to(otel_ctx:new(), Propagator, maps:to_list(Headers)),
+    %% Implicitly puts OTEL context into process dictionary.
+    %% Since cowboy does not reuse process for other requests, we don't care
+    %% about cleaning it up.
+    _Token = otel_ctx:attach(OtelCtx),
+    {ok, Req, Env}.
+
+-define(HEADER_KEY, <<"x-amzn-trace-id">>).
+
+-spec fields(otel_propagator_text_map:propagator_options()) -> [unicode:latin1_binary()].
+fields(_) ->
+    [?HEADER_KEY].
+
+-spec inject(
+    otel_ctx:t(),
+    otel_propagator:carrier(),
+    otel_propagator_text_map:carrier_set(),
+    otel_propagator_text_map:propagator_options()
+) -> no_return().
+inject(_Ctx, _Carrier, _CarrierSet, _Options) ->
+    erlang:error(not_implemented).
+
+-spec extract(
+    otel_ctx:t(),
+    otel_propagator:carrier(),
+    otel_propagator_text_map:carrier_keys(),
+    otel_propagator_text_map:carrier_get(),
+    otel_propagator_text_map:propagator_options()
+) -> otel_ctx:t().
+extract(Ctx, Carrier, _CarrierKeysFun, CarrierGet, _Options) ->
+    case CarrierGet(?HEADER_KEY, Carrier) of
+        undefined ->
+            Ctx;
+        XRayTrace ->
+            case decode(string:trim(XRayTrace)) of
+                undefined ->
+                    Ctx;
+                SpanCtx ->
+                    otel_tracer:set_current_span(Ctx, SpanCtx)
+            end
+    end.
+
+%%
+
+decode(
+    %% NOTE Version is expected to be always single char "1"
+    <<"Root=", _Version:1/binary, "-", Timestamp:8/binary, "-", RootId:24/binary, ";Parent=", ParentId:16/binary,
+        ";Sampled=", Sampled:1/binary, _/binary>>
+) ->
+    try
+        TraceId = binary_to_integer(<<Timestamp/binary, RootId/binary>>, 16),
+        SpanId = binary_to_integer(ParentId, 16),
+        TraceFlags =
+            case Sampled of
+                <<"1">> -> 1;
+                <<"0">> -> 0;
+                _ -> error(badarg)
+            end,
+        otel_tracer:from_remote_span(TraceId, SpanId, TraceFlags)
+    catch
+        %% to integer from base 16 string failed
+        error:badarg ->
+            undefined
+    end;
+decode(_) ->
+    undefined.
diff --git a/apps/capi/src/capi_swagger_server.erl b/apps/capi/src/capi_swagger_server.erl
index e70db7e..06d7253 100644
--- a/apps/capi/src/capi_swagger_server.erl
+++ b/apps/capi/src/capi_swagger_server.erl
@@ -46,7 +46,8 @@ get_cowboy_config(AdditionalRoutes, LogicHandler, SwaggerHandlerOpts) ->
         middlewares => [
             cowboy_router,
             cowboy_cors,
-            cowboy_handler
+            cowboy_handler,
+            capi_otel_middleware
         ],
         stream_handlers => [
             cowboy_access_log_h,