T7363: add session search and update by pid

4 files changed, 193 insertions, 3 deletions

diff --git a/data/vyconf.proto b/data/vyconf.proto
index f77b035..6ad95d8 100644
--- a/data/vyconf.proto
+++ b/data/vyconf.proto
@@ -18,6 +18,14 @@ message Request {
     optional int32 OnBehalfOf = 3;
   }
 
+  message SessionOfPid {
+    required int32 ClientPid = 1;
+  }
+
+  message SessionUpdatePid {
+    required int32 ClientPid = 1;
+  }
+
   message Teardown {
     optional int32 OnBehalfOf = 1;
   }
@@ -160,6 +168,8 @@ message Request {
     Load load = 24;
     Discard discard = 25;
     SessionChanged session_changed = 26;
+    SessionOfPid session_of_pid = 27;
+    SessionUpdatePid session_update_pid = 28;
   }
 }
 
diff --git a/src/vyconf_pbt.ml b/src/vyconf_pbt.ml
index c3a608f..de50dff 100644
--- a/src/vyconf_pbt.ml
+++ b/src/vyconf_pbt.ml
@@ -16,6 +16,14 @@ type request_setup_session = {
   on_behalf_of : int32 option;
 }
 
+type request_session_of_pid = {
+  client_pid : int32;
+}
+
+type request_session_update_pid = {
+  client_pid : int32;
+}
+
 type request_teardown = {
   on_behalf_of : int32 option;
 }
@@ -153,6 +161,8 @@ type request =
   | Load of request_load
   | Discard of request_discard
   | Session_changed of request_session_changed
+  | Session_of_pid of request_session_of_pid
+  | Session_update_pid of request_session_update_pid
 
 type request_envelope = {
   token : string option;
@@ -193,6 +203,18 @@ let rec default_request_setup_session
   on_behalf_of;
 }
 
