handle traceparent and tracestate headers in web

devkit.opam

@@ -40,6 +40,6 @@ depopts: [
 ]
 conflicts: [
   "jemalloc" {< "0.2"}
-  "opentelemetry" {< "0.13"}
+  "opentelemetry" {< "0.92"}
 ]
 available: arch != "arm32" & arch != "x86_32"

possibly_otel.mli

@@ -1,9 +1,19 @@
 module Otrace := Trace_core
 
+module Trace_context : sig
+  val header_names : string list
+  (** Names of the W3C trace context headers ([traceparent], [tracestate]) *)
+
+  val get_ambient_headers : ?explicit_span:Trace_core.span -> unit -> (string * string) list
+  (** Headers propagating the context of [explicit_span], or else of the ambient span:
+      [traceparent] with the span's sampled flag, and [tracestate] if non-empty.
+      Empty if there is no OTEL span. *)
+end
+
 module Traceparent : sig
   val name : string
   val get_ambient : ?explicit_span:Trace_core.span -> unit -> string option
-end
+end [@@deprecated "use Trace_context, which also propagates tracestate"]
 
 val enter_manual_span :
   __FUNCTION__:string ->

possibly_otel.real.ml

@@ -1,18 +1,32 @@
 open Opentelemetry
 
+let ambient_span ?explicit_span () =
+  match explicit_span, Trace_core.current_span () with
+  | Some (Opentelemetry_trace.Extensions.Span_otel span), _
+  | _, Some (Opentelemetry_trace.Extensions.Span_otel span) -> Some span
+  | _ -> None
+
+module Trace_context = struct
+  let header_names = Opentelemetry.Trace_context.[ Traceparent.name; Tracestate.name ]
+
+  let get_ambient_headers ?explicit_span () =
+    match ambient_span ?explicit_span () with
+    | Some span -> Opentelemetry.Trace_context.headers_of_span_ctx (Span.to_span_ctx span)
+    | None -> []
+end
+
 module Traceparent = struct
-  let name = Trace_context.Traceparent.name
+  let name = Opentelemetry.Trace_context.Traceparent.name
 
   let get_ambient ?explicit_span () =
-    match explicit_span, Trace_core.current_span () with
-    | Some (Opentelemetry_trace.Extensions.Span_otel span), _
-    | _, Some (Opentelemetry_trace.Extensions.Span_otel span) ->
-        let ctx = Trace_context.Traceparent.to_value
-          ~trace_id:(Span.trace_id span) ~parent_id:(Span.id span) () in
-        Some ctx
-    | _ -> None
+    match ambient_span ?explicit_span () with
+    | Some span ->
+        Some (Opentelemetry.Trace_context.Traceparent.to_value
+          ~trace_flags:(Span.trace_flags span) ~trace_id:(Span.trace_id span) ~parent_id:(Span.id span) ())
+    | None -> None
 end
 
+(* no explicit [~parent]: the current span if any, otherwise let the collector look for
+   an ambient context (e.g. from an incoming traceparent) instead of forcing a new trace *)
 let enter_manual_span ~__FUNCTION__ ~__FILE__ ~__LINE__ ?data name =
-    let parent = Trace_core.current_span () in
-    Trace_core.enter_span ~parent ~__FUNCTION__ ~__FILE__ ~__LINE__ ?data name
+    Trace_core.enter_span ~__FUNCTION__ ~__FILE__ ~__LINE__ ?data name

possibly_otel.stub.ml

@@ -1,3 +1,9 @@
+module Trace_context = struct
+  let header_names = [ "traceparent"; "tracestate" ]
+
+  let get_ambient_headers ?explicit_span:_ () = []
+end
+
 module Traceparent = struct
   let name = "traceparent"
 

web.ml

@@ -40,6 +40,17 @@ let parse_url_args args =
 
 let url_get_args url = try String.split url "?" |> snd |> parse_url_args with _ -> []
 
