github.com/sabhiram/gsync@v0.0.0-20180214150410-b9089a1b7f13/gsync_server.go (about)

     1  // This Source Code Form is subject to the terms of the Mozilla Public
     2  // License, version 2.0. If a copy of the MPL was not distributed with this
     3  // file, You can obtain one at http://mozilla.org/MPL/2.0/.
     4  
     5  package gsync
     6  
     7  import (
     8  	"context"
     9  	"crypto/sha256"
    10  	"hash"
    11  	"io"
    12  	"os"
    13  
    14  	"github.com/pkg/errors"
    15  )
    16  
    17  // Signatures reads data blocks from reader and pipes out block signatures on the
    18  // returning channel, closing it when done reading or when the context is cancelled.
    19  // This function does not block and returns immediately. The caller must make sure the concrete
    20  // reader instance is not nil or this function will panic.
    21  func Signatures(ctx context.Context, r io.Reader, shash hash.Hash) (<-chan BlockSignature, error) {
    22  	var index uint64
    23  	c := make(chan BlockSignature)
    24  
    25  	bfp := bufferPool.Get().(*[]byte)
    26  	buffer := *bfp
    27  	defer bufferPool.Put(bfp)
    28  
    29  	if r == nil {
    30  		close(c)
    31  		return nil, errors.New("gsync: reader required")
    32  	}
    33  
    34  	if shash == nil {
    35  		shash = sha256.New()
    36  	}
    37  
    38  	go func() {
    39  		defer close(c)
    40  
    41  		for {
    42  			// Allow for cancellation
    43  			select {
    44  			case <-ctx.Done():
    45  				c <- BlockSignature{
    46  					Index: index,
    47  					Error: ctx.Err(),
    48  				}
    49  				return
    50  			default:
    51  				// break out of the select block and continue reading
    52  				break
    53  			}
    54  
    55  			n, err := r.Read(buffer)
    56  			if err == io.EOF {
    57  				break
    58  			}
    59  
    60  			if err != nil {
    61  				c <- BlockSignature{
    62  					Index: index,
    63  					Error: errors.Wrapf(err, "failed reading block"),
    64  				}
    65  				index++
    66  				// let the caller decide whether to interrupt the process or not.
    67  				continue
    68  			}
    69  
    70  			block := buffer[:n]
    71  			shash.Reset()
    72  			shash.Write(block)
    73  			strong := shash.Sum(nil)
    74  			_, _, rhash := rollingHash(block)
    75  
    76  			c <- BlockSignature{
    77  				Index:  index,
    78  				Weak:   rhash,
    79  				Strong: strong,
    80  			}
    81  			index++
    82  		}
    83  	}()
    84  
    85  	return c, nil
    86  }
    87  
    88  // Apply reconstructs a file given a set of operations. The caller must close the ops channel or the context when done or there will be a deadlock.
    89  func Apply(ctx context.Context, dst io.Writer, cache io.ReaderAt, ops <-chan BlockOperation) error {
    90  	bfp := bufferPool.Get().(*[]byte)
    91  	buffer := *bfp
    92  	defer bufferPool.Put(bfp)
    93  
    94  	for o := range ops {
    95  		// Allows for cancellation.
    96  		select {
    97  		case <-ctx.Done():
    98  			return errors.Wrapf(ctx.Err(), "failed applying block operations")
    99  		default:
   100  			// break out of the select block and continue reading ops
   101  			break
   102  		}
   103  
   104  		if o.Error != nil {
   105  			return errors.Wrapf(o.Error, "failed applying operation")
   106  		}
   107  
   108  		var block []byte
   109  
   110  		if len(o.Data) > 0 {
   111  			block = o.Data
   112  		} else {
   113  			if f, ok := cache.(*os.File); ok && f == nil {
   114  				return errors.New("index operation, but cached file was not found")
   115  			}
   116  
   117  			index := int64(o.Index)
   118  			n, err := cache.ReadAt(buffer, (index * DefaultBlockSize))
   119  			if err != nil && err != io.EOF {
   120  				return errors.Wrapf(err, "failed reading cached block")
   121  			}
   122  
   123  			block = buffer[:n]
   124  		}
   125  
   126  		_, err := dst.Write(block)
   127  		if err != nil {
   128  			return errors.Wrapf(err, "failed writing block to destination")
   129  		}
   130  	}
   131  	return nil
   132  }