open Unix;;
open Printf;;

(** Création et attachement d'une prise UDP *)
let open_udp_server port =
  let socket = socket PF_INET SOCK_DGRAM 0 in
  bind socket (ADDR_INET (inet_addr_any, port));
  socket;;

(** Attente de l'acquittement du paquet de numéro n *)
let rec wait_for_ack socket buffer n =
  match read socket buffer 0 4 with
    4 when (buffer.[0] = Char.chr 0 &&
            buffer.[1] = Char.chr 4 &&
            buffer.[2] = Char.chr (n / 256) &&
            buffer.[3] = Char.chr (n mod 256)) -> ()
  | 4 when (buffer.[0] = Char.chr 0 &&
            buffer.[1] = Char.chr 5) ->
             raise (Failure "Error datagram received")
  | _ ->
      prerr_endline "Bad packet received!\n";
      wait_for_ack socket buffer n;;

(** Émission d'un datagramme d'erreur *)
let send_error socket dest code msg =
  let packet = sprintf "\000\005\000%c%s\000" code msg in
  let packet_length = (String.length packet) in
  ignore (sendto socket packet 0 packet_length [] dest);;

(** Création d'une prise UDP et pseudo-connexion *)
let connect_udp_client dest =
  let socket = socket PF_INET SOCK_DGRAM 0 in
  setsockopt_float socket SO_RCVTIMEO 1.;
  connect socket dest;
  socket;;


let retry_nb = 5;;
(** Envoi de paquet avec attente d'acquittement *)
 let rec send_with_ack socket packet packet_size n ack_buffer =
   let rec retry i =
     if i = 0 then raise (Failure "Timeout, maximum retry reached")
     else
       try
         ignore (write socket packet 0 packet_size);
         wait_for_ack socket ack_buffer n
       with Unix_error(EAGAIN,"read",_) -> retry (i-1)
   in
   retry retry_nb;;

(** Lecture du maximum de données *)
let max_read fd buffer start length =
  let rec loop from nb  =
    if length <= 0 then 0 else
    match read fd buffer from nb with
      0 -> from - start
    | r -> loop (from + r) (nb - r) in
  loop start length;;

let unwind_protect f x g y = try let v = f x in g y; v with e -> g y; raise e;;

(** Envoi du fichier *)
let block_size = 512
let header_size = 4
let send_file dest name =
  let fd = openfile name [O_RDONLY] 0 in
  let socket = try connect_udp_client dest with e -> close fd; raise e in
  let packet = String.create (header_size + block_size) in
  packet.[0] <- Char.chr 0;
  packet.[1] <- Char.chr 3;
  let ack_buffer = String.create 4 in
  let rec send i =
    packet.[2] <- Char.chr (i/256);
    packet.[3] <- Char.chr (i mod 256);
    let len = max_read fd packet header_size block_size in
    send_with_ack socket packet (header_size + len) i ack_buffer;
    if len = block_size then send (succ i)
  in
  let close_everything () =
    close fd;
    close socket
  in
  unwind_protect send 1 close_everything ();;

(** Extraction du nom et du mode trouvé dans la requête *)
let extract_name_and_mode buffer len =
  try
    let name_offset = 2 in
    let name_len = String.index_from buffer 2 '\000' - name_offset in
    let name = String.sub buffer name_offset name_len in
    let mode_offset = name_offset + name_len + 1 in
    let mode = String.sub buffer mode_offset (len - mode_offset - 1) in
    name,mode
  with _ -> raise (Failure "extract_name_and_mode");;

(** Traitement d'une requête de lecture *)
let process_read_request socket dest buffer len =
  try
    match extract_name_and_mode buffer len with
      name, "octet" ->
        begin
          eprintf "Sending file %s\n" name;
          flush Pervasives.stderr;
          try
            send_file dest name
          with
            Unix_error (ENOENT,"open",_) ->
              prerr_endline "File not found";
              send_error socket dest '\001' "File not found"
          | Unix_error (EACCES,"open",_) ->
              prerr_endline "Access violation";
              send_error socket dest '\002' "Access violation"
          | Unix_error (ECONNREFUSED,_,_) ->
              prerr_endline "Peer socket closed"
          | e ->
              prerr_endline (Printexc.to_string e);
              send_error socket dest '\000' "Unexpected error"
        end
    | _ ->
        prerr_endline "Mode not supported";
        send_error socket dest '\004' "Mode not supported"
  with Failure "extract_name_and_mode" ->
    send_error socket dest '\000' "Bad request";;

(** Lance un serveur TFTP sur le port passé en argument *)
let establish_tftp_server port =
  let socket = open_udp_server port in
  let buffer_size = 1024 in
  let buffer = String.create buffer_size in
  while true do
    let len,from = recvfrom  socket buffer 0 buffer_size [] in
    match Char.code buffer.[0], Char.code buffer.[1] with
    | 0,1 -> (* Read *)
        process_read_request socket from buffer len
    | 0,2 -> (* Write *)
        prerr_endline "Operation not supported" ;
        send_error socket from '\004' "Operation not supported"
    | _,i -> (* Other *)
        eprintf "Packet with opcode %d received\n" i;
        flush Pervasives.stderr
  done;;

(** Récupère un numéro de port à partir d'une chaîne de caractères *)
let port_of_string port =
  try
    try (getservbyname port "tcp").s_port
    with Not_found -> int_of_string port
  with _ -> raise (Failure "port_of_string");;

(** fonction principale *)
let main () =
  if Array.length Sys.argv <> 2 then
    begin
      prerr_endline ("Usage: "^Sys.argv.(0)^" port");
      exit 1
    end
  else
    try
      let port = port_of_string Sys.argv.(1) in
      establish_tftp_server port
    with
      Failure "port_of_string" ->
        prerr_endline ("Unknown port: "^Sys.argv.(1));
        exit 3;;

handle_unix_error main ();;