Allow configuration of query encoding in uri's

Uri.t can now be created with an optional `query_scheme` parameter that
controls the encoding of the query string. If it's omitted then the
encoding defaults to the uri's own scheme.

An implementation is given for query encoding according to amazon's
standards to be used when writing clients for AWS.

Author: rgrinberg

lib/uri.ml

@@ -74,7 +74,7 @@ module type Scheme = sig
   val canonicalize_path : string list -> string list
 end
 
-module Generic : Scheme = struct
+module Generic = struct
   let sub_delims a =
     let subd = "!$&'()*+,;=" in
     for i = 0 to String.length subd - 1 do
@@ -201,12 +201,32 @@ module Urn : Scheme = struct
 
 end
 
+module Aws : Scheme = struct
+  include Http
+
+  let safe_chars_query_key =
+    let a = Array.copy Generic.safe_chars in
+    a.(Char.code '/') <- true;
+    a
+
+  let safe_chars_for_component = function
+    | `Query_key -> safe_chars_query_key
+    | `Query_value -> Generic.safe_chars
+    | `Query -> Http.safe_chars_for_component `Query
+    | `Path
+    | `Userinfo
+    | `Fragment
+    | `Scheme -> failwith "Aws scheme is only for query encoding"
+    | x -> Http.safe_chars_for_component x
+end
+
 let module_of_scheme = function
   | Some s -> begin match String.lowercase s with
       | "http" -> (module Http : Scheme)
       | "https"  -> (module Https : Scheme)
       | "file" -> (module File : Scheme)
       | "urn"  -> (module Urn : Scheme)
+      | "aws" -> (module Aws : Scheme)
       | _ -> (module Generic : Scheme)
     end
   | None -> (module Generic : Scheme)
