Source file resolver.ml

(*
 * Copyright (c) 2014 Anil Madhavapeddy <anil@recoil.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *
 *)

open Sexplib.Std
open Astring

type service = { name : string; port : int; tls : bool } [@@deriving sexp]

(** Module type for a {{!resolution} resolver} that can map URIs to concrete
    {{!Conduit.endp} endpoints} that stream connections can be established with. *)
module type S = sig
  type +'a io
  (** Abstract type of the cooperative threading library used, normally defined
      via the {!IO} module type *)

  type t [@@deriving sexp]
  (** State handle for a running resolver *)

  type svc [@@deriving sexp]
  (** Abstract type for a service entry, which maps a URI scheme into a protocol
      handler and TCP port *)

  type rewrite_fn = svc -> Uri.t -> Conduit.endp io
  (** A rewrite function resolves a {{!svc} service} and a URI into a concrete
      endpoint. *)

  type service_fn = string -> svc option io
  (** A service function maps the string (such as [http] or [ftp]) from a URI
      scheme into a {{!svc} service} description that includes enough metadata
      about the service to subsequently {{!rewrite_fn} resolve} it into an
      {{!Conduit.endp} endpoint}. *)

  val ( ++ ) : service_fn -> service_fn -> service_fn

  val init :
    ?service:service_fn -> ?rewrites:(string * rewrite_fn) list -> unit -> t
  (** [init ?service ?rewrites] will initialize the resolver and return a state
      handler. The {{!service_fn} service} argument should contain the
      system-specific resolution mechanism for URI schemas.

      The [rewrites] argument can optionally override a subset of the URI domain
      name with the given {!rewrite_fn} to permit custom resolution rules. For
      example, a rewrite rule for ".xen" would let the rewrite function resolve
      hostnames such as "foo.xen" into a shared memory channel for the "foo"
      virtual machine. *)

  val add_rewrite : host:string -> f:rewrite_fn -> t -> unit
  (** [add_rewrite ~host f t] will add to the [t] resolver the [f] rewrite rule
      for all the domain names that shortest-prefix match [host] *)

  val set_service : f:service_fn -> t -> unit
  val service : t -> service_fn

  val resolve_uri :
    ?rewrites:(string * rewrite_fn) list -> uri:Uri.t -> t -> Conduit.endp io
  (** [resolve_uri ?rewrites ~uri t] will use [t] to resolve the [uri] into a
      concrete endpoint. Any [rewrites] that are passed in will be overlayed on
      the existing rules within the [t] resolver, but not otherwise modify it. *)
end

module Make (IO : Conduit.IO) = struct
  open IO

  type svc = service [@@deriving sexp]
  type 'a io = 'a IO.t

  type rewrite_fn = service -> Uri.t -> Conduit.endp IO.t [@@deriving sexp]
  (** A rewrite modifies an input URI with more specialization towards a
      concrete [endp] *)

  type service_fn = string -> service option IO.t [@@deriving sexp]

  type t = {
    default_lookup : rewrite_fn;
    mutable domains : rewrite_fn Conduit_trie.t;
    mutable service : service_fn;
  }
  [@@deriving sexp]

  let default_lookup _ uri =
    (* TODO log *)
    let host = match Uri.host uri with None -> "" | Some host -> host in
    return (`Unknown host)

  let default_service _name =
    (* TODO log *)
    return None

  let host_to_domain_list host =
    (* TODO: slow, specialise the Trie to be a rev string list instead *)
    String.concat ~sep:"." (List.rev (String.cuts ~sep:"." host))

  let add_rewrite ~host ~f t =
    t.domains <- Conduit_trie.insert (host_to_domain_list host) f t.domains

  let set_service ~f t = t.service <- f
  let service t = t.service
  let ( ++ ) f g h = f h >>= function None -> g h | x -> return x

  let init ?(service = default_service) ?(rewrites = []) () =
    let domains = Conduit_trie.empty in
    let t = { domains; default_lookup; service } in
    List.iter (fun (host, f) -> add_rewrite ~host ~f t) rewrites;
    t

  let resolve_uri ?rewrites ~uri t =
    (* Find the service associated with the URI *)
    match Uri.scheme uri with
    | None -> return (`Unknown "no scheme")
    | Some scheme -> (
        t.service scheme >>= function
        | None -> return (`Unknown "unknown scheme")
        | Some service ->
            let host =
              match Uri.host uri with None -> "localhost" | Some host -> host
            in
            let trie =
              (* If there are local rewrites, add them to the trie *)
              match rewrites with
              | None -> t.domains
              | Some rewrites ->
                  List.fold_left
                    (fun acc (host, f) ->
                      Conduit_trie.insert (host_to_domain_list host) f acc)
                    t.domains rewrites
            in
            (* Find the longest prefix function that matches this host *)
            let fn =
              match
                Conduit_trie.longest_prefix (host_to_domain_list host) trie
              with
              | None -> t.default_lookup
              | Some fn -> fn
            in
            fn service uri >>= fun endp ->
            if service.tls then return (`TLS (host, endp)) else return endp)
end