github.com/coreos/mantle@v0.13.0/update/generator/full.go (about)

     1  // Copyright 2016 CoreOS, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package generator
    16  
    17  import (
    18  	"crypto/sha256"
    19  	"errors"
    20  	"fmt"
    21  	"io"
    22  	"os"
    23  
    24  	"github.com/golang/protobuf/proto"
    25  
    26  	"github.com/coreos/mantle/system"
    27  	"github.com/coreos/mantle/update/metadata"
    28  )
    29  
    30  var (
    31  	errShortRead = errors.New("read an incomplete block")
    32  )
    33  
    34  // FullUpdate generates an update Procedure for the given file, embedding its
    35  // entire contents in the payload so it does not depend any previous state.
    36  func FullUpdate(path string) (*Procedure, error) {
    37  	source, err := os.Open(path)
    38  	if err != nil {
    39  		return nil, err
    40  	}
    41  	defer source.Close()
    42  
    43  	info, err := NewInstallInfo(source)
    44  	if err != nil {
    45  		return nil, err
    46  	}
    47  
    48  	payload, err := system.PrivateFile("")
    49  	if err != nil {
    50  		return nil, err
    51  	}
    52  
    53  	scanner := fullScanner{payload: payload, source: source}
    54  	for err == nil {
    55  		err = scanner.Scan()
    56  	}
    57  	if err != nil && err != io.EOF {
    58  		payload.Close()
    59  		if err == errShortRead {
    60  			err = fmt.Errorf("%s: %v", path, err)
    61  		}
    62  		return nil, err
    63  	}
    64  
    65  	if _, err := payload.Seek(0, os.SEEK_SET); err != nil {
    66  		payload.Close()
    67  		return nil, err
    68  	}
    69  
    70  	return &Procedure{
    71  		InstallProcedure: metadata.InstallProcedure{
    72  			NewInfo:    info,
    73  			Operations: scanner.operations,
    74  		},
    75  		ReadCloser: payload,
    76  	}, nil
    77  }
    78  
    79  type fullScanner struct {
    80  	payload    io.Writer
    81  	source     io.Reader
    82  	offset     uint64
    83  	operations []*metadata.InstallOperation
    84  }
    85  
    86  func (f *fullScanner) readChunk() ([]byte, error) {
    87  	chunk := make([]byte, ChunkSize)
    88  	n, err := io.ReadFull(f.source, chunk)
    89  	if (err == io.EOF || err == io.ErrUnexpectedEOF) && n != 0 {
    90  		err = nil
    91  	}
    92  	return chunk[:n], err
    93  }
    94  
    95  func (f *fullScanner) Scan() error {
    96  	chunk, err := f.readChunk()
    97  	if err != nil {
    98  		return err
    99  	}
   100  	if len(chunk)%BlockSize != 0 {
   101  		return errShortRead
   102  	}
   103  
   104  	startBlock := uint64(f.offset) / BlockSize
   105  	numBlocks := uint64(len(chunk)) / BlockSize
   106  	f.offset += uint64(len(chunk))
   107  
   108  	// Try bzip2 compressing the data, hopefully it will shrink!
   109  	opType := metadata.InstallOperation_REPLACE_BZ
   110  	opData, err := Bzip2(chunk)
   111  	if err != nil {
   112  		return err
   113  	}
   114  
   115  	if len(opData) >= len(chunk) {
   116  		// That was disappointing, use the uncompressed data instead.
   117  		opType = metadata.InstallOperation_REPLACE
   118  		opData = chunk
   119  	}
   120  
   121  	if _, err := f.payload.Write(opData); err != nil {
   122  		return err
   123  	}
   124  
   125  	// Operation.DataOffset is filled in by Generator.updateOffsets
   126  	sum := sha256.Sum256(opData)
   127  	op := &metadata.InstallOperation{
   128  		Type: opType.Enum(),
   129  		DstExtents: []*metadata.Extent{&metadata.Extent{
   130  			StartBlock: proto.Uint64(startBlock),
   131  			NumBlocks:  proto.Uint64(numBlocks),
   132  		}},
   133  		DataLength:     proto.Uint32(uint32(len(opData))),
   134  		DataSha256Hash: sum[:],
   135  	}
   136  
   137  	f.operations = append(f.operations, op)
   138  
   139  	return nil
   140  }