+(** Header name is a non-empty HTTP token. See https://www.rfc-editor.org/rfc/rfc9110.html#name-collected-abnf *)
+let valid_header_name n =
+  n <> "" && Stdlib.String.for_all (function
+    | 'a'..'z' | 'A'..'Z' | '0'..'9'
+    | '!' | '#' | '$' | '%' | '&' | '\'' | '*' | '+' | '-' | '.' | '^' | '_' | '`' | '|' | '~' -> true
+    | _ -> false) n
+
+(** Header value is not too long and has no CR, LF or NUL (which could inject into the request) *)
+let valid_header_value v =
+  String.length v <= 8192 && Stdlib.String.for_all (function '\r' | '\n' | '\000' -> false | _ -> true) v
+
 let () = Curl.global_init Curl.CURLINIT_GLOBALALL
 
 let curl_times h =
@@ -119,6 +130,9 @@ module type CURL = sig
   val perform : Curl.t -> Curl.curlCode t
 end
 
+(** [?headers] are sent as given, with one exception: if the request has an OpenTelemetry span, any [traceparent]
+    and [tracestate] headers are removed (names match case-insensitively) and replaced with ones generated from that span.
+    To continue an incoming trace, make it the ambient context; don't pass these headers yourself. *)
 type ('body,'ret) http_request_ =
   ?verbose:bool ->
   ?ua:string ->
@@ -324,17 +338,29 @@ module Http (IO : IO_TYPE) (Curl_IO : CURL with type 'a t = 'a IO.t) : HTTP with
       in
       log #info ~pairs "http error"
 
+  let validate_header (n, v) =
+    if not (valid_header_name n) then invalid_arg (sprintf "Web: invalid header name %S" n);
+    if not (valid_header_value v) then invalid_arg (sprintf "Web: invalid value for header %s: %S" n (Stre.shorten 64 v));
+    ()
+
   let verbose_curl_result nr_http action t hr =
     match Log.State.get_cur_format () with
     | `Plain, _ -> verbose_curl_result_plain nr_http action t hr
     | `Logfmt, _ -> verbose_curl_result_logfmt nr_http action t hr
 
-  (* Given a list of strings, check pre-existing entry starting with `~name`; and adds the concatenation of `~name` and `~value` if not. *)
-  let add_if_absent ~name ~value strs =
-    match strs with
-    | Some strs when List.exists (fun s -> Stre.starts_with s (name^":")) strs -> strs
-    | Some strs -> (String.concat ": " [name; value]) :: strs
-    | None -> [String.concat ": " [name; value]]
+  (* Remove from the list of "Name: value" headers those whose name (case-insensitive) is in `names`, then add `pairs`.
+     Also handles curl's "Name;" syntax for sending an empty header. *)
+  let replace_headers ~names pairs strs =
+    let header_name s =
+      let i = try String.index s ':' with Not_found -> String.length s in
+      let i = try min i (String.index s ';') with Not_found -> i in
+      String.strip (String.sub s 0 i)
+    in
+    let is_replaced s = let n = header_name s in List.exists (Stre.iequal n) names in
+    let strs = Option.default [] strs in
+    let strs = if List.exists is_replaced strs then List.filter (fun s -> not (is_replaced s)) strs else strs in
+    let header (n, v) = validate_header (n, v); sprintf "%s: %s" n v in
+    List.map header pairs @ strs
 
   (* NOTE don't forget to set http_1_0=true when sending requests to a Httpev-based server *)
   (* Don't use curl_setheaders when using ?headers option *)
@@ -414,9 +440,12 @@ module Http (IO : IO_TYPE) (Curl_IO : CURL with type 'a t = 'a IO.t) : HTTP with
       Possibly_otel.enter_manual_span
         ~__FUNCTION__:"Devkit.Web.Http.http_request_k" ~__FILE__ ~__LINE__ ~data:describe span_name in
 
-    let headers = match Possibly_otel.Traceparent.get_ambient ~explicit_span () with
-    | None -> headers
-    | Some value -> Some (add_if_absent ~name:(Possibly_otel.Traceparent.name) ~value headers)
+    (* trace context headers always override the caller's *)
+    let headers = match Possibly_otel.Trace_context.get_ambient_headers ~explicit_span () with
+    | [] -> headers
+    | pairs ->
+      try Some (replace_headers ~names:Possibly_otel.Trace_context.header_names pairs headers)
+      with exn -> Trace_core.exit_span explicit_span; Stdlib.raise exn
     in
 
     let set_body_and_headers h ct body =