diff --git a/ocaml/xapi/wlb_reports.ml b/ocaml/xapi/wlb_reports.ml
index b97331c1008..719ca2d5d62 100644
--- a/ocaml/xapi/wlb_reports.ml
+++ b/ocaml/xapi/wlb_reports.ml
@@ -98,8 +98,8 @@ open Stdext.Xstringext
module D = Debug.Make(struct let name="wlb_reports" end)
open D
-let report_tokens = ("", "")
-let diagnostics_tokens = ("", "")
+let report_tag = "XmlDataSet"
+let diagnostics_tag = "DiagnosticData"
let bufsize = 16384
@@ -107,13 +107,8 @@ let hex_entity s =
(*debug "hex_entity %s" s; *)
char_of_int (int_of_string ("0" ^ (String.sub s 1 (String.length s - 1))))
-let trim_and_send method_name (start_str, end_str) recv_sock send_sock =
+let trim_and_send method_name tag recv_sock send_sock =
let recv_buf = Buffer.create bufsize in
- let send_buf = Buffer.create bufsize in
- let recv_state = ref 1 in
- let send_state = ref 1 in
- let entity = ref "" in
-
let fill () =
let s = String.create bufsize in
let n = Unix.read recv_sock s 0 bufsize in
@@ -121,106 +116,41 @@ let trim_and_send method_name (start_str, end_str) recv_sock send_sock =
Buffer.add_string recv_buf (String.sub s 0 n);
n
in
-
+ (* Since we use xml parser to parse the reponse message, we don't need to escape the xml content in `send` *)
let send s =
- let s_len = String.length s in
- let rec send' i =
- let c = s.[i] in
- (* debug "%c" c; *)
- if !send_state = 1 then
- begin
- if c = '&' then
- send_state := 2
- else
- Buffer.add_char send_buf c
- end
- else
- begin
- if c = ';' then
- let e = !entity in
- Buffer.add_char send_buf
- (if e = "lt" then
- '<'
- else if e = "gt" then
- '>'
- else if e = "amp" then
- '&'
- else if e = "apos" then
- '\''
- else if e = "quot" then
- '"'
- else
- hex_entity e);
- send_state := 1;
- entity := ""
- else
- entity := !entity ^ (String.of_char c)
- end;
- if i < s_len - 1 then
- send' (i + 1)
- else
- ()
- in
- send' 0;
- ignore (Unix.write send_sock (Buffer.contents send_buf) 0
- (Buffer.length send_buf));
- Buffer.clear send_buf
+ ignore (Unix.write send_sock s 0 (String.length s))
in
-
- let rec pump () =
+ let rec recv_all ()=
let n = fill() in
- if Buffer.length recv_buf > 0 then
- begin
- let s = Buffer.contents recv_buf in
- (* debug "%s %d" s !recv_state; *)
- if !recv_state = 1 then
- match String.find_all start_str s with
- | n :: _ ->
- Buffer.clear recv_buf;
- let i = n + String.length start_str in
- Buffer.add_substring recv_buf s i (String.length s - i);
- recv_state := 2
- | [] ->
- ()
- else if !recv_state = 2 then
- match String.find_all end_str s with
- | n :: _ ->
- send (String.sub s 0 n);
- Buffer.clear recv_buf;
- recv_state := 3
- | [] ->
- send s;
- Buffer.clear recv_buf
- else
- Buffer.clear recv_buf;
- if n > 0 then
- pump()
- else if !recv_state != 3 then
- (* if in state 1 we are still looking for the opening tag of the data set, expect xml to be valid
- if in state 2 we are still looking for the closing tag of the data set, expect xml to be truncated *)
- let rec_data = (Buffer.contents recv_buf) in
- if !recv_state = 1 then
- begin
- try
- let xml_data = Xml.parse_string rec_data in
- Workload_balancing.parse_result_code
- method_name
- (Workload_balancing.retrieve_inner_xml method_name xml_data true)
- "Failed to detect end of XML, data could be truncated"
- rec_data
- true
- with
- | Xml.Error err ->
- Workload_balancing.raise_malformed_response' method_name (Xml.error err) rec_data
- end
- else
- Workload_balancing.raise_malformed_response' method_name "Expected data is truncated." rec_data
- end
+ if n > 0 then
+ recv_all()
+ else
+ ()
in
- pump()
-
-
-let handle req bio method_name tokens (method_name, request_func) =
+ recv_all();
+ let s = Buffer.contents recv_buf in
+ debug "receive len: %d, content: %s" (String.length s) s;
+ try
+ let xml_data = Xml.parse_string s in
+ let report_result_xml = Workload_balancing.retrieve_inner_xml method_name xml_data true in
+ try
+ let xml_data_set_content = Workload_balancing.data_from_leaf (Workload_balancing.descend_and_match [tag] report_result_xml) in
+ debug "send conent: %s" xml_data_set_content;
+ send xml_data_set_content
+ with
+ | Workload_balancing.Xml_parse_failure error ->
+ Workload_balancing.parse_result_code
+ method_name
+ report_result_xml
+ "Failed to detect end of XML, data could be truncated"
+ s
+ true
+ with
+ | Xml.Error err ->
+ Workload_balancing.raise_malformed_response' method_name "Expected data is truncated." s
+
+
+let handle req bio method_name tag (method_name, request_func) =
let client_sock = Buf_io.fd_of bio in
Buf_io.assert_buffer_empty bio;
debug "handle: fd = %d" (Stdext.Unixext.int_of_file_descr client_sock);
@@ -236,7 +166,7 @@ let handle req bio method_name tokens (method_name, request_func) =
let parse response wlb_sock =
Http_svr.headers client_sock (Http.http_200_ok ());
- trim_and_send method_name tokens wlb_sock client_sock
+ trim_and_send method_name tag wlb_sock client_sock
in
try
request_func ~__context ~handler:parse
@@ -267,11 +197,11 @@ let report_handler (req: Request.t) (bio: Buf_io.t) _ =
not (List.mem k ["session_id"; "task_id"; "report"]))
req.Request.query
in
- handle req bio "ExecuteReport" report_tokens
+ handle req bio "ExecuteReport" report_tag
(Workload_balancing.wlb_report_request report params)
(* GET /wlb_diagnostics?session_id=&task_id= *)
let diagnostics_handler (req: Request.t) (bio: Buf_io.t) _ =
- handle req bio "GetDiagnostics" diagnostics_tokens
+ handle req bio "GetDiagnostics" diagnostics_tag
Workload_balancing.wlb_diagnostics_request