extra/bittorrent/bittorrent.factor

   1 ! Copyright (C) 2020 John Benediktsson
   2 ! See http://factorcode.org/license.txt for BSD license
   3 USING: accessors arrays assocs bencode byte-arrays checksums
   4 checksums.sha combinators fry grouping http.client io io.binary
   5 io.encodings.binary io.files io.pathnames io.sockets
   6 io.streams.byte-array kernel literals make math math.bitwise
   7 math.parser math.ranges namespaces random sequences splitting
   8 strings urls ;
   9
  10 IN: bittorrent
  11
  12 <<
  13 CONSTANT: ALPHANUMERIC $[
  14     [
  15         CHAR: a CHAR: z [a,b] %
  16         CHAR: A CHAR: Z [a,b] %
  17         CHAR: 0 CHAR: 9 [a,b] %
  18         ".-_~" %
  19     ] { } make
  20 ]
  21
  22 : random-peer-id ( -- bytes )
  23     20 [ ALPHANUMERIC random ] B{ } replicate-as ;
  24 >>
  25
  26 SYMBOL: torrent-peer-id
  27 torrent-peer-id [ random-peer-id ] initialize
  28
  29 SYMBOL: torrent-port
  30 torrent-port [ 6881 ] initialize
  31
  32
  33 ! bitfield
  34
  35 : bitfield-index ( n -- j i )
  36     8 /mod 7 swap - ;
  37
  38 : set-bitfield ( elt n bitfield -- )
  39     [ bitfield-index rot ] dip -rot
  40     '[ _ _ [ set-bit ] [ clear-bit ] if ] change-nth ;
  41
  42 : check-bitfield ( n bitfield -- ? )
  43     [ bitfield-index swap ] dip nth bit? ;
  44
  45
  46 ! http
  47
  48 : http-get-bencode ( url -- obj )
  49     <get-request> BV{ } clone [
  50         '[ _ push-all ] with-http-request* check-response drop
  51     ] keep B{ } like bencode> ;
  52
  53
  54 ! metainfo
  55
  56 GENERIC: load-metainfo ( obj -- metainfo )
  57
  58 M: url load-metainfo http-get-bencode ;
  59
  60 M: pathname load-metainfo
  61     binary [ read-bencode ] with-file-reader ;
  62
  63 M: string load-metainfo
  64     dup "http" head? [ >url ] [ <pathname> ] if load-metainfo ;
  65
  66 : info-hash ( metainfo -- hash )
  67     "info hash" swap dup '[
  68         drop _ "info" of >bencode sha1 checksum-bytes
  69     ] cache ;
  70
  71 : announce-url ( metainfo -- url )
  72     dup "announce-list" of [ nip first random ] [ "announce" of ] if* ;
  73
  74 : scrape-url ( metainfo -- url/f )
  75     announce-url "announce" over path>> subseq? [
  76         [ "announce" "scrape" replace ] change-path
  77     ] [ drop f ] if ;
  78
  79
  80
  81 ! tracker
  82
  83 : tracker-url ( metainfo -- url )
  84     {
  85         [ announce-url >url ]
  86         [
  87             info-hash "info_hash" set-query-param
  88             torrent-peer-id get "peer_id" set-query-param
  89             torrent-port get "port" set-query-param
  90             0 "uploaded" set-query-param
  91             0 "downloaded" set-query-param
  92             1 "compact" set-query-param
  93         ]
  94         [
  95             { "info" "length" } [ of ] each
  96             "left" set-query-param
  97         ]
  98     } cleave ;
  99
 100 : parse-peer4 ( peerbin -- inet4 )
 101     4 cut [
 102         [ number>string ] { } map-as "." join
 103     ] dip be> <inet4> ;
 104
 105 : parse-peer4s ( peersbin -- inet4s )
 106     dup array? [
 107         [ [ "ip" of ] [ "port" of ] bi <inet4> ] map
 108     ] [
 109         6 <groups> [ parse-peer4 ] map
 110     ] if ;
 111
 112 : parse-peer6 ( peerbin -- inet6 )
 113     16 cut [
 114         2 <groups> [ be> number>string ] map ":" join
 115     ] dip be> <inet6> ;
 116
 117 : parse-peer6s ( peersbin -- inet6s )
 118     18 <groups> [ parse-peer6 ] map ;
 119
 120 : load-tracker ( torrent -- response )
 121     tracker-url http-get-bencode
 122     "peers" over [ parse-peer4s ] change-at ;
 123
 124 : send-event ( torrent event -- response )
 125     [ tracker-url ] [ "event" set-query-param ] bi*
 126     http-get-bencode ;
 127
 128
 129
 130 ! messages
 131
 132 TUPLE: handshake string reserved info-hash peer-id ;
 133
 134 : <handshake> ( info-hash peer-id -- handshake )
 135     handshake new
 136         "BitTorrent protocol" >byte-array >>string
 137         8 <byte-array> >>reserved
 138         swap >>peer-id
 139         swap >>info-hash ;
 140
 141 : read-handshake ( -- handshake/f )
 142     read1 [
 143         [ 48 + read ] keep cut 8 cut 20 cut handshake boa
 144     ] [ f ] if* ;
 145
 146 : write-handshake ( handshake -- )
 147     {
 148         [ string>> [ length write1 ] [ write ] bi ]
 149         [ reserved>> write ]
 150         [ info-hash>> write ]
 151         [ peer-id>> write ]
 152     } cleave flush ;
 153
 154 TUPLE: keep-alive ;
 155 TUPLE: choke ;
 156 TUPLE: unchoke ;
 157 TUPLE: interested ;
 158 TUPLE: not-interested ;
 159 TUPLE: have index ;
 160 TUPLE: bitfield bitfield ;
 161 TUPLE: request index begin length ;
 162 TUPLE: piece index begin block ;
 163 TUPLE: cancel index begin length ;
 164 TUPLE: port port ;
 165 TUPLE: suggest-piece index ;
 166 TUPLE: have-all ;
 167 TUPLE: have-none ;
 168 TUPLE: reject-request index begin length ;
 169 TUPLE: allowed-fast index ;
 170 TUPLE: extended id payload ;
 171 TUPLE: unknown id payload ;
 172
 173 : read-int ( -- n/f ) 4 read [ be> ] [ f ] if* ;
 174
 175 : parse-message ( bytes -- message/f )
 176     unclip {
 177         ! Core Protocol
 178         { 0 [ drop choke boa ] }
 179         { 1 [ drop unchoke boa ] }
 180         { 2 [ drop interested boa ] }
 181         { 3 [ drop not-interested boa ] }
 182         { 4 [ 4 head be> have boa ] }
 183         { 5 [ bitfield boa ] }
 184         { 6 [ 4 cut 4 cut 4 head [ be> ] tri@ request boa ] }
 185         { 7 [ 4 cut 4 cut [ [ be> ] bi@ ] dip piece boa ] }
 186         { 8 [ 4 cut 4 cut 4 head [ be> ] tri@ cancel boa ] }
 187
 188         ! DHT Extension
 189         { 9 [ be> port boa ] }
 190
 191         ! Fast Extensions
 192         { 0x0D [ 4 head be> suggest-piece boa ] }
 193         { 0x0E [ drop have-all boa ] }
 194         { 0x0F [ drop have-none boa ] }
 195         { 0x10 [ 4 cut 4 cut 4 head [ be> ] tri@ reject-request boa ] }
 196         { 0x11 [ 4 head be> allowed-fast boa ] }
 197
 198         ! Extension Protocol
 199         { 0x14 [ unclip swap extended boa ] }
 200
 201         ! Hash Transfer Protocol
 202         ! { 0x15 [ "HashRequest" ] }
 203         ! { 0x16 [ "Hashes" ] }
 204         ! { 0x17 [ "HashReject" ] }
 205         [ swap unknown boa ]
 206     } case ;
 207
 208 : read-message ( -- message )
 209     read-int {
 210         { f [ f ] }
 211         { 0 [ keep-alive boa ] }
 212         [ read [ parse-message ] [ f ] if* ]
 213     } case ;
 214
 215 : write-int ( n -- ) 4 >be write ;
 216
 217 GENERIC: write-message ( message -- )
 218
 219 M: keep-alive write-message drop 0 write-int ;
 220
 221 M: choke write-message drop 1 write-int 0 write1 ;
 222
 223 M: unchoke write-message drop 1 write-int 1 write1 ;
 224
 225 M: interested write-message drop 1 write-int 2 write1 ;
 226
 227 M: not-interested write-message drop 1 write-int 3 write1 ;
 228
 229 M: have write-message
 230     5 write-int 4 write1 index>> write-int ;
 231
 232 M: bitfield write-message
 233     field>> dup length 1 + write-int 5 write1 write ;
 234
 235 M: request write-message
 236     [ index>> ] [ begin>> ] [ length>> ] tri
 237     13 write-int 6 write1 [ write-int ] tri@ ;
 238
 239 M: piece write-message
 240     [ index>> ] [ offset>> ] [ block>> ] tri
 241     dup length 9 + write-int 7 write1
 242     [ write-int ] [ write-int ] [ write ] tri* ;
 243
 244 M: cancel write-message
 245     [ index>> ] [ offset>> ] [ length>> ] tri
 246     13 write-int 8 write1 [ write-int ] tri@ ;
 247
 248 M: port write-message
 249     5 write-int 9 write1 port>> write-int ;
 250
 251 M: suggest-piece write-message
 252     5 write-int 0x0D write1 index>> write-int ;
 253
 254 M: have-all write-message drop 1 write-int 0x0E write1 ;
 255
 256 M: have-none write-message drop 1 write-int 0x0F write1 ;
 257
 258 M: reject-request write-message
 259     [ index>> ] [ begin>> ] [ length>> ] tri
 260     13 write-int 0x10 write1 [ write-int ] tri@ ;
 261
 262 M: allowed-fast write-message
 263     5 write-int 0x11 write1 index>> write-int ;
 264
 265 M: extended write-message
 266     [ payload>> ] [ id>> ] bi
 267     over length 2 + write-int 0x14 write1 write1 write ;
 268
 269 M: unknown write-message
 270     [ payload>> ] [ id>> ] bi
 271     over length 1 + write-int write1 write ;
 272
 273 : >message ( bytes -- message )
 274     binary [ read-message ] with-byte-reader ;
 275
 276 : message> ( message -- bytes )
 277     binary [ write-message ] with-byte-writer ;