+let rec default_request_session_of_pid 
+  ?client_pid:((client_pid:int32) = 0l)
+  () : request_session_of_pid  = {
+  client_pid;
+}
+
+let rec default_request_session_update_pid 
+  ?client_pid:((client_pid:int32) = 0l)
+  () : request_session_update_pid  = {
+  client_pid;
+}
+
 let rec default_request_teardown 
   ?on_behalf_of:((on_behalf_of:int32 option) = None)
   () : request_teardown  = {
@@ -401,6 +423,22 @@ let default_request_setup_session_mutable () : request_setup_session_mutable = {
   on_behalf_of = None;
 }
 
+type request_session_of_pid_mutable = {
+  mutable client_pid : int32;
+}
+
+let default_request_session_of_pid_mutable () : request_session_of_pid_mutable = {
+  client_pid = 0l;
+}
+
+type request_session_update_pid_mutable = {
+  mutable client_pid : int32;
+}
+
+let default_request_session_update_pid_mutable () : request_session_update_pid_mutable = {
+  client_pid = 0l;
+}
+
 type request_teardown_mutable = {
   mutable on_behalf_of : int32 option;
 }
@@ -665,6 +703,18 @@ let rec pp_request_setup_session fmt (v:request_setup_session) =
   in
   Pbrt.Pp.pp_brk pp_i fmt ()
 
+let rec pp_request_session_of_pid fmt (v:request_session_of_pid) = 
+  let pp_i fmt () =
+    Pbrt.Pp.pp_record_field ~first:true "client_pid" Pbrt.Pp.pp_int32 fmt v.client_pid;
+  in
+  Pbrt.Pp.pp_brk pp_i fmt ()
+
+let rec pp_request_session_update_pid fmt (v:request_session_update_pid) = 
+  let pp_i fmt () =
+    Pbrt.Pp.pp_record_field ~first:true "client_pid" Pbrt.Pp.pp_int32 fmt v.client_pid;
+  in
+  Pbrt.Pp.pp_brk pp_i fmt ()
+
 let rec pp_request_teardown fmt (v:request_teardown) = 
   let pp_i fmt () =
     Pbrt.Pp.pp_record_field ~first:true "on_behalf_of" (Pbrt.Pp.pp_option Pbrt.Pp.pp_int32) fmt v.on_behalf_of;
@@ -855,6 +905,8 @@ let rec pp_request fmt (v:request) =
   | Load x -> Format.fprintf fmt "@[<hv2>Load(@,%a)@]" pp_request_load x
   | Discard x -> Format.fprintf fmt "@[<hv2>Discard(@,%a)@]" pp_request_discard x
   | Session_changed x -> Format.fprintf fmt "@[<hv2>Session_changed(@,%a)@]" pp_request_session_changed x
+  | Session_of_pid x -> Format.fprintf fmt "@[<hv2>Session_of_pid(@,%a)@]" pp_request_session_of_pid x
+  | Session_update_pid x -> Format.fprintf fmt "@[<hv2>Session_update_pid(@,%a)@]" pp_request_session_update_pid x
 
 let rec pp_request_envelope fmt (v:request_envelope) = 
   let pp_i fmt () =
@@ -918,6 +970,16 @@ let rec encode_pb_request_setup_session (v:request_setup_session) encoder =
   end;
   ()
 
+let rec encode_pb_request_session_of_pid (v:request_session_of_pid) encoder = 
+  Pbrt.Encoder.int32_as_varint v.client_pid encoder;
+  Pbrt.Encoder.key 1 Pbrt.Varint encoder; 
+  ()
+
+let rec encode_pb_request_session_update_pid (v:request_session_update_pid) encoder = 
+  Pbrt.Encoder.int32_as_varint v.client_pid encoder;
+  Pbrt.Encoder.key 1 Pbrt.Varint encoder; 
+  ()
+
 let rec encode_pb_request_teardown (v:request_teardown) encoder = 
   begin match v.on_behalf_of with
   | Some x -> 
@@ -1242,6 +1304,12 @@ let rec encode_pb_request (v:request) encoder =
   | Session_changed x ->
     Pbrt.Encoder.nested encode_pb_request_session_changed x encoder;
     Pbrt.Encoder.key 26 Pbrt.Bytes encoder; 
+  | Session_of_pid x ->
+    Pbrt.Encoder.nested encode_pb_request_session_of_pid x encoder;
+    Pbrt.Encoder.key 27 Pbrt.Bytes encoder; 
+  | Session_update_pid x ->
+    Pbrt.Encoder.nested encode_pb_request_session_update_pid x encoder;
+    Pbrt.Encoder.key 28 Pbrt.Bytes encoder; 
   end
 
 let rec encode_pb_request_envelope (v:request_envelope) encoder = 
@@ -1344,6 +1412,46 @@ let rec decode_pb_request_setup_session d =
     on_behalf_of = v.on_behalf_of;
   } : request_setup_session)
 
+let rec decode_pb_request_session_of_pid d =
+  let v = default_request_session_of_pid_mutable () in
+  let continue__= ref true in
+  let client_pid_is_set = ref false in
+  while !continue__ do
+    match Pbrt.Decoder.key d with
+    | None -> (
+    ); continue__ := false
+    | Some (1, Pbrt.Varint) -> begin
+      v.client_pid <- Pbrt.Decoder.int32_as_varint d; client_pid_is_set := true;
+    end
+    | Some (1, pk) -> 
+      Pbrt.Decoder.unexpected_payload "Message(request_session_of_pid), field(1)" pk
+    | Some (_, payload_kind) -> Pbrt.Decoder.skip d payload_kind
+  done;
+  begin if not !client_pid_is_set then Pbrt.Decoder.missing_field "client_pid" end;
+  ({
+    client_pid = v.client_pid;
+  } : request_session_of_pid)
+
+let rec decode_pb_request_session_update_pid d =
+  let v = default_request_session_update_pid_mutable () in
+  let continue__= ref true in
+  let client_pid_is_set = ref false in
+  while !continue__ do
+    match Pbrt.Decoder.key d with
+    | None -> (
+    ); continue__ := false
+    | Some (1, Pbrt.Varint) -> begin
+      v.client_pid <- Pbrt.Decoder.int32_as_varint d; client_pid_is_set := true;
+    end
+    | Some (1, pk) -> 
+      Pbrt.Decoder.unexpected_payload "Message(request_session_update_pid), field(1)" pk
+    | Some (_, payload_kind) -> Pbrt.Decoder.skip d payload_kind
+  done;
+  begin if not !client_pid_is_set then Pbrt.Decoder.missing_field "client_pid" end;
+  ({
+    client_pid = v.client_pid;
+  } : request_session_update_pid)
+
 let rec decode_pb_request_teardown d =
   let v = default_request_teardown_mutable () in
   let continue__= ref true in
