metainfo/builder.go

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