github.com/grailbio/base@v0.0.11/embedbin/create.go (about)

     1  // Copyright 2019 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package embedbin
     6  
     7  import (
     8  	"archive/zip"
     9  	"io"
    10  	"os"
    11  )
    12  
    13  // WriteOpt is an option to NewWriter.
    14  type WriteOpt func(*Writer)
    15  
    16  // Deflate compresses embedded files.
    17  var Deflate WriteOpt = func(w *Writer) {
    18  	w.embedMethod = zip.Deflate
    19  }
    20  
    21  // Writer is used to append embedbin files to an existing binary.
    22  type Writer struct {
    23  	w io.Writer
    24  
    25  	embedOffset int64
    26  	embedZ      *zip.Writer
    27  	embedMethod uint16 // no compression by default
    28  }
    29  
    30  // NewFileWriter returns a writer that can be used to append embedbin
    31  // files to the binary represented by the provided file.
    32  // NewFileWriter removes any existing embedbin files that may be
    33  // attached to the binary. It relies on content sniffing (see Sniff)
    34  // to determine its offset.
    35  func NewFileWriter(file *os.File) (*Writer, error) {
    36  	info, err := file.Stat()
    37  	if err != nil {
    38  		return nil, err
    39  	}
    40  	embedOffset, err := Sniff(file, info.Size())
    41  	if err != nil {
    42  		return nil, err
    43  	}
    44  	if err = file.Truncate(embedOffset); err != nil {
    45  		return nil, err
    46  	}
    47  	_, err = file.Seek(0, io.SeekEnd)
    48  	if err != nil {
    49  		return nil, err
    50  	}
    51  	return NewWriter(file, embedOffset), nil
    52  }
    53  
    54  // NewWriter returns a writer that may be used to append embedbin
    55  // files to the writer w. The writer should be positioned at the end
    56  // of the base binary image.
    57  func NewWriter(w io.Writer, embedOffset int64, opts ...WriteOpt) *Writer {
    58  	ew := Writer{w: w, embedOffset: embedOffset, embedZ: zip.NewWriter(w)}
    59  	for _, opt := range opts {
    60  		opt(&ew)
    61  	}
    62  	return &ew
    63  }
    64  
    65  // Create returns a Writer into which the named file should be written.
    66  // The image's contents must be written before the next call to Create or Close.
    67  func (w *Writer) Create(name string) (io.Writer, error) {
    68  	return w.embedZ.CreateHeader(&zip.FileHeader{
    69  		Name:   name,
    70  		Method: w.embedMethod,
    71  	})
    72  }
    73  
    74  // Flush flushes the unwritten data to the underlying file.
    75  func (w *Writer) Flush() error {
    76  	return w.embedZ.Flush()
    77  }
    78  
    79  // Close should be called after all embedded files have been written.
    80  // No more files can be written after a call to Close.
    81  func (w *Writer) Close() error {
    82  	if err := w.embedZ.Close(); err != nil {
    83  		return err
    84  	}
    85  	_, err := writeFooter(w.w, w.embedOffset)
    86  	return err
    87  }