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