vcap.ml

   1 open Format;;
   2
   3 type c =
   4     { w : int
   5     ; h : int
   6     ; pitch : int
   7     ; audio_bytes_per_frame : int
   8     ; audio_bytes_processed : int
   9     ; silence : string
  10     ; audio_running : bool
  11     ; depth : int
  12     ; scaledsize : int
  13     ; video : string
  14     ; lastout : string
  15     ; update_requested : bool
  16     ; last_video_update : float
  17     }
  18
  19 type encset =
  20     { out_w : int
  21     ; out_h : int
  22     ; fps : int
  23     ; freq : int
  24     ; nchannels : int
  25     ; audiofmt : int
  26     ; useaudio : bool
  27     }
  28
  29 module Scale = struct
  30   type t = (int * int * int * int * int)
  31   external reset : t -> int = "ml_scale_reset"
  32   external scale : string -> string -> unit = "ml_scale"
  33
  34   let reset ~w ~h ~out_w ~out_h ~pitch =
  35     reset (w, h, out_w, out_h, pitch)
  36 end
  37
  38 let fail fmt = Printf.kprintf (fun s -> failwith s) fmt;;
  39 let log = printf;;
  40 let log fmt = Printf.kprintf ignore fmt;;
  41
  42 let r8 s o = Char.code s.[o];;
  43 let r16 s o = let b1 = r8 s o and b2 = r8 s (o+1) in (b1 lsl 8) lor b2;;
  44 let r16s s o = let i = r16 s 0 in i - ((i lor 0x8000) lsl 1);;
  45 let r32 s o =
  46   let w1 = r16 s o
  47   and w2 = r16 s (o+2) in
  48   Int32.logor (Int32.shift_left (Int32.of_int w1) 16) (Int32.of_int w2)
  49 ;;
  50 let rint s o = let l = r32 s o in Int32.to_int l;;
  51 let w8 s o i = s.[o] <- Char.chr (i land 255);;
  52 let w16 s o i = w8 s (o+1) i; w8 s o (i lsr 8);;
  53 let wint s o i = w16 s (o+2) i; w16 s o (i asr 16);;
  54 let w32 s o l =
  55   w16 s (o+2) (Int32.to_int l);
  56   w16 s o (Int32.to_int (Int32.shift_right l 16));
  57 ;;
  58
  59 let encodings t =
  60   let b = Buffer.create 20 in
  61   Buffer.add_string b "\002xyy";
  62   let c = "xxxx" in
  63   let nencodings = List.fold_left
  64     (fun n i -> w32 c 0 i; Buffer.add_string b c; succ n)
  65     0 (0l :: -223l :: (if t.useaudio then [-259l] else []))
  66   in
  67   let s = Buffer.contents b in
  68   w16 s 2 nencodings;
  69   s
  70 ;;
  71
  72 let setpixfmt =
  73   let s = String.create 20 in
  74   w8 s 0 0;
  75   w8 s 4 16;
  76   w8 s 5 2;
  77   w8 s 6 1;
  78   w8 s 7 1;
  79   w16 s 8 31;
  80   w16 s 10 63;
  81   w16 s 12 31;
  82   w8 s 14 11;
  83   w8 s 15 5;
  84   w8 s 16 0;
  85   s
  86 ;;
  87
  88 let setaudiofmt ~fmt ~nchannels ~freq =
  89   let s = String.create 10 in
  90   w8  s 0 255;
  91   w8  s 1 001;
  92   w16 s 2 002;
  93   w8  s 4 fmt;
  94   w8  s 5 nchannels;
  95   w32 s 6 freq;
  96   s
  97 ;;
  98
  99 let useaudio =
 100   let s = String.create 4 in
 101   w8  s 0 255;
 102   w8  s 1 001;
 103   w16 s 2 000;
 104   s
 105 ;;
 106
 107 let resize t w h out_w out_h =
 108   let pitch = w * t.depth in
 109   let scaledsize = Scale.reset ~w ~h ~out_w ~out_h ~pitch in
 110   let video = String.create (pitch * h) in
 111   { t with video = video
 112     ; w = w
 113     ; h = h
 114     ; pitch = pitch
 115     ; scaledsize = scaledsize
 116     ; lastout = ""
 117   }
 118 ;;
 119
 120 let create_context ~w ~h ~encset =
 121   let bytes_per_frame =
 122     let sample_size =
 123       match encset.audiofmt with
 124       | 0 | 1 -> 1
 125       | 2 | 3 -> 2
 126       | 4 | 5 -> 4
 127       | _ -> fail "invalid audiofmt %d" encset.audiofmt
 128     in
 129     let frame_size = sample_size * encset.nchannels in
 130     (encset.freq * frame_size) / encset.fps
 131   in
 132   resize
 133     { video = ""; w = -1; h = -1; pitch = -1; depth = 2
 134     ; audio_bytes_per_frame = bytes_per_frame
 135     ; audio_bytes_processed = 0
 136     ; scaledsize = -1
 137     ; update_requested = false
 138     ; audio_running = false
 139     ; lastout = ""
 140     ; last_video_update = 0.0
 141     ; silence =
 142         String.make bytes_per_frame '\000' (* wrong for unsigned *)
 143     }
 144     w h encset.out_w encset.out_h
 145 ;;
 146
 147 let reason pump =
 148   let n = rint (pump#get 4) 0 in
 149   let s = pump#get n in
 150   fail "Protocol failure, reason: `%S'" s;
 151 ;;
 152
 153 let rec handshake encset pump =
 154   let s = pump#get 12 in
 155   let maj, min =
 156     try
 157       let dec a b c = a*100 + b*10 + c in
 158       Scanf.sscanf s "RFB %1i%1i%1i.%1i%1i%1i\n"
 159         (fun a b c d e f -> dec a b c, dec d e f)
 160     with exn ->
 161       fail "can't parse ProtocolVersion `%S'" s
 162   in
 163
 164   if maj < 3 || (maj == 3 && min < 8)
 165   then fail "Too old of a protocol %d.%d" maj min;
 166
 167   pump#put"RFB 003.008\n";
 168   let s = let nsectypes = r8 (pump#get 1) 0 in pump#get nsectypes in
 169   (try ignore (String.index s '\001')
 170     with Not_found -> fail "None auth not found");
 171   pump#put "\001";
 172   if r32 (pump#get 4) 0 <> 0l then reason pump;
 173   pump#put "\001";
 174   serverinit encset pump
 175
 176 and serverinit encset pump =
 177   let s = pump#get 24 in
 178   let w = r16 s 0
 179   and h = r16 s 2
 180   and namelen = rint s 20 in
 181   let name = pump#get namelen in
 182   let () =
 183     printf "server is `%S'@." name;
 184     pump#put (encodings encset);
 185     pump#put setpixfmt;
 186   in
 187   let context = create_context ~w ~h ~encset in
 188   boot encset context pump
 189
 190 and boot encset t pump =
 191   let request_update t =
 192     if not t.update_requested then (
 193       let s = "\003\001\000\000\000\000wwhh" in
 194       w16 s 6 t.w;
 195       w16 s 8 t.h;
 196       pump#put s;
 197     );
 198     { t with update_requested = true }
 199   in
 200
 201   let rec srv1 t =
 202     pump#run;
 203     match r8 (pump#get 1) 0 with
 204     | 000 -> framebuffer_update t
 205     | 255 -> aliguori t
 206     | msg -> fail "unexpected server message %x" msg
 207
 208   and drain t =
 209     let t =
 210       if t.audio_running
 211       then
 212         if t.audio_bytes_processed > t.audio_bytes_per_frame
 213         then request_update t else t
 214       else
 215         let fdur = 1.0 /. float encset.fps in
 216         let deadline = t.last_video_update +. fdur in
 217         let now = Unix.gettimeofday () in
 218         let rec loop t delta =
 219           pump#run;
 220           let got_input = if delta > 0.0 then pump#wait delta else false in
 221           if got_input
 222           then t
 223           else (
 224             let t =
 225               if t.update_requested
 226               then (pump#add_audio t.silence; pump#add_video t.lastout; t)
 227               else request_update t
 228             in
 229             loop t fdur
 230           )
 231         in
 232         loop t (deadline -. now)
 233     in
 234     srv1 t
 235
 236   and audio_ack t =
 237     let fmt = encset.audiofmt
 238     and nchannels = encset.nchannels
 239     and freq = Int32.of_int encset.freq in
 240     pump#put (setaudiofmt ~fmt ~nchannels ~freq);
 241     pump#put useaudio;
 242     t
 243
 244   and aliguori t =
 245     match r8 (pump#get 1) 0 with
 246     | 001 -> audio t
 247     | msg -> fail "unexpected aliguori message %x" msg
 248
 249   and framebuffer_update t =
 250     let nrects = r16 (pump#get 3) 1 in
 251     let rec loop t i = if i = nrects then t else
 252         let s = pump#get 12 in
 253         let x = r16 s 0
 254         and y = r16 s 2
 255         and w = r16 s 4
 256         and h = r16 s 6
 257         and e = r32 s 8 in
 258         match e with
 259         | 0l -> raw_rect t x y w h; loop t (i+1)
 260         | -223l -> loop (resize t w h encset.out_w encset.out_h) (i+1)
 261         | -259l -> loop (audio_ack t) (i+1)
 262         | _  -> fail "rect with unhandled encoding %lx, %S" e s
 263     in
 264     let t = loop t 0 in
 265     let out = String.create t.scaledsize in
 266     Scale.scale t.video out;
 267     pump#add_video out;
 268     if not t.audio_running then pump#add_audio t.silence;
 269     drain { t with
 270       last_video_update = Unix.gettimeofday ();
 271       lastout = out;
 272       update_requested = false;
 273       audio_bytes_processed = t.audio_bytes_processed - t.audio_bytes_per_frame }
 274
 275   and raw_rect t x y w h =
 276     let src = pump#get (w*h*t.depth) in
 277     let dst = t.video in
 278     let len = w*t.depth in
 279     let pitch = t.pitch in
 280     let rec loop row src_pos dst_pos =
 281       if row = h then () else (
 282         StringLabels.blit ~src ~dst ~src_pos ~dst_pos ~len;
 283         loop (row+1) (src_pos+len) (dst_pos+pitch)
 284       )
 285     in
 286     loop 0 0 (x*t.depth + y*t.pitch)
 287
 288   and audio t =
 289     match r16 (pump#get 2) 0 with
 290     | 000 -> drain { t with audio_running = false }
 291     | 001 -> drain { t with audio_running = true }
 292     | 002 -> audio_data t
 293     | msg -> fail "unexpected audio message %x" msg
 294
 295   and audio_data t =
 296     let nbytes = rint (pump#get 4) 0 in
 297     let audio = pump#get nbytes in
 298     pump#add_audio audio;
 299     let t =
 300       { t with audio_bytes_processed = t.audio_bytes_processed + nbytes }
 301     in
 302     drain t
 303   in
 304
 305   srv1 (request_update t);
 306 ;;
 307
 308 let vncopen addr =
 309   let sock, addr =
 310     if String.length addr > 5 && String.sub addr 0 5 = "unix:"
 311     then
 312       let path = String.sub addr 5 (String.length addr - 5) in
 313       let sock = Unix.socket Unix.PF_UNIX Unix.SOCK_STREAM 0 in
 314       sock, Unix.ADDR_UNIX path
 315     else
 316       let sock = Unix.socket Unix.PF_INET Unix.SOCK_STREAM 0 in
 317       let addr, port =
 318         try
 319           let p = String.index addr ':' in
 320           let s = String.sub addr (p + 1) (String.length addr - p - 1) in
 321           let port =
 322             try int_of_string s with exn ->
 323               fail "can't parse port in `%S': %s"
 324                 addr (Printexc.to_string exn)
 325           in
 326           let addr = String.sub addr 0 p in
 327           addr, port
 328         with Not_found ->
 329           addr, 5900
 330       in
 331       let addr = (Unix.gethostbyname addr).Unix.h_addr_list.(0) in
 332       sock, Unix.ADDR_INET (addr, port)
 333   in
 334   Unix.connect sock addr;
 335   sock
 336 ;;
 337
 338 let pump encset sock =
 339   let null = Unix.openfile "/dev/null" [] 0 in
 340   let rapipe, wapipe = Unix.pipe () in
 341   let rvpipe, wvpipe = Unix.pipe () in
 342   let () =
 343     Unix.set_close_on_exec wapipe;
 344     Unix.set_close_on_exec wvpipe;
 345     Unix.set_close_on_exec sock;
 346     Unix.set_nonblock wapipe;
 347     Unix.set_nonblock wvpipe;
 348   in
 349   let _pid = Unix.create_process
 350     "/bin/sh"
 351     [| "/bin/sh"; "./qcap.sh"
 352     ; string_of_int encset.out_w
 353     ; string_of_int encset.out_h
 354     ; string_of_int (Obj.magic rvpipe)
 355     ; string_of_int (Obj.magic rapipe)
 356     ; string_of_int encset.fps
 357     |]
 358     null
 359     null
 360     Unix.stderr
 361   in
 362   let rw c f s =
 363     let rec loop p l =
 364       let n = f s p l in
 365       if n = 0 then fail "%s: end of file" c;
 366       if n = l then () else loop (p + n) (l - n)
 367     in loop 0 (String.length s)
 368   in
 369   let rec write s fd buf pos len () =
 370     let n =
 371       try
 372         let n = Unix.write fd buf pos len in
 373         if n = 0 then fail "EOF while writing %s" s;
 374         n
 375       with Unix.Unix_error (Unix.EAGAIN, _, _) -> 0
 376     in
 377     if n = 0 then `again else
 378       if n = len
 379       then `completed
 380       else `more (write s fd buf (pos+n) (len-n))
 381   in object (self)
 382     val al = ref []
 383     val vl = ref []
 384
 385     method run =
 386       let l =
 387         match !al, !vl with
 388         | [], [] -> []
 389         | [], _ -> [wvpipe]
 390         | _, [] -> [wapipe]
 391         | _ -> [wvpipe; wapipe]
 392       in
 393       if l = [] then () else
 394         let _, l, _ = Unix.select [] l [] 0.0 in
 395         let rec loop = function
 396           | fd :: tl ->
 397               let lr = if fd = wapipe then al else vl in
 398               let rec wrall accu = function
 399                 | [] -> accu
 400                 | f :: tl ->
 401                     match f () with
 402                     | `completed -> wrall accu tl
 403                     | `again -> f :: accu
 404                     | `more f -> f :: accu
 405               in
 406               lr := List.rev (wrall [] !lr);
 407               loop tl
 408           | [] -> ()
 409         in
 410         loop l
 411
 412     method add_video data =
 413       if String.length data = 0 then fail "empty data passed to add_video";
 414       vl := !vl @ [write "video" wvpipe data 0 (String.length data)]
 415
 416     method add_audio data =
 417       al := !al @ [write "audio" wapipe data 0 (String.length data)]
 418
 419     method wait dur =
 420       let r, _, _ = Unix.select [sock] [] [] dur in
 421       r != []
 422
 423     method get n = let s = String.create n in rw "get" (Unix.read sock) s; s
 424     method put s = rw "put" (Unix.write sock) s
 425   end
 426 ;;
 427
 428 let main =
 429   let out_w = ref 640
 430   and out_h = ref 400
 431   and fps = ref 25
 432   and noaudio = ref false
 433   and host = ref "localhost:5900" in
 434   let setdim s = Scanf.sscanf s "%dx%d" (fun w h -> out_w := w; out_h := h) in
 435   Arg.parse (Arg.align
 436                 ["-s", Arg.String setdim , "<int>x<int> output dimensions"
 437                 ;"-fps", Arg.Set_int fps, "<int> desired fps"
 438                 ;"-noaudio", Arg.Set noaudio, "<bool> do not capture audio"])
 439     (fun s -> host := s)
 440     "Usage: vcap [options] [host]\nOptions are:";
 441   let encset =
 442     { out_w = !out_w
 443     ; out_h = !out_h
 444     ; nchannels = 2
 445     ; audiofmt = 3
 446     ; freq = 44100
 447     ; fps = !fps
 448     ; useaudio = not !noaudio
 449     }
 450   in
 451   try
 452     let sock = vncopen !host in
 453     let pump = pump encset sock in
 454     handshake encset pump;
 455   with Unix.Unix_error (e, s1, s2) ->
 456     eprintf "%s(%s): %s@." s1 s2 (Unix.error_message e)
 457 ;;