github.com/coreos/mantle@v0.13.0/update/generator/generator.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  	"encoding/binary"
    19  	"errors"
    20  	"io"
    21  	"os"
    22  
    23  	"github.com/coreos/pkg/capnslog"
    24  	"github.com/golang/protobuf/proto"
    25  
    26  	"github.com/coreos/mantle/lang/destructor"
    27  	"github.com/coreos/mantle/update/metadata"
    28  	"github.com/coreos/mantle/update/signature"
    29  )
    30  
    31  const (
    32  	// Default block size to use for all generated payloads.
    33  	BlockSize = 4096
    34  
    35  	// Default data size limit to process in a single operation.
    36  	ChunkSize = BlockSize * 256
    37  )
    38  
    39  var (
    40  	plog = capnslog.NewPackageLogger("github.com/coreos/mantle", "update/generator")
    41  
    42  	// ErrProcedureExists indicates that a given procedure type has
    43  	// already been added to the Generator.
    44  	ErrProcedureExists = errors.New("generator: procedure already exists")
    45  )
    46  
    47  // Generator assembles an update payload from a number of sources. Each of
    48  // its methods must only be called once, ending with Write.
    49  type Generator struct {
    50  	destructor.MultiDestructor
    51  	manifest metadata.DeltaArchiveManifest
    52  	payloads []io.Reader
    53  }
    54  
    55  // Procedure represent independent update within a payload.
    56  type Procedure struct {
    57  	metadata.InstallProcedure
    58  	io.ReadCloser
    59  }
    60  
    61  // Partition adds the given /usr update Procedure to the payload.
    62  // It must always be the first procedure added to the Generator.
    63  func (g *Generator) Partition(proc *Procedure) error {
    64  	if len(g.payloads) != 0 {
    65  		return ErrProcedureExists
    66  	}
    67  
    68  	g.AddCloser(proc)
    69  	g.manifest.PartitionOperations = proc.Operations
    70  	g.manifest.OldPartitionInfo = proc.OldInfo
    71  	g.manifest.NewPartitionInfo = proc.NewInfo
    72  	g.payloads = append(g.payloads, proc)
    73  	return nil
    74  }
    75  
    76  // Write finalizes the payload, writing it out to the given file path.
    77  func (g *Generator) Write(path string) (err error) {
    78  	if err = g.updateOffsets(); err != nil {
    79  		return
    80  	}
    81  
    82  	plog.Infof("Writing payload to %s", path)
    83  
    84  	f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
    85  	if err != nil {
    86  		return
    87  	}
    88  	defer func() {
    89  		if e := f.Close(); e != nil && err == nil {
    90  			err = e
    91  		}
    92  	}()
    93  
    94  	// All payload data up until the signatures must be hashed.
    95  	hasher := signature.NewSignatureHash()
    96  	w := io.MultiWriter(f, hasher)
    97  
    98  	if err = g.writeHeader(w); err != nil {
    99  		return
   100  	}
   101  
   102  	if err = g.writeManifest(w); err != nil {
   103  		return
   104  	}
   105  
   106  	for _, payload := range g.payloads {
   107  		if _, err = io.Copy(w, payload); err != nil {
   108  			return
   109  		}
   110  	}
   111  
   112  	// Hashed writes complete, write signatures to payload file.
   113  	err = g.writeSignatures(f, hasher.Sum(nil))
   114  	return
   115  }
   116  
   117  func (g *Generator) updateOffsets() error {
   118  	var offset uint32
   119  	updateOps := func(ops []*metadata.InstallOperation) {
   120  		for _, op := range ops {
   121  			if op.DataLength == nil {
   122  				op.DataOffset = nil
   123  			} else {
   124  				op.DataOffset = proto.Uint32(offset)
   125  				offset += *op.DataLength
   126  			}
   127  		}
   128  	}
   129  
   130  	updateOps(g.manifest.PartitionOperations)
   131  	for _, proc := range g.manifest.Procedures {
   132  		updateOps(proc.Operations)
   133  	}
   134  
   135  	sigSize, err := signature.SignaturesSize()
   136  	g.manifest.SignaturesOffset = proto.Uint64(uint64(offset))
   137  	g.manifest.SignaturesSize = proto.Uint64(uint64(sigSize))
   138  	return err
   139  }
   140  
   141  func (g *Generator) writeHeader(w io.Writer) error {
   142  	manifestSize := proto.Size(&g.manifest)
   143  	header := metadata.DeltaArchiveHeader{
   144  		Version:      metadata.Version,
   145  		ManifestSize: uint64(manifestSize),
   146  	}
   147  	copy(header.Magic[:], []byte(metadata.Magic))
   148  
   149  	return binary.Write(w, binary.BigEndian, &header)
   150  }
   151  
   152  func (g *Generator) writeManifest(w io.Writer) error {
   153  	buf, err := proto.Marshal(&g.manifest)
   154  	if err != nil {
   155  		return err
   156  	}
   157  
   158  	_, err = w.Write(buf)
   159  	return err
   160  }
   161  
   162  func (g *Generator) writeSignatures(w io.Writer, sum []byte) error {
   163  	signatures, err := signature.Sign(sum)
   164  	if err != nil {
   165  		return err
   166  	}
   167  
   168  	buf, err := proto.Marshal(signatures)
   169  	if err != nil {
   170  		return err
   171  	}
   172  
   173  	_, err = w.Write(buf)
   174  	return err
   175  }