irc-bot/client.ml

open Irc

(* ==========================================
 * Client stuff
 *)
type t = {outq: Command.t Queue.t;
          unsent: string ref;
          ibuf: string;
          ibuf_len: int ref;
          output_ready: unit -> unit;
          handle_command: (t -> Command.t -> unit) ref;
          nick: string ref;
          username: string ref;
          realname: string ref}

let modes = "l"

let dbg msg a = prerr_endline msg; a

let by_file_descr = Hashtbl.create 25
let by_nick = Hashtbl.create 25

let ibuf_max = 4096
let max_outq = 50
let obuf_max = 4096

let lookup nick =
  Hashtbl.find by_nick nick

let lookup_fd fd =
  Hashtbl.find by_file_descr fd

let close cli ues g fd =
  Hashtbl.remove by_nick !(cli.nick);
  Hashtbl.remove by_file_descr fd;
  Unix.close fd;
  Unixqueue.remove_resource ues g (Unixqueue.Wait_in fd);
  try
    Unixqueue.remove_resource ues g (Unixqueue.Wait_out fd);
  with Not_found ->
    ()

let write cli cmd =
  let was_empty = Queue.is_empty cli.outq in
    Queue.add cmd cli.outq;
    if (was_empty && (!(cli.unsent) = ""))  then
      cli.output_ready ()

let handle_close cli =
  ()

let handle_command cli cmd =
  write cli cmd

let handle_command_login cli cmd =
  (* Handle a command during the login phase *)
  begin
    match (Command.as_tuple cmd) with
      | (None, "USER", [username; _; _], Some realname) ->
          cli.username := username;
          cli.realname := Irc.truncate realname 40
      | (None, "NICK", [nick], None) ->
          cli.nick := nick;
      | _ ->
          write cli (Command.create 
                       (Some !Irc.name) "451" ["*"] (Some "Register first."))
  end;
  match (!(cli.username), !(cli.nick)) with
    | ("", _)
    | (_, "") ->
	()
    | (_, nick) ->
	cli.handle_command := handle_command;
	let command name text =
	  write cli (Command.create 
		       (Some !(Irc.name))
		       name
		       [nick]
		       (Some text))
	in
	  command "001" "Welcome to IRC.";
	  command "002" ("I am " ^ !(Irc.name) ^ 
			   " running version " ^ Irc.version);
	  command "003" "This server was created sometime";
	  command "004" (!(Irc.name) ^
			   " " ^ Irc.version ^
			   " " ^ modes ^
			   " " ^ Channel.modes)

let crlf = Str.regexp "\r?\n"

let handle_input cli =
  let buf = Str.string_before cli.ibuf !(cli.ibuf_len) in
  let lines = Str.split_delim crlf buf in
  let rec loop l =
    match l with
      | [] ->
	  ()
      | [leftover] ->
	  cli.ibuf_len := (String.length leftover);
	  String.blit leftover 0 cli.ibuf 0 !(cli.ibuf_len)
      | line :: tl ->
	  let parsed = Command.from_string line in
	    !(cli.handle_command) cli parsed;
	    loop tl
  in
    loop lines

let handle_event ues esys e =
  match e with
    | Unixqueue.Input_arrived (g, fd) ->
        let cli = lookup_fd fd in
        let size = ibuf_max - !(cli.ibuf_len) in
        let len = Unix.read fd cli.ibuf !(cli.ibuf_len) size in
          if (len > 0) then
            begin
              cli.ibuf_len := !(cli.ibuf_len) + len;
              try
                handle_input cli
              with Not_found ->
                if (!(cli.ibuf_len) = ibuf_max) then
                  (* No newline found, and the buffer is full *)
                  raise (Failure "Buffer overrun");
            end
          else
            close cli ues g fd
    | Unixqueue.Output_readiness (g, fd) ->
        (* XXX: Could be optimized to try and fill the output buffer *)
        let cli = lookup_fd fd in
        let buf =
          if (!(cli.unsent) = "") then
            let cmd = Queue.pop cli.outq in
              (Command.as_string cmd) ^ "\r\n"
          else
            !(cli.unsent)
        in
	let buflen = String.length buf in
        let n = Unix.single_write fd buf 0 buflen in
	  if n < buflen then
            cli.unsent := Str.string_after buf n
	  else if Queue.is_empty cli.outq then
	    Unixqueue.remove_resource ues g (Unixqueue.Wait_out fd)
    | Unixqueue.Out_of_band (g, fd) ->
        print_endline "oob"
    | Unixqueue.Timeout (g, op) ->
        print_endline "timeout"
    | Unixqueue.Signal ->
        print_endline "signal"
    | Unixqueue.Extra exn ->
        print_endline "extra"


let create ues g fd =
  let cli =
    {outq = Queue.create ();
     unsent = ref "";
     ibuf = String.create ibuf_max;
     ibuf_len = ref 0;
     output_ready = 
	begin
          fun () -> Unixqueue.add_resource ues g (Unixqueue.Wait_out fd, -.1.0)
	end;
     handle_command = ref handle_command_login;
     nick = ref "";
     username = ref "";
     realname = ref ""}
  in
    Hashtbl.replace by_file_descr fd cli;
    cli

let set_nick cli nick =
  Hashtbl.remove by_nick !(cli.nick);
  Hashtbl.replace by_nick nick cli;
  cli.nick := nick