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 byte-vectors
   4 calendar checksums checksums.sha combinators destructors fry
   5 grouping http.client io io.binary io.encodings.binary io.files
   6 io.pathnames io.sockets io.streams.byte-array io.streams.duplex
   7 kernel literals locals make math math.bitwise math.functions
   8 math.order math.parser math.ranges multiline namespaces
   9 prettyprint random sequences splitting strings timers
  10 tools.annotations ui ui.gadgets.panes ui.tools.listener urls ;
  11
  12 IN: bittorrent
  13
  14 <<
  15 CONSTANT: ALPHANUMERIC $[
  16     [
  17         CHAR: a CHAR: z [a,b] %
  18         CHAR: A CHAR: Z [a,b] %
  19         CHAR: 0 CHAR: 9 [a,b] %
  20         ".-_~" %
  21     ] { } make
  22 ]
  23
  24 : random-peer-id ( -- bytes )
  25     20 [ ALPHANUMERIC random ] B{ } replicate-as ;
  26 >>
  27
  28 SYMBOL: torrent-peer-id
  29 torrent-peer-id [ random-peer-id ] initialize
  30
  31 SYMBOL: torrent-port
  32 torrent-port [ 6881 ] initialize
  33
  34
  35 ! bitfield
  36
  37 : bitfield-index ( n -- j i )
  38     8 /mod 7 swap - ;
  39
  40 : set-bitfield ( elt n bitfield -- )
  41     [ bitfield-index rot ] dip -rot
  42     '[ _ _ [ set-bit ] [ clear-bit ] if ] change-nth ;
  43
  44 : check-bitfield ( n bitfield -- ? )
  45     [ bitfield-index swap ] dip nth swap bit? ;
  46
  47
  48 ! http
  49
  50 : http-get-bencode ( url -- obj )
  51     <get-request> BV{ } clone [
  52         '[ _ push-all ] with-http-request* check-response drop
  53     ] keep B{ } like bencode> ;
  54
  55
  56 ! metainfo
  57
  58 GENERIC: load-metainfo ( obj -- metainfo )
  59
  60 M: url load-metainfo http-get-bencode ;
  61
  62 M: pathname load-metainfo
  63     binary [ read-bencode ] with-file-reader ;
  64
  65 M: string load-metainfo
  66     dup "http" head? [ >url ] [ <pathname> ] if load-metainfo ;
  67
  68 : info-hash ( metainfo -- hash )
  69     "info hash" swap dup '[
  70         drop _ "info" of >bencode sha1 checksum-bytes
  71     ] cache ;
  72
  73 : announce-url ( metainfo -- url )
  74     dup "announce-list" of [ nip first random ] [ "announce" of ] if* ;
  75
  76 : scrape-url ( metainfo -- url/f )
  77     announce-url "announce" over path>> subseq? [
  78         [ "announce" "scrape" replace ] change-path
  79     ] [ drop f ] if ;
  80
  81
  82
  83 ! tracker
  84
  85 : tracker-url ( metainfo -- url )
  86     {
  87         [ announce-url >url ]
  88         [
  89             info-hash "info_hash" set-query-param
  90             torrent-peer-id get "peer_id" set-query-param
  91             torrent-port get "port" set-query-param
  92             0 "uploaded" set-query-param
  93             0 "downloaded" set-query-param
  94             1 "compact" set-query-param
  95         ]
  96         [
  97             { "info" "length" } [ of ] each
  98             "left" set-query-param
  99         ]
 100     } cleave ;
 101
 102 : parse-peer4 ( peerbin -- inet4 )
 103     4 cut [
 104         [ number>string ] { } map-as "." join
 105     ] dip be> <inet4> ;
 106
 107 : parse-peer4s ( peersbin -- inet4s )
 108     dup array? [
 109         [ [ "ip" of ] [ "port" of ] bi <inet4> ] map
 110     ] [
 111         6 <groups> [ parse-peer4 ] map
 112     ] if ;
 113
 114 : parse-peer6 ( peerbin -- inet6 )
 115     16 cut [
 116         2 <groups> [ be> number>string ] map ":" join
 117     ] dip be> <inet6> ;
 118
 119 : parse-peer6s ( peersbin -- inet6s )
 120     18 <groups> [ parse-peer6 ] map ;
 121
 122 : load-tracker ( torrent -- response )
 123     tracker-url http-get-bencode
 124     "peers" over [ parse-peer4s ] change-at ;
 125
 126 : send-event ( torrent event -- response )
 127     [ tracker-url ] [ "event" set-query-param ] bi*
 128     http-get-bencode ;
 129
 130
 131
 132 ! messages
 133
 134 TUPLE: handshake string reserved info-hash peer-id ;
 135
 136 : <handshake> ( info-hash peer-id -- handshake )
 137     handshake new
 138         "BitTorrent protocol" >byte-array >>string
 139         8 <byte-array> >>reserved
 140         swap >>peer-id
 141         swap >>info-hash ;
 142
 143 : read-handshake ( -- handshake/f )
 144     read1 [
 145         [ 48 + read ] keep cut 8 cut 20 cut handshake boa
 146     ] [ f ] if* ;
 147
 148 : write-handshake ( handshake -- )
 149     {
 150         [ string>> [ length write1 ] [ write ] bi ]
 151         [ reserved>> write ]
 152         [ info-hash>> write ]
 153         [ peer-id>> write ]
 154     } cleave flush ;
 155
 156 TUPLE: keep-alive ;
 157 TUPLE: choke ;
 158 TUPLE: unchoke ;
 159 TUPLE: interested ;
 160 TUPLE: not-interested ;
 161 TUPLE: have index ;
 162 TUPLE: bitfield bitfield ;
 163 TUPLE: request index begin length ;
 164 TUPLE: piece index begin block ;
 165 TUPLE: cancel index begin length ;
 166 TUPLE: port port ;
 167 TUPLE: suggest-piece index ;
 168 TUPLE: have-all ;
 169 TUPLE: have-none ;
 170 TUPLE: reject-request index begin length ;
 171 TUPLE: allowed-fast index ;
 172 TUPLE: extended id payload ;
 173 TUPLE: unknown id payload ;
 174
 175 : read-int ( -- n/f ) 4 read [ be> ] [ f ] if* ;
 176
 177 : parse-message ( bytes -- message/f )
 178     unclip {
 179         ! Core Protocol
 180         { 0 [ drop choke boa ] }
 181         { 1 [ drop unchoke boa ] }
 182         { 2 [ drop interested boa ] }
 183         { 3 [ drop not-interested boa ] }
 184         { 4 [ 4 head be> have boa ] }
 185         { 5 [ bitfield boa ] }
 186         { 6 [ 4 cut 4 cut 4 head [ be> ] tri@ request boa ] }
 187         { 7 [ 4 cut 4 cut [ [ be> ] bi@ ] dip piece boa ] }
 188         { 8 [ 4 cut 4 cut 4 head [ be> ] tri@ cancel boa ] }
 189
 190         ! DHT Extension
 191         { 9 [ be> port boa ] }
 192
 193         ! Fast Extensions
 194         { 0x0D [ 4 head be> suggest-piece boa ] }
 195         { 0x0E [ drop have-all boa ] }
 196         { 0x0F [ drop have-none boa ] }
 197         { 0x10 [ 4 cut 4 cut 4 head [ be> ] tri@ reject-request boa ] }
 198         { 0x11 [ 4 head be> allowed-fast boa ] }
 199
 200         ! Extension Protocol
 201         { 0x14 [ unclip swap extended boa ] }
 202
 203         ! Hash Transfer Protocol
 204         ! { 0x15 [ "HashRequest" ] }
 205         ! { 0x16 [ "Hashes" ] }
 206         ! { 0x17 [ "HashReject" ] }
 207         [ swap unknown boa ]
 208     } case ;
 209
 210 : read-message ( -- message )
 211     read-int {
 212         { f [ f ] }
 213         { 0 [ keep-alive boa ] }
 214         [ read [ parse-message ] [ f ] if* ]
 215     } case ;
 216
 217 : write-int ( n -- ) 4 >be write ;
 218
 219 GENERIC: write-message ( message -- )
 220
 221 M: keep-alive write-message drop 0 write-int ;
 222
 223 M: choke write-message drop 1 write-int 0 write1 ;
 224
 225 M: unchoke write-message drop 1 write-int 1 write1 ;
 226
 227 M: interested write-message drop 1 write-int 2 write1 ;
 228
 229 M: not-interested write-message drop 1 write-int 3 write1 ;
 230
 231 M: have write-message
 232     5 write-int 4 write1 index>> write-int ;
 233
 234 M: bitfield write-message
 235     field>> dup length 1 + write-int 5 write1 write ;
 236
 237 M: request write-message
 238     [ index>> ] [ begin>> ] [ length>> ] tri
 239     13 write-int 6 write1 [ write-int ] tri@ ;
 240
 241 M: piece write-message
 242     [ index>> ] [ offset>> ] [ block>> ] tri
 243     dup length 9 + write-int 7 write1
 244     [ write-int ] [ write-int ] [ write ] tri* ;
 245
 246 M: cancel write-message
 247     [ index>> ] [ offset>> ] [ length>> ] tri
 248     13 write-int 8 write1 [ write-int ] tri@ ;
 249
 250 M: port write-message
 251     5 write-int 9 write1 port>> write-int ;
 252
 253 M: suggest-piece write-message
 254     5 write-int 0x0D write1 index>> write-int ;
 255
 256 M: have-all write-message drop 1 write-int 0x0E write1 ;
 257
 258 M: have-none write-message drop 1 write-int 0x0F write1 ;
 259
 260 M: reject-request write-message
 261     [ index>> ] [ begin>> ] [ length>> ] tri
 262     13 write-int 0x10 write1 [ write-int ] tri@ ;
 263
 264 M: allowed-fast write-message
 265     5 write-int 0x11 write1 index>> write-int ;
 266
 267 M: extended write-message
 268     [ payload>> ] [ id>> ] bi
 269     over length 2 + write-int 0x14 write1 write1 write ;
 270
 271 M: unknown write-message
 272     [ payload>> ] [ id>> ] bi
 273     over length 1 + write-int write1 write ;
 274
 275 : >message ( bytes -- message )
 276     binary [ read-message ] with-byte-reader ;
 277
 278 : message> ( message -- bytes )
 279     binary [ write-message ] with-byte-writer ;
 280
 281
 282
 283
 284 SYMBOL: torrent-max-block-size
 285 torrent-max-block-size [ 16384 ] initialize
 286
 287 SYMBOL: torrent-directory
 288 torrent-directory [ "~/Desktop" ] initialize
 289
 290
 291
 292 : with-debug ( quot -- )
 293     ui-running? [
 294         get-listener output>> <pane-stream> swap
 295         '[ @ flush ] with-output-stream*
 296     ] [
 297         '[ @ flush ] with-global
 298     ] if ; inline
 299
 300
 301 ! connection
 302
 303 TUPLE: torrent metainfo tracker peers ;
 304
 305 :: <torrent> ( metainfo -- torrent )
 306     torrent new
 307         metainfo >>metainfo
 308         metainfo load-tracker >>tracker
 309         V{ } clone >>peers
 310     ;
 311
 312 : torrent-path ( torrent -- path )
 313     metainfo>> { "info" "name" } [ of ] each
 314     torrent-directory get prepend-path ;
 315
 316
 317 DEFER: <peer>
 318
 319 : random-peer ( torrent -- peer )
 320     dup tracker>> "peers" of random <peer> ;
 321
 322 : update-tracker ( client -- client )
 323     dup metainfo>> load-tracker >>tracker ;
 324
 325
 326 ! peers
 327
 328 TUPLE: peer < disposable torrent remote stream local handshake
 329 self-choking self-interested peer-choking peer-interested timer
 330 #pieces piece-length bitfield current-index current-piece ;
 331
 332 DEFER: with-peer
 333
 334 :: <peer> ( torrent remote -- peer )
 335     peer new
 336         torrent >>torrent
 337         remote >>remote
 338
 339         t >>self-choking
 340         f >>self-interested
 341         t >>peer-choking
 342         f >>peer-interested
 343
 344         dup '[
 345             _ [ T{ keep-alive } write-message flush ] with-peer
 346         ] 30 seconds dup <timer> >>timer
 347
 348         torrent metainfo>> "info" of
 349         "pieces" of length 20 / [ >>#pieces ] keep
 350
 351         8 / ceiling <byte-array> >>bitfield
 352
 353         torrent metainfo>> "info" of
 354         "piece length" of [ >>piece-length ] keep
 355
 356         <byte-vector> >>current-piece
 357     ;
 358
 359 M: peer dispose
 360     dup timer>> stop-timer
 361     [ dispose f ] change-stream
 362     f >>local f >>remote drop ;
 363
 364 :: with-peer ( peer quot -- )
 365     [
 366         peer remote>> remote-address set
 367         peer local>> local-address set
 368         peer stream>> quot with-stream*
 369     ] with-scope ; inline
 370
 371 : connect-peer ( peer -- peer )
 372     dup remote>> binary <client> [ >>stream ] [ >>local ] bi* ;
 373
 374 : handshake-peer ( peer -- peer )
 375     dup torrent>> metainfo>> info-hash
 376     torrent-peer-id get <handshake> write-handshake
 377     read-handshake >>handshake ;
 378
 379 : fast-peer? ( peer -- ? )
 380     handshake>> reserved>> 7 swap nth 3 swap bit? ;
 381
 382 : unchoke-peer ( peer -- peer )
 383     T{ unchoke } write-message f >>self-choking
 384     T{ interested } write-message t >>self-interested
 385     flush ;
 386
 387 : choke-peer ( peer -- peer )
 388     T{ choke } write-message t >>self-choking
 389     T{ not-interested } write-message f >>self-interested
 390     flush ;
 391
 392 :: verify-block ( peer -- peer )
 393     peer current-piece>> sha1 checksum-bytes
 394     peer current-index>> 20 * dup 20 +
 395     peer torrent>> metainfo>>
 396     { "info" "pieces" } [ of ] each
 397     B{ } subseq-as assert=
 398     peer ;
 399
 400 :: save-block ( peer -- peer )
 401     peer torrent>> torrent-path binary [
 402         peer current-index>>
 403         peer piece-length>> *
 404         seek-absolute seek-output
 405         peer current-piece>> write
 406     ] with-file-appender peer ;
 407
 408 :: next-block ( peer -- peer )
 409     peer current-index>> [ 1 + ] [ 0 ] if*
 410     peer #pieces>>
 411     peer bitfield>> '[ _ check-bitfield ] (find-integer)
 412     peer current-index<<
 413     0 peer current-piece>> set-length
 414     peer ;
 415
 416 :: request-piece ( peer -- peer )
 417     peer current-index>>
 418     [ peer next-block current-index>> ] unless*
 419     peer current-piece>> length
 420     peer piece-length>> over - torrent-max-block-size get min
 421     [
 422         2drop
 423         peer
 424         verify-block
 425         save-block
 426         next-block
 427         request-piece
 428     ] [
 429         request boa write-message flush peer
 430     ] if-zero ;
 431
 432 GENERIC: handle-message ( peer message -- peer )
 433
 434 M: object handle-message drop ;
 435
 436 M: choke handle-message
 437     drop t >>peer-choking ;
 438
 439 M: unchoke handle-message
 440     drop f >>peer-choking
 441     dup self-choking>> [ next-block request-piece ] unless ;
 442
 443 M: interested handle-message
 444     drop t >>peer-interested ;
 445
 446 M: not-interested handle-message
 447     drop f >>peer-interested ;
 448
 449 M: have handle-message
 450     t swap index>> pick bitfield>> set-nth ;
 451
 452 M: bitfield handle-message
 453     2dup [ bitfield>> length ] bi@ assert=
 454     bitfield>> >>bitfield ;
 455
 456 M: request handle-message
 457     [ index>> ] [ begin>> ] [ length>> ] tri
 458     reject-request boa write-message ;
 459
 460 M: have-all handle-message
 461     drop [ length [ 255 ] B{ } replicate-as ] change-bitfield ;
 462
 463 M: have-none handle-message
 464     drop [ length <byte-array> ] change-bitfield ;
 465
 466 M:: piece handle-message ( peer message -- peer )
 467     peer current-index>> message index>> assert=
 468     peer current-piece>> length message begin>> assert=
 469     message [ block>> ] [ begin>> ] bi peer current-piece>> copy
 470     peer request-piece ;
 471
 472 : read-messages ( peer -- peer )
 473     [ read-message ] [ handle-message ] while* ;