extra/protocols/tftp/tftp.factor

   1 ! Copyright (C) 2019 Doug Coleman.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors arrays assocs combinators
   4 combinators.short-circuit continuations destructors endian
   5 io io.directories io.encodings.binary io.encodings.latin1
   6 io.encodings.string io.encodings.utf8 io.files io.files.info
   7 io.sockets kernel math math.parser namespaces pack prettyprint
   8 random sequences sequences.extras splitting strings ;
   9 IN: protocols.tftp
  10
  11 CONSTANT: TFTP-RRQ 1 ! Read request (RRQ)
  12 CONSTANT: TFTP-WRQ 2 ! Write request (WRQ)
  13 CONSTANT: TFTP-DATA 3 ! Data (DATA)
  14 CONSTANT: TFTP-ACK 4 ! Acknowledgment (ACK)
  15 CONSTANT: TFTP-ERROR 5 ! Error (ERROR)
  16
  17 GENERIC: get-tftp-host ( server -- host )
  18 M: string get-tftp-host resolve-host random host>> 69 <inet4> ;
  19 M: integer get-tftp-host "127.0.0.1" swap <inet4> ;
  20 M: inet4 get-tftp-host ;
  21 M: f get-tftp-host drop "127.0.0.1" 69 <inet4> ;
  22
  23 : tftp-get ( filename encoding server -- bytes )
  24     '[
  25         TFTP-RRQ _ _ 3array "Saa" pack-be
  26         _ get-tftp-host
  27         f 0 <inet4> <datagram> &dispose
  28         [ send ] keep
  29         dup
  30         '[
  31             _ receive
  32             [ 4 cut swap 2 cut nip be> TFTP-ACK swap 2array "SS" pack-be ] dip
  33             _ send
  34             dup length 511 >
  35         ] loop>array* concat
  36     ] with-destructors ;
  37
  38 : tftp-get-netascii ( filename server/port/inet4/f -- bytes )
  39     "netascii" swap tftp-get latin1 decode ;
  40
  41 : tftp-get-octet ( filename server/port/inet4/f -- bytes )
  42     "octet" swap tftp-get ;
  43
  44 SYMBOL: tftp-server
  45 SYMBOL: tftp-client
  46 SYMBOL: clients
  47 SYMBOL: tftp-servers
  48 tftp-servers [ H{ } clone ] initialize
  49 TUPLE: read-file path encoding block ;
  50
  51 : send-client ( bytes -- )
  52     tftp-client get tftp-server get send ;
  53
  54 : send-error ( message -- )
  55     [ TFTP-ERROR 1 ] dip 3array "SSa" pack-be send-client ;
  56
  57 : send-file-block ( bytes block -- )
  58     TFTP-DATA swap 2array "SS" pack-be B{ } prepend-as
  59     send-client ;
  60
  61 : read-file-block ( path n -- bytes )
  62     binary swap
  63     '[ _ 512 * seek-absolute seek-input 512 read ] with-file-reader ;
  64
  65 : handle-send-file-next ( block -- )
  66     drop
  67     tftp-client get clients get ?at [
  68         [ [ path>> ] [ block>> ] bi read-file-block ]
  69         [ [ 1 + ] change-block block>> ] bi
  70         send-file-block
  71     ] [
  72         drop
  73     ] if ;
  74
  75 : handle-send-file ( bytes -- )
  76     "\0" split harvest first2 [ utf8 decode ] bi@
  77     over { [ file-exists? ] [ file-info directory? not ] } 1&& [
  78         "netascii" sequence= utf8 binary ? 0 read-file boa
  79         tftp-client get clients get set-at
  80         0 handle-send-file-next
  81     ] [
  82         2drop "File not found" send-error
  83     ] if ;
  84
  85 : read-tftp-command ( -- )
  86     tftp-server get receive tftp-client [
  87         2 cut swap be> {
  88             { TFTP-RRQ [ handle-send-file ] }
  89             { TFTP-ACK [ be> handle-send-file-next ] }
  90             [ number>string " unimplemented" append throw ]
  91         } case
  92     ] with-variable ;
  93
  94 : start-tftp-server ( directory port/f -- )
  95     get-tftp-host
  96     '[
  97         H{ } clone clients [
  98             _ <datagram> tftp-server [
  99                 tftp-server get dup addr>> port>> tftp-servers get-global set-at
 100                 [
 101                     [ read-tftp-command t ]
 102                     [ [ . flush ] with-global f ] recover
 103                 ] loop
 104             ] with-variable
 105         ] with-variable
 106     ] with-directory ;
 107
 108 ERROR: tftp-server-not-running port ;
 109 : stop-tftp-server ( port -- )
 110     tftp-servers get-global ?delete-at [
 111         dispose
 112     ] [
 113         tftp-server-not-running
 114     ] if ;