Source file protocol.ml

(*
 * Copyright (C) Citrix Systems Inc.
 *
 * 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; version 2.1 only. with the special
 * exception on linking described in file LICENSE.
 *
 * 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.
 *)

(* NBD client library *)

open Sexplib.Std

(* We need to serialise/deserialise result values *)
type ('a, 'b) _result = [`Ok of 'a | `Error of 'b] [@@deriving sexp]

let result_of_sexp a b s =
  match _result_of_sexp a b s with `Ok x -> Ok x | `Error y -> Error y

let sexp_of_result a b r =
  sexp_of__result a b (match r with Ok x -> `Ok x | Error y -> `Error y)

let _nbd_cmd_read = 0l

let _nbd_cmd_write = 1l

let _nbd_cmd_disc = 2l

let _nbd_cmd_flush = 3l

let _nbd_cmd_trim = 4l

let nbd_request_magic = 0x25609513l

let nbd_reply_magic = 0x67446698l

let nbd_flag_has_flags = 1

let nbd_flag_read_only = 2

let nbd_flag_send_flush = 4

let nbd_flag_send_fua = 8

let nbd_flag_rotational = 16

let nbd_flag_send_trim = 32

let nbd_flag_fixed_newstyle = 1

let nbd_flag_no_zeroes = 2

let nbd_flag_c_fixed_newstyle = 1

let nbd_flag_c_no_zeroes = 2

let zero buf =
  for i = 0 to Cstruct.length buf - 1 do
    Cstruct.set_uint8 buf i 0
  done

module PerExportFlag = struct
  type t = Read_only | Send_flush | Send_fua | Rotational | Send_trim
  [@@deriving sexp]

  let to_string t = Sexplib.Sexp.to_string (sexp_of_t t)

  let of_int32 x =
    let flags = Int32.to_int x in
    let is_set i mask = i land mask = mask in
    List.map snd
      (List.filter
         (fun (mask, _) -> is_set flags mask)
         [
           (nbd_flag_read_only, Read_only)
         ; (nbd_flag_send_flush, Send_flush)
         ; (nbd_flag_send_fua, Send_fua)
         ; (nbd_flag_rotational, Rotational)
         ; (nbd_flag_send_trim, Send_trim)
         ]
      )

  let to_int flags =
    let one = function
      | Read_only ->
          nbd_flag_read_only
      | Send_flush ->
          nbd_flag_send_flush
      | Send_fua ->
          nbd_flag_send_fua
      | Rotational ->
          nbd_flag_rotational
      | Send_trim ->
          nbd_flag_send_trim
    in
    List.fold_left ( lor ) nbd_flag_has_flags (List.map one flags)

  let to_int32 flags = Int32.of_int (to_int flags)
end

module GlobalFlag = struct
  type t = Fixed_newstyle | No_zeroes [@@deriving sexp]

  let to_string t = Sexplib.Sexp.to_string (sexp_of_t t)

  let of_int flags =
    let is_set i mask = i land mask = mask in
    List.map snd
      (List.filter
         (fun (mask, _) -> is_set flags mask)
         [
           (nbd_flag_fixed_newstyle, Fixed_newstyle)
         ; (nbd_flag_no_zeroes, No_zeroes)
         ]
      )

  let to_int flags =
    let one = function
      | Fixed_newstyle ->
          nbd_flag_fixed_newstyle
      | No_zeroes ->
          nbd_flag_no_zeroes
    in
    List.fold_left ( lor ) 0 (List.map one flags)
end

module ClientFlag = struct
  type t = Fixed_newstyle | No_zeroes [@@deriving sexp]

  let to_string t = Sexplib.Sexp.to_string (sexp_of_t t)

  let of_int32 flags =
    let flags = Int32.to_int flags in
    let is_set mask = mask land flags <> 0 in
    List.map snd
      (List.filter
         (fun (mask, _) -> is_set mask)
         [
           (nbd_flag_c_fixed_newstyle, Fixed_newstyle)
         ; (nbd_flag_c_no_zeroes, No_zeroes)
         ]
      )

  let to_int32 flags =
    let one = function
      | Fixed_newstyle ->
          nbd_flag_c_fixed_newstyle
      | No_zeroes ->
          nbd_flag_c_no_zeroes
    in
    Int32.of_int (List.fold_left ( lor ) 0 (List.map one flags))
end

module Error = struct
  type t = [`EPERM | `EIO | `ENOMEM | `EINVAL | `ENOSPC | `Unknown of int32]
  [@@deriving sexp]

  let to_string t = Sexplib.Sexp.to_string (sexp_of_t t)

  let of_int32 = function
    | 1l ->
        `EPERM
    | 5l ->
        `EIO
    | 12l ->
        `ENOMEM
    | 22l ->
        `EINVAL
    | 28l ->
        `ENOSPC
    | x ->
        `Unknown x

  let to_int32 = function
    | `EPERM ->
        1l
    | `EIO ->
        5l
    | `ENOMEM ->
        12l
    | `EINVAL ->
        22l
    | `ENOSPC ->
        28l
    | `Unknown x ->
        x
end

module Command = struct
  type t = Read | Write | Disc | Flush | Trim | Unknown of int32
  [@@deriving sexp]

  let to_string t = Sexplib.Sexp.to_string (sexp_of_t t)

  let of_int32 = function
    | 0l ->
        Read
    | 1l ->
        Write
    | 2l ->
        Disc
    | 3l ->
        Flush
    | 4l ->
        Trim
    | c ->
        Unknown c

  let to_int32 = function
    | Read ->
        0l
    | Write ->
        1l
    | Disc ->
        2l
    | Flush ->
        3l
    | Trim ->
        4l
    | Unknown c ->
        c
end

module Option = struct
  type t = ExportName | Abort | List | StartTLS | Unknown of int32
  [@@deriving sexp]

  let to_string t = Sexplib.Sexp.to_string (sexp_of_t t)

  let of_int32 = function
    | 1l ->
        ExportName
    | 2l ->
        Abort
    | 3l ->
        List
    (* 4 is not in use in the NBD protocol. *)
    | 5l ->
        StartTLS
    (* 6, 7, 8 are not supported in this implementation. *)
    | c ->
        Unknown c

  let to_int32 = function
    | ExportName ->
        1l
    | Abort ->
        2l
    | List ->
        3l
    | StartTLS ->
        5l
    | Unknown c ->
        c
end

module OptionResponse = struct
  type t =
    | Ack
    | Server
    | Unsupported
    | Policy
    | Invalid
    | Platform
    | TlsReqd
    | Unknown of int32
  [@@deriving sexp]

  let to_string t = Sexplib.Sexp.to_string (sexp_of_t t)

  let of_int32 = function
    | 1l ->
        Ack
    | 2l ->
        Server
    | -2147483647l ->
        Unsupported
    | -2147483646l ->
        Policy
    | -2147483645l ->
        Invalid
    | -2147483644l ->
        Platform
    | -2147483643l ->
        TlsReqd
    | x ->
        Unknown x

  let to_int32 = function
    | Ack ->
        1l
    | Server ->
        2l
    | Unsupported ->
        -2147483647l
    | Policy ->
        -2147483646l
    | Invalid ->
        -2147483645l
    | Platform ->
        -2147483644l
    | TlsReqd ->
        -2147483643l
    | Unknown x ->
        x
end

(* Sent by the server to the client which includes an initial
   protocol choice *)
module Announcement = struct
  type t = [`V1 | `V2] [@@deriving sexp]

  type%cstruct t = {passwd: uint8_t [@len 8]; magic: uint64_t} [@@big_endian]

  let sizeof = sizeof_t

  let expected_passwd = "NBDMAGIC"

  let v1_magic = 0x00420281861253L

  let v2_magic = 0x49484156454F5054L (* Ascii encoding of "IHAVEOPT" *)

  let marshal buf t =
    set_t_passwd expected_passwd 0 buf ;
    set_t_magic buf (match t with `V1 -> v1_magic | `V2 -> v2_magic)

  let unmarshal buf =
    let passwd = Cstruct.to_string (get_t_passwd buf) in
    if passwd <> expected_passwd then
      Error (Failure "Bad magic in negotiate")
    else
      let magic = get_t_magic buf in
      if magic = v1_magic then
        Ok `V1
      else if magic = v2_magic then
        Ok `V2
      else
        Error
          (Failure
             (Printf.sprintf "Bad magic; expected %Ld or %Ld got %Ld" v1_magic
                v2_magic magic
             )
          )
end

module Negotiate = struct
  type v1 = {size: int64; flags: PerExportFlag.t list} [@@deriving sexp]

  type v2 = GlobalFlag.t list [@@deriving sexp]

  type t = V1 of v1 | V2 of v2 [@@deriving sexp]

  let to_string t = Sexplib.Sexp.to_string (sexp_of_t t)

  type%cstruct v1 = {
      size: uint64_t
    ; flags: uint32_t
    ; padding: uint8_t [@len 124]
  }
  [@@big_endian]

  type%cstruct v2 = {flags: uint16_t} [@@big_endian]

  let sizeof = function `V1 -> sizeof_v1 | `V2 -> sizeof_v2

  let marshal buf t =
    zero buf ;
    match t with
    | V1 t ->
        set_v1_size buf t.size ;
        set_v1_flags buf (PerExportFlag.to_int32 t.flags)
    | V2 t ->
        set_v2_flags buf (GlobalFlag.to_int t)

  let unmarshal buf t =
    match t with
    | `V1 ->
        let size = get_v1_size buf in
        let flags = PerExportFlag.of_int32 (get_v1_flags buf) in
        Ok (V1 {size; flags})
    | `V2 ->
        let flags = GlobalFlag.of_int (get_v2_flags buf) in
        Ok (V2 flags)
end

module NegotiateResponse = struct
  type t = ClientFlag.t list [@@deriving sexp]

  let sizeof = 4

  let marshal buf t = Cstruct.BE.set_uint32 buf 0 (ClientFlag.to_int32 t)

  let unmarshal buf = ClientFlag.of_int32 (Cstruct.BE.get_uint32 buf 0)
end

(* In the 'new' and 'new fixed' protocols, options are preceeded by
   a common header which includes a type and a length. *)
module OptionRequestHeader = struct
  type t = {ty: Option.t; length: int32} [@@deriving sexp]

  type%cstruct t = {magic: uint64_t; ty: uint32_t; length: uint32_t}
  [@@big_endian]

  let sizeof = sizeof_t

  let marshal buf t =
    set_t_magic buf Announcement.v2_magic ;
    set_t_ty buf (Option.to_int32 t.ty) ;
    set_t_length buf t.length

  let unmarshal buf =
    let open Rresult in
    let magic = get_t_magic buf in
    ( if Announcement.v2_magic <> magic then
        Error
          (Failure
             (Printf.sprintf "Bad reply magic: expected %Ld, got %Ld"
                Announcement.v2_magic magic
             )
          )
    else
      Ok ()
    )
    >>= fun () ->
    let ty = Option.of_int32 (get_t_ty buf) in
    let length = get_t_length buf in
    Ok {ty; length}
end

(* This is the option sent by the client to select a particular disk
   export. *)
module ExportName = struct
  type t = string [@@deriving sexp]

  let sizeof = String.length

  let marshal buf x = Cstruct.blit_from_string x 0 buf 0 (String.length x)
end

(* In both the 'new' style handshake and the 'fixed new' style handshake,
   the server will reply to an ExportName option with either a connection
   close or a DiskInfo: *)
module DiskInfo = struct
  type t = {size: int64; flags: PerExportFlag.t list} [@@deriving sexp]

  type%cstruct t = {
      size: uint64_t
    ; flags: uint16_t
    ; padding: uint8_t [@len 124]
  }
  [@@big_endian]

  let sizeof = sizeof_t

  let unmarshal buf =
    let size = get_t_size buf in
    let flags = PerExportFlag.of_int32 (Int32.of_int (get_t_flags buf)) in
    Ok {size; flags}

  let marshal buf t =
    set_t_size buf t.size ;
    set_t_flags buf (PerExportFlag.to_int t.flags)
end

(* In the 'fixed new' style handshake, all options apart from ExportName
   should result in reply packets as follows: *)
module OptionResponseHeader = struct
  type%cstruct t = {
      magic: uint64_t
    ; request_type: uint32_t
    ; response_type: uint32_t
    ; length: uint32_t
  }
  [@@big_endian]

  type t = {
      request_type: Option.t
    ; response_type: OptionResponse.t
    ; length: int32
  }
  [@@deriving sexp]

  let to_string t = Sexplib.Sexp.to_string (sexp_of_t t)

  let sizeof = sizeof_t

  let expected_magic = 0x3e889045565a9L

  let unmarshal buf =
    let open Rresult in
    let magic = get_t_magic buf in
    ( if expected_magic <> magic then
        Error
          (Failure
             (Printf.sprintf "Bad reply magic: expected %Ld, got %Ld"
                expected_magic magic
             )
          )
    else
      Ok ()
    )
    >>= fun () ->
    let request_type = Option.of_int32 (get_t_request_type buf) in
    let response_type = OptionResponse.of_int32 (get_t_response_type buf) in
    let length = get_t_length buf in
    Ok {request_type; response_type; length}

  let marshal buf t =
    set_t_magic buf expected_magic ;
    set_t_request_type buf (Option.to_int32 t.request_type) ;
    set_t_response_type buf (OptionResponse.to_int32 t.response_type) ;
    set_t_length buf t.length
end

(* A description of an export, sent in response to a List option *)
module Server = struct
  type t = {name: string} [@@deriving sexp]

  type%cstruct t = {length: uint32_t} [@@big_endian]

  let sizeof t = sizeof_t + String.length t.name

  let unmarshal buf =
    let length = Int32.to_int (get_t_length buf) in
    let buf = Cstruct.shift buf sizeof_t in
    let name = Cstruct.to_string (Cstruct.sub buf 0 length) in
    Ok {name}
end

module Request = struct
  type t = {ty: Command.t; handle: int64; from: int64; len: int32}
  [@@deriving sexp]

  let to_string t =
    Printf.sprintf "{ Command = %s; handle = %Ld; from = %Ld; len = %ld }"
      (Command.to_string t.ty) t.handle t.from t.len

  type%cstruct t = {
      magic: uint32_t
    ; ty: uint32_t
    ; handle: uint64_t
    ; from: uint64_t
    ; len: uint32_t
  }
  [@@big_endian]

  let unmarshal (buf : Cstruct.t) =
    let open Rresult in
    let magic = get_t_magic buf in
    ( if nbd_request_magic <> magic then
        Error
          (Failure
             (Printf.sprintf "Bad request magic: expected %ld, got %ld" magic
                nbd_request_magic
             )
          )
    else
      Ok ()
    )
    >>= fun () ->
    let ty = Command.of_int32 (get_t_ty buf) in
    let handle = get_t_handle buf in
    let from = get_t_from buf in
    let len = get_t_len buf in
    Ok {ty; handle; from; len}

  let sizeof = sizeof_t

  let marshal (buf : Cstruct.t) t =
    set_t_magic buf nbd_request_magic ;
    set_t_ty buf (Command.to_int32 t.ty) ;
    set_t_handle buf t.handle ;
    set_t_from buf t.from ;
    set_t_len buf t.len
end

module Reply = struct
  type t = {error: (unit, Error.t) result; handle: int64} [@@deriving sexp]

  let to_string t = Sexplib.Sexp.to_string (sexp_of_t t)

  type%cstruct t = {magic: uint32_t; error: uint32_t; handle: uint64_t}
  [@@big_endian]

  let unmarshal (buf : Cstruct.t) =
    let open Rresult in
    let magic = get_t_magic buf in
    ( if nbd_reply_magic <> magic then
        Error
          (Failure
             (Printf.sprintf "Bad reply magic: expected %ld, got %ld" magic
                nbd_reply_magic
             )
          )
    else
      Ok ()
    )
    >>= fun () ->
    let error = get_t_error buf in
    let error = if error = 0l then Ok () else Error (Error.of_int32 error) in
    let handle = get_t_handle buf in
    Ok {error; handle}

  let sizeof = sizeof_t

  let marshal (buf : Cstruct.t) t =
    set_t_magic buf nbd_reply_magic ;
    let error =
      match t.error with Ok () -> 0l | Error e -> Error.to_int32 e
    in
    set_t_error buf error ; set_t_handle buf t.handle
end