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