Source file `iam_provider.ml`

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
module Config = struct
  type t = {
    refresh_interval : float;
    user_id : string;
    cluster_id : string;
    region : string;
  }

  let default ~user_id ~cluster_id ~region =
    { refresh_interval = 600.0; user_id; cluster_id; region }
end

(* A [registration] wraps a caller-supplied enumerator closure.
   Physical-identity is used for unregister lookup, so the
   record has no other fields — the closure is the payload. *)
type registration = {
  enumerate : unit -> Connection.t list;
}

type t = {
  credentials : Iam_credentials.t;
  config : Config.t;
  clock : float Eio.Time.clock_ty Eio.Resource.t;
  cached_token : string Atomic.t;
  registrations : registration list Atomic.t;
}

let sign_now t =
  Iam_sigv4.presigned_elasticache_token
    ~credentials:t.credentials
    ~region:t.config.region
    ~cluster_id:t.config.cluster_id
    ~user_id:t.config.user_id
    ~now:(Eio.Time.now t.clock)

let current_token t = Atomic.get t.cached_token

(* Collect connections from every registered enumerator, skip
   [Dead _] entries. Duplicates (same Connection.t returned by
   more than one enumerator) are removed via [List.memq] so we
   don't send redundant AUTH commands. Enumeration is typically
   tens of connections at most, so an O(n²) dedup is fine. *)
let collect_live_connections t =
  let regs = Atomic.get t.registrations in
  List.fold_left
    (fun acc r ->
      let enum =
        try r.enumerate () with _ -> []
      in
      List.fold_left
        (fun acc conn ->
          match Connection.state conn with
          | Connection.Dead _ -> acc
          | _ when List.memq conn acc -> acc
          | _ -> conn :: acc)
        acc enum)
    [] regs

let push_auth_to_all t token =
  let live = collect_live_connections t in
  List.iter
    (fun conn ->
      match
        Connection.refresh_auth conn
          ~user:t.config.user_id ~password:token
      with
      | Ok () -> ()
      | Error e -> Observability.record_auth_refresh_failure e)
    live

let force_refresh t =
  let token = sign_now t in
  Atomic.set t.cached_token token;
  push_auth_to_all t token

let spawn_refresh_fiber ~sw t =
  (* [fork_daemon]: refresh-ticking is background-forever work;
     the switch cancels it once all non-daemon fibers (the
     user's actual workload) have exited. A plain [fork] would
     keep the switch alive indefinitely. *)
  Eio.Fiber.fork_daemon ~sw (fun () ->
    let rec loop () =
      Eio.Time.sleep t.clock t.config.refresh_interval;
      (try force_refresh t with
       | Eio.Cancel.Cancelled _ as exn -> raise exn
       | exn ->
           (* Never let the refresh fiber die from a
              transient exception — log and keep ticking.
              Cancellation is the only legitimate exit. *)
           let err =
             Connection.Error.Terminal (Printexc.to_string exn)
           in
           Observability.record_auth_refresh_failure err);
      loop ()
    in
    (try loop ()
     with Eio.Cancel.Cancelled _ -> ());
    `Stop_daemon)

let create ~sw ~clock ~credentials ~config =
  let clock : float Eio.Time.clock_ty Eio.Resource.t =
    (clock :> float Eio.Time.clock_ty Eio.Resource.t)
  in
  (* Eager first sign so [current_token] returns a usable token
     from the moment [create] returns. *)
  let initial =
    Iam_sigv4.presigned_elasticache_token
      ~credentials
      ~region:config.Config.region
      ~cluster_id:config.Config.cluster_id
      ~user_id:config.Config.user_id
      ~now:(Eio.Time.now clock)
  in
  let t = {
    credentials;
    config;
    clock;
    cached_token = Atomic.make initial;
    registrations = Atomic.make [];
  } in
  spawn_refresh_fiber ~sw t;
  t

let auth_provider t =
  Connection.Auth.custom ~name:"iam" (fun () ->
    t.config.user_id, current_token t)

let register t enumerate =
  let reg = { enumerate } in
  let rec loop () =
    let old = Atomic.get t.registrations in
    let next = reg :: old in
    if Atomic.compare_and_set t.registrations old next then reg
    else loop ()
  in
  loop ()

let unregister t reg =
  let rec loop () =
    let old = Atomic.get t.registrations in
    let next = List.filter (fun r -> r != reg) old in
    if Atomic.compare_and_set t.registrations old next then ()
    else loop ()
  in
  loop ()