@@ -509,6 +529,7 @@ let encoded_of_query ?scheme = Query.encoded_of_query ?scheme
 
 (* Type of the URI, with most bits being optional *)
 type t = {
+  query_scheme: Pct.decoded sexp_option;
   scheme: Pct.decoded sexp_option;
   userinfo: Userinfo.t sexp_option;
   host: Pct.decoded sexp_option;
@@ -519,6 +540,7 @@ type t = {
 } with sexp
 
 let empty = {
+  query_scheme = None;
   scheme = None;
   userinfo = None;
   host = None;
@@ -572,7 +594,7 @@ let normalize schem uri =
  * casting/uncasting (which isn't fully identity due to the option box), but it is
  * no big deal for now.
 *)
-let make ?scheme ?userinfo ?host ?port ?path ?query ?fragment () =
+let make ?query_scheme ?scheme ?userinfo ?host ?port ?path ?query ?fragment () =
   let decode = function
     |Some x -> Some (Pct.cast_decoded x) |None -> None in
   let host = match userinfo, host, port with
@@ -593,9 +615,10 @@ let make ?scheme ?userinfo ?host ?port ?path ?query ?fragment () =
     | Some p -> Query.KV p
   in
   let scheme = decode scheme in
+  let query_scheme = decode query_scheme in
   normalize scheme
-    { scheme; userinfo;
-      host=decode host; port; path; query; fragment=decode fragment }
+    { scheme ; query_scheme ; userinfo ; host=decode host ; port ; path
+    ; query ; fragment=decode fragment }
 
 (** Parse a URI string into a structure *)
 let of_string s =
@@ -645,7 +668,8 @@ let of_string s =
     | None -> Query.Raw (None, Lazy.from_val [])
   in
   let fragment = get_opt subs 9 in
-  normalize scheme { scheme; userinfo; host; port; path; query; fragment }
+  normalize scheme { query_scheme=None ; scheme ; userinfo ; host ; port
+                   ; path ; query ; fragment }
 
 (** Convert a URI structure into a percent-encoded string
     <http://tools.ietf.org/html/rfc3986#section-5.3>
@@ -687,6 +711,14 @@ let to_string uri =
     Buffer.add_char buf ':';
     Buffer.add_string buf (string_of_int port)
   );
+  let scheme =
+    match uri.query_scheme with
+    | None -> scheme
+    | Some s ->
+      let s' = Pct.uncast_decoded s in
+      (if s' = "aws" then
+         failwith @@ Printf.sprintf "using scheme %s\n" s');
+      Some (Pct.uncast_decoded s) in
   (match uri.path with (* Handle relative paths correctly *)
   | [] -> ()
   | "/"::_ ->
@@ -725,6 +757,11 @@ let with_scheme uri =
   |Some scheme -> { uri with scheme=Some (Pct.cast_decoded scheme) }
   |None -> { uri with scheme=None }
 
+let with_query_scheme uri =
+  function
+  |Some scheme -> { uri with query_scheme=Some (Pct.cast_decoded scheme) }
+  |None -> { uri with query_scheme=None }
+
 let host uri = get_decoded_opt uri.host
 let with_host uri =
   function
@@ -818,16 +855,21 @@ let remove_query_param uri k = Query.(
   { uri with query=KV (List.filter (fun (k',_) -> k<>k') (kv uri.query)) }
 )
 
+let q_scheme uri =
+  match uri.query_scheme with
+  | None -> uri.scheme
+  | t -> t
+
 (* Construct encoded path and query components *)
 let path_and_query uri =
   match (path uri), (query uri) with
   |"", [] -> "/" (* TODO: What about same document? (/) *)
   |"", q -> (* TODO: What about same document? (/) *)
-    let scheme = uncast_opt uri.scheme in
+    let scheme = uncast_opt (q_scheme uri) in
     Printf.sprintf "/?%s" (encoded_of_query ?scheme q)
   |p, [] -> p
   |p, q ->
-    let scheme = uncast_opt uri.scheme in
+    let scheme = uncast_opt (q_scheme uri) in
     Printf.sprintf "%s?%s" p (encoded_of_query ?scheme q)
 
 (* TODO: functions to add and remove from a URI *)

lib/uri.mli

@@ -79,7 +79,7 @@ val canonicalize : t -> t
     supplied without host, an empty host is added. If path is supplied
     and userinfo, host, or port is also supplied, path is made
     absolute but not resolved. *)
-val make : ?scheme:string -> ?userinfo:string -> ?host:string ->
+val make : ?query_scheme:string -> ?scheme:string -> ?userinfo:string -> ?host:string ->
   ?port:int -> ?path:string -> ?query:(string * string list) list ->
   ?fragment:string -> unit -> t
 
@@ -192,6 +192,8 @@ val scheme : t -> string option
     Input URI is not modified *)
 val with_scheme : t -> string option -> t
 
+val with_query_scheme : t -> string option -> t
+
 (** Get the userinfo component of a URI *)
 val userinfo : t -> string option
 

lib_test/test_runner.ml

@@ -20,7 +20,7 @@ open OUnit
 open Printf
 
 (* Tuples of decoded and encoded strings. The first element is a number to
-   refer to the test, as the pcts_large version duplicates the second field 
+   refer to the test, as the pcts_large version duplicates the second field
    to a large size, so it cant be used as the name of the test *)
 let pcts = [
   (1, "hello world!", "hello%20world!");
@@ -46,7 +46,7 @@ let pcts_large =
     done;
     (n, Buffer.contents a', Buffer.contents b')
   ) pcts
- 
+
 (* Tuple of string URI and the decoded version *)
 let uri_encodes = [
   "https://user:pass@foo.com:123/wh/at/ever?foo=1&bar=5#5",
@@ -86,7 +86,7 @@ let map_pcts_tests size name test args =
   ) args
 
 let test_pct_small =
-  (map_pcts_tests "small" "encode" (fun a b -> b, (Uri.pct_encode a)) pcts) @ 
+  (map_pcts_tests "small" "encode" (fun a b -> b, (Uri.pct_encode a)) pcts) @
   (map_pcts_tests "small" "decode" (fun a b -> (Uri.pct_decode b), a) pcts)
 
 let test_pct_large =
@@ -608,6 +608,50 @@ let test_canonicalize =
     )
   ) canonical_map
 
+let aws_encode_val =
+  let safe_chars =
+    let a = Array.make 256 false in
+    let always_safe =
+      "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789_.-~" in
+    for i = 0 to String.length always_safe - 1 do
+      let c = Char.code always_safe.[i] in
+      a.(c) <- true
+    done;
+    a
+  in
+  fun s ->
+    let len = (String.length s) in
+    let b = Buffer.create len in
+    for i = 0 to len - 1 do
+      let code = Char.code s.[i] in
+      if safe_chars.(code) then
+        Buffer.add_char b s.[i]
+      else
+        Buffer.add_string b (Printf.sprintf "%%%02X" code);
+    done;
+    Buffer.contents b
+
+let aws_encode kv =
+  kv
+  |> List.map (fun (x, y) -> x ^ "=" ^ aws_encode_val y)
+  |> String.concat "&"
+
+let test_aws_encode =
+  let f q =
+    let query = q |> List.map (fun (k, v) -> (k, [v])) in
+    let u = Uri.make ~query_scheme:"aws" ~query () in
+    Uri.path_and_query u
+  in
+  let test q =
+    let q' = "/?" ^ aws_encode q in
+    (q, q')
+  in
+  [ test ["q","*"]
+  ; test ["q","star wars"; "q.options", "{fields: ['title^5.0','description']}"]
+  ] |> List.map (fun (in_, out) ->
+    out >:: (fun () -> assert_equal ~printer:(fun x -> x) out (f in_))
+  )
+
 (* Returns true if the result list contains successes only.
    Copied from oUnit source as it isnt exposed by the mli *)
 let rec was_successful =
@@ -639,6 +683,7 @@ let _ =
     @ test_sexping
     @ test_with_change
     @ test_canonicalize
+    @ test_aws_encode
   ) in
   let verbose = ref false in
   let set_verbose _ = verbose := true in