(* O'Caml Freenet Client Protocol client module                     *)
(* by Travis Bemann and Eric Norige                                 *)
(*                                                                  *)
(* This program is free software; you can redistribute it and/or    *)
(* modify it under the terms of the GNU Lesser General Public       *)
(* License as published by the Free Software Foundation; either     *)
(* version 2 of the License, or (at your option) any later version. *)
(*                                                                  *)
(* This program is distributed in the hope that it will be useful,  *)
(* but WITHOUT ANY WARRANTY; without even the implied warranty of   *)
(* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the    *)
(* GNU Lesser General Public License for more details.              *)

type fcp_value = 
    VString of string 
  | VInt of int

type fcp_param = string * fcp_value

let get_param p s = List.assoc s p

let get_string_param p s =
  match get_param p s with
    VString s -> s
  | VInt i -> string_of_int i

let get_int_param p s = 
  match get_param p s with
    VString s -> Strutil.int_of_hex s
  | VInt i -> i

let param_to_string = function
    (t, VString s) ->
      Printf.sprintf "%s=%s\n" t s
  | (t, VInt v) ->
      Printf.sprintf "%s=%X\n" t v

type fcp_mesg = 
    Small of string * fcp_param list
  | Large of string * fcp_param list * string
  | Large_no_data of string * fcp_param list

exception Closed
exception Restarted
exception Finished
exception Uri_error of string
exception Data_not_found
exception Route_not_found of fcp_mesg
exception Key_collision of string
exception Size_error
exception Connect_error
exception Io_error
exception Node_format_error of string (* Node bugs only *)
exception Format_error of string (* Shouldn't happen *)
exception Node_error of string (* Node bugs only *)
exception Unexpected_message of string (* Node bugs only *)

class transact connect = object (self)
  val in_chan = Unix.in_channel_of_descr connect
  val out_chan = Unix.out_channel_of_descr connect
  val mutable opened = true
  method private close =
    if opened then begin
      opened <- false;
      Unix.shutdown connect Unix.SHUTDOWN_ALL
    end else ()
  method cancel =
    (* CURRENTLY this closes the connection, but this may change in
     * the future if persistent FCP connections are introduced and
     * as a result the architecture of this FCP module changes *)
    self#close
  method private send data = 
    if not opened then raise Closed
    else begin
      try output_string out_chan data; flush out_chan
      with Sys_error _ -> raise Io_error
    end
  method private recv len = 
    if not opened then raise Closed
    else begin
      let data = String.create len in
      begin try really_input in_chan data 0 len
      with Sys_error _ -> raise Io_error
      end;
      data
    end
  method private recv_soft len =
    if not opened then raise Closed
    else
      let data = String.create len in
      let rlen =
	try input in_chan data 0 len
	with Sys_error _ -> raise Io_error
      in String.sub data 0 rlen
  method private recv_line =
    if not opened then raise Closed
    else try input_line in_chan
    with Sys_error _ -> raise Io_error
  method private recv_line_strip =
    let l = Strutil.strip (self#recv_line) in
    if l <> "" then l
    else self#recv_line_strip
  method private send_msg = function 
      Small (n,p) -> 
	self#send (Printf.sprintf "%s\n" n);
	List.iter (fun x -> self#send (param_to_string x)) p;
	self#send "EndMessage\n"
    | Large_no_data (n, p) ->
	self#send (Printf.sprintf "%s\n" n);
	List.iter (fun x -> self#send (param_to_string x)) p;
	self#send "Data\n"
    | Large (n, p, d) -> 
	let rec hasdl = function 
	    ("DataLength", _) :: _ -> true
	  | _ :: y -> hasdl y
	  | [] -> false in
	self#send (Printf.sprintf "%s\n" n);
	List.iter (fun x -> self#send (param_to_string x)) p;
	if not (hasdl p) then 
	  self#send (Printf.sprintf "DataLength=%X\n" (String.length d));
	self#send "Data\n";
	self#send d
  method private recv_msg =
    let rec parse_header accum =
      match self#recv_line_strip with
	"EndMessage" -> accum, false
      | "Data" -> accum, true
      | l -> try
	  let (t,s) = Strutil.splitstrip2 l '=' in
	  parse_header ((t, VString s)::accum)
      with Not_found -> parse_header accum in
    let n = self#recv_line_strip in
    let p, d = parse_header [] in
    if not d then Small (n, p)
    else try Large (n, p, self#recv (get_int_param p "Length"))
    with
      Not_found ->
	raise (Node_format_error
		 "Missing Length node message parameter")
    | Failure m ->
	raise (Node_format_error
		 "Length node message parameter is not a number")
	    
  method private handle_other m = 
    try
      self#close;
      match m with 
	Small ("Failed", p) ->
	  raise (Node_error (get_string_param p "Reason"))
      | Small ("FormatError", p) -> 
	  raise (Format_error (get_string_param p "Reason"))
      | Small (n, _) | Large (n, _, _) | Large_no_data (n, _) ->
	  raise (Unexpected_message n)
    with Not_found -> raise (Node_error "No reason specified")
  initializer
    self#send "\000\000\000\002"
end

class request_transact connect_init uri htl = object (self)
  inherit transact connect_init
  val mutable len = 0
  val mutable metadata_len = 0
  val buf = Buffer.create 1024
  val mutable non_buf_len = 0
  method uri = uri
  method htl = htl
  method len = len
  method metadata_len = metadata_len
  method read_len = (Buffer.length buf) + non_buf_len
  method private wait_found =
    match self#recv_msg with
      Small ("DataFound", p) ->
	begin try
	  len <- get_int_param p "DataLength";
	with
	  Not_found ->
	    self#close;
	    raise (Node_format_error
		     "Missing DataLength parameter of DataFound")
	| Failure m ->
	    self#close;
	    raise (Node_format_error
		     ("DataLength parameter of DataFound is not a"^
		      " number"))
	end;
	begin try
	  metadata_len <- get_int_param p "MetadataLength"
	with
	  Not_found ->()
	| Failure m ->
	    self#close;
	    raise (Node_format_error
		     ("MetadataLength parameter of DataFound is not"^
		      " a number"))
	end
    | Small ("Restarted", _) -> self#wait_found
    | Small ("URIError", _) ->	self#close; raise (Uri_error uri)
    | Small ("DataNotFound", _) -> self#close; raise Data_not_found
    | Small ("RouteNotFound", p) -> self#close; 
	raise (Route_not_found (Small ("RouteNotFound", p)));
    | m -> self#handle_other m; assert false
  method recv_block =
    if self#read_len >= len then begin self#close; raise Finished end;
    match self#recv_msg with
      Large ("DataChunk", _, d) ->
	if Buffer.length buf < metadata_len then
	  Buffer.add_string buf d
	else non_buf_len <- non_buf_len + (String.length d);
	d
    | Small ("DataChunk", _) -> self#close; 
	raise (Node_format_error "DataChunk must have data")
    | Small ("Restarted", _) -> Buffer.clear buf; raise Restarted
    | m -> self#handle_other m; assert false
  method recv_metadata =
    while (Buffer.length buf) < metadata_len do
      try ignore self#recv_block
      with Restarted -> ()
    done;
    String.sub (Buffer.contents buf) 0 metadata_len
  method recv_data_prev =
    let start = metadata_len 
    and length = max (Buffer.length buf - metadata_len) 0 in
    String.sub (Buffer.contents buf) start length
  initializer
    let request_msg = Small ("ClientGet", 
			     [("URI", VString uri); ("HopsToLive", VInt htl)]) in
    self#send_msg request_msg;
    self#wait_found
end

type insert_state = Sending | Waiting | Done

class insert_transact connect_init uri htl metadata_len data_len =
  object (self)
    inherit transact connect_init as super
    val mutable metadata_len_cur = 0
    val mutable data_len_cur = 0
    val mutable state = Sending
    method uri = uri
    method htl = htl
    method private close =
      state <- Done;
      super#close
    method metadata_block s =
      if state = Sending then
	let len = String.length s in
	if data_len_cur = 0 then
	  if metadata_len_cur + len <= metadata_len then begin
	    metadata_len_cur <- metadata_len_cur + len;
	    self#send s
	  end else
	    raise (Failure "More metadata given than specified")
	else raise (Failure "Metadata must be specified before data")
      else raise (Failure "Not sending")
    method block s =
      if state = Sending then
	let len = String.length s in
	if metadata_len_cur = metadata_len then
	  if data_len_cur + len <= data_len then begin
	    data_len_cur <- data_len_cur + len;
	    self#send s
	  end else raise (Failure "More data given than specified")
	else raise (Failure "Data must be specified after metadata")
      else raise (Failure "Not sending")
    method private check_done_sending =
      if (metadata_len_cur = metadata_len)
	  && (data_len_cur = data_len)
	  && (state = Sending) then
	state <- Waiting
      else ()
    method wait =
      self#check_done_sending;
      if state <> Waiting then raise (Failure "Not waiting");
      match self#recv_msg with
	Small ("Success", p) ->
	  self#close;
	  begin try Some (get_string_param p "URI")
	  with Not_found ->
	    raise (Node_format_error
		     "Missing URI parameter of Success")
	  end
      | Small ("Pending", _) -> None
      | Small ("URIError", _) -> self#close; raise (Uri_error uri)
      | Small ("Restarted", _) -> None
      | Small ("RouteNotFound", p) -> self#close;
	  raise (Route_not_found (Small ("RouteNotFound", p)));
      | Small ("KeyCollision", p) ->
	  self#close;
	  begin try
	    raise (Key_collision (get_string_param p "URI"))
	  with Not_found ->
	    raise (Node_format_error
		     "Missing URI parameter of KeyCollision")
	  end
      | Small ("SizeError", _) -> self#close; raise Size_error
      | m -> self#handle_other m; assert false
    method wait_all =
      match (self#wait) with
	Some ruri -> ruri
      |	None -> self#wait_all
    initializer
      if (uri = "SVK@") || (uri = "freenet:SVK@") then
	(* This is NECESSARY, because although the FCP protocol allows
	 * SVK@ and freenet:SVK@ as valid SVKs, and says that a SVK
	 * key pair may be returned by Success as a result, this is
	 * redundant (the appropriate way to do it is with a node
	 * object's generate_svk method), and the method wait of
	 * insert_transact DOES NOT return SVK key pairs *)
	raise (Uri_error uri)
      else
	let prefix = Large_no_data
	    ("ClientPut",
	     [("HopsToLive", VInt htl);
	      ("URI", VString uri);
	      ("DataLength", VInt (metadata_len + data_len));
	      ("MetadataLength", VInt metadata_len)]) in
	self#send_msg prefix
  end

class gen_chk_transact connect_init metadata_len data_len = object
    (self)  
  inherit transact connect_init
  val mutable metadata_len_cur = 0
  val mutable data_len_cur = 0
  val mutable uri = None
  method metadata_block s =
    match uri with
      None ->
	let len = String.length s in
	if data_len_cur = 0 then
	  if metadata_len_cur + len <= metadata_len then begin
	    metadata_len_cur <- metadata_len_cur + len;
	      self#send s
	  end else
	    raise (Failure "More metadata given than specified")
	else
	  raise (Failure "Metadata must be specified before data")
    | Some chk ->
	raise (Failure "CHK already derived from metadata and data")
  method block s =
    match uri with
      None ->
	let len = String.length s in
	if metadata_len_cur = metadata_len then
	  if data_len_cur + len <= data_len then begin
	    data_len_cur <- data_len_cur + len;
	    self#send s
	  end else
	    raise (Failure "More data given than specified")
	else
	  raise (Failure "Metadata must be specified before data")
    | Some chk ->
	raise (Failure "CHK already derived from metadata and data")
  method private wait =
    match self#recv_msg with
      Small ("Success", p) ->
	self#close;
	begin try get_string_param p "URI"
	with Not_found ->
	  raise (Node_format_error "Missing URI parameter of Success")
	end
    | m -> self#handle_other m; assert false
  method generate_chk =
    match uri with
      None ->
	if (metadata_len_cur = metadata_len)
	    && (data_len_cur = data_len) then
	  begin
	    let chk = self#wait in
	    uri <- Some chk;
	    self#close;
	    chk
	  end
	else
	  raise (Failure "Not all metadata and data sent")
    | Some chk -> chk
    initializer
      let message = Large_no_data
	  ("GenerateCHK",
	   [("DataLength", VInt (metadata_len + data_len));
	    ("MetadataLength", VInt metadata_len)]) in
      self#send_msg message
end

class gen_svk_pair_transact connect_init = object (self)
  inherit transact connect_init
  val mutable prk = None
  val mutable pbk = None
  method private wait =
    match self#recv_msg with
      Small ("Success", p) ->
	self#close;
	begin try
	  let x = get_string_param p "PrivateKey"
	  and y = get_string_param p "PublicKey" in
	  prk <- Some x;
	  pbk <- Some y
	with Not_found ->
	  raise (Node_format_error
		   "Missing PrivateKey or PublicKey parameter of Success")
	end
    | m -> self#handle_other m; assert false
  method private_key =
    match prk with
      Some key -> key
    | None -> raise (Failure "SVK pair generation failed")
  method public_key =
    match pbk  with
      Some key -> key
    | None -> raise (Failure "SVK pair generation failed")
  initializer
    self#send_msg (Small ("GenerateSVKPair", []));
    self#wait
end

class node addr port = object (self)
  method private connect =
    try
      let s = Unix.socket Unix.PF_INET Unix.SOCK_STREAM 0 in
      let h = Unix.gethostbyname addr in
      Unix.connect s (Unix.ADDR_INET ((h.Unix.h_addr_list.(0)), port));
      s
    with Unix.Unix_error (_, _, _) -> raise Connect_error
  method request ~uri ~htl =
    new request_transact (self#connect) uri htl
  method insert ~uri ~htl ~metadata_len ~data_len =
    new insert_transact (self#connect) uri htl metadata_len data_len
  method generate_chk ~metadata_len ~data_len =
    new gen_chk_transact (self#connect) metadata_len data_len
  method generate_svk =
    let c = new gen_svk_pair_transact (self#connect) in
    c#private_key, c#public_key
end