@@ -1933,6 +2041,8 @@ let rec decode_pb_request d =
       | Some (24, _) -> (Load (decode_pb_request_load (Pbrt.Decoder.nested d)) : request) 
       | Some (25, _) -> (Discard (decode_pb_request_discard (Pbrt.Decoder.nested d)) : request) 
       | Some (26, _) -> (Session_changed (decode_pb_request_session_changed (Pbrt.Decoder.nested d)) : request) 
+      | Some (27, _) -> (Session_of_pid (decode_pb_request_session_of_pid (Pbrt.Decoder.nested d)) : request) 
+      | Some (28, _) -> (Session_update_pid (decode_pb_request_session_update_pid (Pbrt.Decoder.nested d)) : request) 
       | Some (n, payload_kind) -> (
         Pbrt.Decoder.skip d payload_kind; 
         loop () 
diff --git a/src/vyconf_pbt.mli b/src/vyconf_pbt.mli
index 667a64d..907f78c 100644
--- a/src/vyconf_pbt.mli
+++ b/src/vyconf_pbt.mli
@@ -23,6 +23,14 @@ type request_setup_session = {
   on_behalf_of : int32 option;
 }
 
+type request_session_of_pid = {
+  client_pid : int32;
+}
+
+type request_session_update_pid = {
+  client_pid : int32;
+}
+
 type request_teardown = {
   on_behalf_of : int32 option;
 }
@@ -160,6 +168,8 @@ type request =
   | Load of request_load
   | Discard of request_discard
   | Session_changed of request_session_changed
+  | Session_of_pid of request_session_of_pid
+  | Session_update_pid of request_session_update_pid
 
 type request_envelope = {
   token : string option;
@@ -204,6 +214,18 @@ val default_request_setup_session :
   request_setup_session
 (** [default_request_setup_session ()] is the default value for type [request_setup_session] *)
 
+val default_request_session_of_pid : 
+  ?client_pid:int32 ->
+  unit ->
+  request_session_of_pid
+(** [default_request_session_of_pid ()] is the default value for type [request_session_of_pid] *)
+
+val default_request_session_update_pid : 
+  ?client_pid:int32 ->
+  unit ->
+  request_session_update_pid
+(** [default_request_session_update_pid ()] is the default value for type [request_session_update_pid] *)
+
 val default_request_teardown : 
   ?on_behalf_of:int32 option ->
   unit ->
@@ -397,6 +419,12 @@ val pp_request_prompt : Format.formatter -> request_prompt -> unit
 val pp_request_setup_session : Format.formatter -> request_setup_session -> unit 
 (** [pp_request_setup_session v] formats v *)
 
+val pp_request_session_of_pid : Format.formatter -> request_session_of_pid -> unit 
+(** [pp_request_session_of_pid v] formats v *)
+
+val pp_request_session_update_pid : Format.formatter -> request_session_update_pid -> unit 
+(** [pp_request_session_update_pid v] formats v *)
+
 val pp_request_teardown : Format.formatter -> request_teardown -> unit 
 (** [pp_request_teardown v] formats v *)
 
@@ -496,6 +524,12 @@ val encode_pb_request_prompt : request_prompt -> Pbrt.Encoder.t -> unit
 val encode_pb_request_setup_session : request_setup_session -> Pbrt.Encoder.t -> unit
 (** [encode_pb_request_setup_session v encoder] encodes [v] with the given [encoder] *)
 
+val encode_pb_request_session_of_pid : request_session_of_pid -> Pbrt.Encoder.t -> unit
+(** [encode_pb_request_session_of_pid v encoder] encodes [v] with the given [encoder] *)
+
+val encode_pb_request_session_update_pid : request_session_update_pid -> Pbrt.Encoder.t -> unit
+(** [encode_pb_request_session_update_pid v encoder] encodes [v] with the given [encoder] *)
+
 val encode_pb_request_teardown : request_teardown -> Pbrt.Encoder.t -> unit
 (** [encode_pb_request_teardown v encoder] encodes [v] with the given [encoder] *)
 
@@ -595,6 +629,12 @@ val decode_pb_request_prompt : Pbrt.Decoder.t -> request_prompt
 val decode_pb_request_setup_session : Pbrt.Decoder.t -> request_setup_session
 (** [decode_pb_request_setup_session decoder] decodes a [request_setup_session] binary value from [decoder] *)
 
+val decode_pb_request_session_of_pid : Pbrt.Decoder.t -> request_session_of_pid
+(** [decode_pb_request_session_of_pid decoder] decodes a [request_session_of_pid] binary value from [decoder] *)
+
+val decode_pb_request_session_update_pid : Pbrt.Decoder.t -> request_session_update_pid
+(** [decode_pb_request_session_update_pid decoder] decodes a [request_session_update_pid] binary value from [decoder] *)
+
 val decode_pb_request_teardown : Pbrt.Decoder.t -> request_teardown
 (** [decode_pb_request_teardown decoder] decodes a [request_teardown] binary value from [decoder] *)
 
diff --git a/src/vyconfd.ml b/src/vyconfd.ml
index b0b4e52..edc4b9d 100644
--- a/src/vyconfd.ml
+++ b/src/vyconfd.ml
@@ -45,18 +45,46 @@ let usage = "Usage: " ^ Sys.argv.(0) ^ " [options]"
 
 let response_tmpl = {status=Success; output=None; error=None; warning=None}
 
+let find_session token = Hashtbl.find sessions token
+
+let find_session_by_pid pid =
+    let exception E of string in
+    let find_k k v acc =
+        if v.Session.client_pid = pid then
+            raise_notrace (E k)
+        else acc
+    in
+    try
+        Hashtbl.fold find_k sessions None
+    with E x -> Some x
+
 let make_session_token () =
     Sha1.string (string_of_int (Random.bits ())) |> Sha1.to_hex
 
-let setup_session world req =
+let setup_session world (req: request_setup_session) =
     let token = make_session_token () in
-    let user = "unknown user" in
     let pid = req.client_pid in
+    let user = "unknown user" in
     let client_app = Option.value req.client_application ~default:"unknown client" in
     let () = Hashtbl.add sessions token (Session.make world client_app user pid) in
     {response_tmpl with output=(Some token)}
 
-let find_session token = Hashtbl.find sessions token
+let session_of_pid _world (req: request_session_of_pid) =
+    let pid = req.client_pid in
+    let extant = find_session_by_pid pid in
+    {response_tmpl with output=extant}
+
+let session_update_pid _world token (req: request_session_update_pid) =
+    let pid = req.client_pid in
+    try
+        begin
+        let s = Hashtbl.find sessions token in
+        if s.client_pid <> pid then
+            let session = {s with client_pid=pid} in
+            Hashtbl.replace sessions token session
+        end;
+        {response_tmpl with output=(Some token)}
+    with Not_found -> {response_tmpl with status=Fail; output=None}
 
 let enter_conf_mode req token =
     let open Session in
@@ -261,8 +289,10 @@ let rec handle_connection world ic oc () =
                     (match req with
                     | _, Prompt -> response_tmpl
                     | _, Setup_session r -> setup_session world r
+                    | _, Session_of_pid r -> session_of_pid world r
                     | _, Reload_reftree r -> reload_reftree world r
                     | None, _ -> {response_tmpl with status=Fail; output=(Some "Operation requires session token")}
+                    | Some t, Session_update_pid r -> session_update_pid world t r
                     | Some t, Teardown _ -> teardown t
                     | Some t, Enter_configuration_mode r -> enter_conf_mode r t
                     | Some t, Exit_configuration_mode -> exit_conf_mode world t