metainfo/builder.go

   1 package metainfo
   2
   3 import (
   4         "crypto/sha1"
   5         "errors"
   6         "hash"
   7         "io"
   8         "os"
   9         "path/filepath"
  10         "sort"
  11         "time"
  12
  13         "github.com/anacrolix/missinggo"
  14
  15         "github.com/anacrolix/torrent/bencode"
  16 )
  17
  18 //----------------------------------------------------------------------------
  19 // Build
  20 //----------------------------------------------------------------------------
  21
  22 // The Builder type is responsible for .torrent files construction. Just
  23 // instantiate it, call necessary methods and then call the .Build method. While
  24 // waiting for completion you can use 'status' channel to get status reports.
  25 type Builder struct {
  26         batch_state
  27         filesmap map[string]bool
  28 }
  29
  30 // Adds a file to the builder queue. You may add one or more files.
  31 func (b *Builder) AddFile(filename string) {
  32         if b.filesmap == nil {
  33                 b.filesmap = make(map[string]bool)
  34         }
  35
  36         filename, err := filepath.Abs(filename)
  37         if err != nil {
  38                 panic(err)
  39         }
  40         b.filesmap[filename] = true
  41 }
  42
  43 // Defines a name of the future torrent file. For single file torrents it's the
  44 // recommended name of the contained file. For multiple files torrents it's the
  45 // recommended name of the directory in which all of them will be
  46 // stored. Calling this function is not required. In case if no name was
  47 // specified, the builder will try to automatically assign it. It will use the
  48 // name of the file if there is only one file in the queue or it will try to
  49 // find the rightmost common directory of all the queued files and use its name as
  50 // a torrent name. In case if name cannot be assigned automatically, it will use
  51 // "unknown" as a torrent name.
  52 func (b *Builder) SetName(name string) {
  53         b.name = name
  54 }
  55
  56 // Sets the length of a piece in the torrent file in bytes. The default is
  57 // 256kb.
  58 func (b *Builder) SetPieceLength(length int64) {
  59         b.piece_length = length
  60 }
  61
  62 // Sets the "private" flag. The default is false.
  63 func (b *Builder) SetPrivate(v bool) {
  64         b.private = v
  65 }
  66
  67 // Add announce URL group. TODO: better explanation.
  68 func (b *Builder) AddAnnounceGroup(group []string) {
  69         b.announce_list = append(b.announce_list, group)
  70 }
  71
  72 // Add DHT nodes URLs for trackerless mode
  73 func (b *Builder) AddDhtNodes(group []string) {
  74         b.node_list = append(b.node_list, group...)
  75 }
  76
  77 // Sets creation date. The default is time.Now() when the .Build method was
  78 // called.
  79 func (b *Builder) SetCreationDate(date time.Time) {
  80         b.creation_date = date
  81 }
  82
  83 // Sets the comment. The default is no comment.
  84 func (b *Builder) SetComment(comment string) {
  85         b.comment = comment
  86 }
  87
  88 // Sets the "created by" parameter. The default is "libtorgo".
  89 func (b *Builder) SetCreatedBy(createdby string) {
  90         b.created_by = createdby
  91 }
  92
  93 // Sets the "encoding" parameter. The default is "UTF-8".
  94 func (b *Builder) SetEncoding(encoding string) {
  95         b.encoding = encoding
  96 }
  97
  98 // Add WebSeed URL to the list.
  99 func (b *Builder) AddWebSeedURL(url string) {
 100         b.urls = append(b.urls, url)
 101 }
 102
 103 // Finalizes the Builder state and makes a Batch out of it. After calling that
 104 // method, Builder becomes empty and you can use it to create another Batch if
 105 // you will.
 106 func (b *Builder) Submit() (*Batch, error) {
 107         err := b.check_parameters()
 108         if err != nil {
 109                 return nil, err
 110         }
 111         b.set_defaults()
 112
 113         batch := &Batch{
 114                 batch_state: b.batch_state,
 115         }
 116
 117         const non_regular = os.ModeDir | os.ModeSymlink |
 118                 os.ModeDevice | os.ModeNamedPipe | os.ModeSocket
 119
 120         // convert a map to a slice, calculate sizes and split paths
 121         batch.total_size = 0
 122         batch.files = make([]file, 0, 10)
 123         for f, _ := range b.filesmap {
 124                 var file file
 125                 fi, err := os.Stat(f)
 126                 if err != nil {
 127                         return nil, err
 128                 }
 129
 130                 if fi.Mode()&non_regular != 0 {
 131                         return nil, errors.New(f + " is not a regular file")
 132                 }
 133
 134                 file.abspath = f
 135                 file.splitpath = split_path(f)
 136                 file.size = fi.Size()
 137                 batch.files = append(batch.files, file)
 138                 batch.total_size += file.size
 139         }
 140
 141         // find the rightmost common directory
 142         if len(batch.files) == 1 {
 143                 sp := batch.files[0].splitpath
 144                 batch.default_name = sp[len(sp)-1]
 145         } else {
 146                 common := batch.files[0].splitpath
 147                 for _, f := range batch.files {
 148                         if len(common) > len(f.splitpath) {
 149                                 common = common[:len(f.splitpath)]
 150                         }
 151
 152                         for i, n := 0, len(common); i < n; i++ {
 153                                 if common[i] != f.splitpath[i] {
 154                                         common = common[:i]
 155                                         break
 156                                 }
 157                         }
 158
 159                         if len(common) == 0 {
 160                                 break
 161                         }
 162                 }
 163
 164                 if len(common) == 0 {
 165                         return nil, errors.New("no common rightmost folder was found for a set of queued files")
 166                 }
 167
 168                 // found the common folder, let's strip that part from splitpath
 169                 // and setup the default name
 170                 batch.default_name = common[len(common)-1]
 171
 172                 lcommon := len(common)
 173                 for i := range batch.files {
 174                         f := &batch.files[i]
 175                         f.splitpath = f.splitpath[lcommon:]
 176                 }
 177
 178                 // and finally sort the files
 179                 sort.Sort(file_slice(batch.files))
 180         }
 181
 182         // reset the builder state
 183         b.batch_state = batch_state{}
 184         b.filesmap = nil
 185
 186         return batch, nil
 187 }
 188
 189 func (b *Builder) set_defaults() {
 190         if b.piece_length == 0 {
 191                 b.piece_length = 256 * 1024
 192         }
 193
 194         if b.creation_date.IsZero() {
 195                 b.creation_date = time.Now()
 196         }
 197
 198         if b.created_by == "" {
 199                 b.created_by = "libtorgo"
 200         }
 201
 202         if b.encoding == "" {
 203                 b.encoding = "UTF-8"
 204         }
 205 }
 206
 207 func emptyStringsFiltered(ss []string) (ret []string) {
 208         for _, s := range ss {
 209                 if s != "" {
 210                         ret = append(ret, s)
 211                 }
 212         }
 213         return
 214 }
 215
 216 func (b *Builder) check_parameters() error {
 217         // should be at least one file
 218         if len(b.filesmap) == 0 {
 219                 return errors.New("no files were queued")
 220         }
 221
 222         // let's clean up the announce_list and node_list
 223         b.announce_list = cleanUpLists(b.announce_list)
 224         b.node_list = emptyStringsFiltered(b.node_list)
 225
 226         if len(b.announce_list) == 0 && len(b.node_list) == 0 {
 227                 return errors.New("no announce group or DHT nodes specified")
 228         }
 229
 230         // Either the node_list or announce_list can be present
 231         // Never the both!
 232         if len(b.announce_list) > 0 && len(b.node_list) > 0 {
 233                 return errors.New("announce group and nodes are mutually exclusive")
 234         }
 235
 236         // and clean up the urls
 237         b.urls = remove_empty_strings(b.urls)
 238
 239         return nil
 240 }
 241
 242 func cleanUpLists(list [][]string) [][]string {
 243         newList := make([][]string, 0, len(list))
 244         for _, l := range list {
 245                 l = remove_empty_strings(l)
 246
 247                 // discard empty announce groups
 248                 if len(l) == 0 {
 249                         continue
 250                 }
 251                 newList = append(newList, l)
 252         }
 253         return newList
 254 }
 255
 256 //----------------------------------------------------------------------------
 257 // Batch
 258 //----------------------------------------------------------------------------
 259
 260 // Batch represents a snapshot of a builder state, ready for transforming it
 261 // into a torrent file. Note that Batch contains two accessor methods you might
 262 // be interested in. The TotalSize is the total size of all the files queued for
 263 // hashing, you will use it for status reporting. The DefaultName is an
 264 // automatically determined name of the torrent metainfo, you might want to use
 265 // it for naming the .torrent file itself.
 266 type Batch struct {
 267         batch_state
 268         files        []file
 269         total_size   int64
 270         default_name string
 271 }
 272
 273 // Get a total size of all the files queued for hashing. Useful in conjunction
 274 // with status reports.
 275 func (b *Batch) TotalSize() int64 {
 276         return b.total_size
 277 }
 278
 279 // Get an automatically determined name of the future torrent metainfo. You can
 280 // use it for a .torrent file in case user hasn't provided it specifically.
 281 func (b *Batch) DefaultName() string {
 282         return b.default_name
 283 }
 284
 285 // Starts a process of building the torrent file. This function does everything
 286 // in a separate goroutine and uses up to 'nworkers' of goroutines to perform
 287 // SHA1 hashing. Therefore it will return almost immedately. It returns two
 288 // channels, the first one is for completion awaiting, the second one is for
 289 // getting status reports. Status report is a number of bytes hashed, you can
 290 // get the total amount of bytes by inspecting the Batch.TotalSize method return
 291 // value.
 292 func (b *Batch) Start(w io.Writer, nworkers int) (<-chan error, <-chan int64) {
 293         if nworkers <= 0 {
 294                 nworkers = 1
 295         }
 296
 297         completion := make(chan error)
 298         status := make(chan int64)
 299
 300         go func() {
 301                 // prepare workers
 302                 workers := make([]*worker, nworkers)
 303                 free_workers := make(chan *worker, nworkers)
 304                 for i := 0; i < nworkers; i++ {
 305                         workers[i] = new_worker(free_workers)
 306                 }
 307                 stop_workers := func() {
 308                         for _, w := range workers {
 309                                 w.stop()
 310                         }
 311                         for _, w := range workers {
 312                                 w.wait_for_stop()
 313                         }
 314                 }
 315
 316                 // prepare files for reading
 317                 fr := files_reader{files: b.files}
 318                 npieces := (b.total_size + b.piece_length - 1) / b.piece_length
 319                 b.pieces = make([]byte, 20*npieces)
 320                 hashed := int64(0)
 321
 322                 // read all the pieces passing them to workers for hashing
 323                 var data []byte
 324                 for i := int64(0); i < npieces; i++ {
 325                         if data == nil {
 326                                 data = make([]byte, b.piece_length)
 327                         }
 328
 329                         nr, err := fr.Read(data)
 330                         if err != nil {
 331                                 // EOF is not an eror if it was the last piece
 332                                 if err == io.EOF {
 333                                         if i != npieces-1 {
 334                                                 stop_workers()
 335                                                 completion <- err
 336                                                 return
 337                                         }
 338                                 } else {
 339                                         stop_workers()
 340                                         completion <- err
 341                                         return
 342                                 }
 343                         }
 344
 345                         // cut the data slice to the amount of actual data read
 346                         data = data[:nr]
 347                         w := <-free_workers
 348                         data = w.queue(data, b.pieces[20*i:20*i+20])
 349
 350                         // update and try to send the status report
 351                         if data != nil {
 352                                 hashed += int64(len(data))
 353                                 data = data[:cap(data)]
 354
 355                                 select {
 356                                 case status <- hashed:
 357                                 default:
 358                                 }
 359                         }
 360                 }
 361                 stop_workers()
 362
 363                 // at this point the hash was calculated and we're ready to
 364                 // write the torrent file
 365                 err := b.write_torrent(w)
 366                 if err != nil {
 367                         completion <- err
 368                         return
 369                 }
 370                 completion <- nil
 371         }()
 372         return completion, status
 373 }
 374
 375 func (b *Batch) write_torrent(w io.Writer) error {
 376         var td MetaInfo
 377
 378         // Either announce or node lists are allowed - not both
 379         if len(b.announce_list) != 0 {
 380                 td.Announce = b.announce_list[0][0]
 381                 if len(b.announce_list) != 1 || len(b.announce_list[0]) != 1 {
 382                         td.AnnounceList = b.announce_list
 383                 }
 384         }
 385
 386         missinggo.CastSlice(&td.Nodes, b.node_list)
 387         td.CreationDate = b.creation_date.Unix()
 388         td.Comment = b.comment
 389         td.CreatedBy = b.created_by
 390         td.Encoding = b.encoding
 391         switch {
 392         case len(b.urls) == 0:
 393         case len(b.urls) == 1:
 394                 td.URLList = b.urls[0]
 395         default:
 396                 td.URLList = b.urls
 397         }
 398
 399         td.Info.PieceLength = b.piece_length
 400         td.Info.Pieces = b.pieces
 401         if b.name == "" {
 402                 td.Info.Name = b.default_name
 403         } else {
 404                 td.Info.Name = b.name
 405         }
 406         if len(b.files) == 1 {
 407                 td.Info.Length = b.files[0].size
 408         } else {
 409                 td.Info.Files = make([]FileInfo, len(b.files))
 410                 for i, f := range b.files {
 411                         td.Info.Files[i] = FileInfo{
 412                                 Path:   f.splitpath,
 413                                 Length: f.size,
 414                         }
 415                 }
 416         }
 417         td.Info.Private = b.private
 418
 419         e := bencode.NewEncoder(w)
 420         return e.Encode(&td)
 421 }
 422
 423 //----------------------------------------------------------------------------
 424 // misc stuff
 425 //----------------------------------------------------------------------------
 426
 427 // splits path into components (dirs and files), works only on absolute paths
 428 func split_path(path string) []string {
 429         var dir, file string
 430         s := make([]string, 0, 5)
 431
 432         dir = path
 433         for {
 434                 dir, file = filepath.Split(filepath.Clean(dir))
 435                 if file == "" {
 436                         break
 437                 }
 438                 s = append(s, file)
 439         }
 440
 441         // reverse the slice
 442         for i, n := 0, len(s)/2; i < n; i++ {
 443                 i2 := len(s) - i - 1
 444                 s[i], s[i2] = s[i2], s[i]
 445         }
 446
 447         return s
 448 }
 449
 450 // just a common data between the Builder and the Batch
 451 type batch_state struct {
 452         name          string
 453         piece_length  int64
 454         pieces        []byte
 455         private       bool
 456         announce_list [][]string
 457         node_list     []string
 458         creation_date time.Time
 459         comment       string
 460         created_by    string
 461         encoding      string
 462         urls          []string
 463 }
 464
 465 type file struct {
 466         abspath   string
 467         splitpath []string
 468         size      int64
 469 }
 470
 471 type file_slice []file
 472
 473 func (s file_slice) Len() int           { return len(s) }
 474 func (s file_slice) Less(i, j int) bool { return s[i].abspath < s[j].abspath }
 475 func (s file_slice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
 476
 477 func remove_empty_strings(slice []string) []string {
 478         j := 0
 479         for i, n := 0, len(slice); i < n; i++ {
 480                 if slice[i] == "" {
 481                         continue
 482                 }
 483                 slice[j] = slice[i]
 484                 j++
 485         }
 486         return slice[:j]
 487 }
 488
 489 //----------------------------------------------------------------------------
 490 // worker
 491 //----------------------------------------------------------------------------
 492
 493 type worker struct {
 494         msgbox chan bool
 495         hash   hash.Hash
 496
 497         // request
 498         sha1 []byte
 499         data []byte
 500 }
 501
 502 // returns existing 'data'
 503 func (w *worker) queue(data, sha1 []byte) []byte {
 504         d := w.data
 505         w.data = data
 506         w.sha1 = sha1
 507         w.msgbox <- false
 508         return d
 509 }
 510
 511 func (w *worker) stop() {
 512         w.msgbox <- true
 513 }
 514
 515 func (w *worker) wait_for_stop() {
 516         <-w.msgbox
 517 }
 518
 519 func new_worker(out chan<- *worker) *worker {
 520         w := &worker{
 521                 msgbox: make(chan bool),
 522                 hash:   sha1.New(),
 523         }
 524         go func() {
 525                 var sha1 [20]byte
 526                 for {
 527                         if <-w.msgbox {
 528                                 w.msgbox <- true
 529                                 return
 530                         }
 531                         w.hash.Reset()
 532                         w.hash.Write(w.data)
 533                         w.hash.Sum(sha1[:0])
 534                         copy(w.sha1, sha1[:])
 535                         out <- w
 536                 }
 537         }()
 538         out <- w
 539         return w
 540 }
 541
 542 //----------------------------------------------------------------------------
 543 // files_reader
 544 //----------------------------------------------------------------------------
 545
 546 type files_reader struct {
 547         files   []file
 548         cur     int
 549         curfile *os.File
 550         off     int64
 551 }
 552
 553 func (f *files_reader) Read(data []byte) (int, error) {
 554         if f.cur >= len(f.files) {
 555                 return 0, io.EOF
 556         }
 557
 558         if len(data) == 0 {
 559                 return 0, nil
 560         }
 561
 562         read := 0
 563         for len(data) > 0 {
 564                 file := &f.files[f.cur]
 565                 if f.curfile == nil {
 566                         var err error
 567                         f.curfile, err = os.Open(file.abspath)
 568                         if err != nil {
 569                                 return read, err
 570                         }
 571                 }
 572
 573                 // we need to read up to 'len(data)' bytes from current file
 574                 n := int64(len(data))
 575
 576                 // unless there is not enough data in this file
 577                 if file.size-f.off < n {
 578                         n = file.size - f.off
 579                 }
 580
 581                 // if there is no data in this file, try next one
 582                 if n == 0 {
 583                         err := f.curfile.Close()
 584                         if err != nil {
 585                                 return read, err
 586                         }
 587
 588                         f.curfile = nil
 589                         f.off = 0
 590                         f.cur++
 591                         if f.cur >= len(f.files) {
 592                                 return read, io.EOF
 593                         }
 594                         continue
 595                 }
 596
 597                 // read, handle errors
 598                 nr, err := f.curfile.Read(data[:n])
 599                 read += nr
 600                 f.off += int64(nr)
 601                 if err != nil {
 602                         return read, err
 603                 }
 604
 605                 // ok, we've read nr bytes out of len(data), cut the data slice
 606                 data = data[nr:]
 607         }
 608
 609         return read, nil
 610 }