github.com/grailbio/base@v0.0.11/fatbin/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 fatbin
     6  
     7  import (
     8  	"archive/zip"
     9  	"io"
    10  	"os"
    11  )
    12  
    13  // Writer is used to append fatbin images to an existing binary.
    14  type Writer struct {
    15  	z            *zip.Writer
    16  	w            io.Writer
    17  	off          int64
    18  	goos, goarch string
    19  }
    20  
    21  // NewFileWriter returns a writer that can be used to append fatbin
    22  // images to the binary represented by the provided file.
    23  // NewFileWriter removes any existing fatbin images that may be
    24  // attached to the binary. It relies on content sniffing (see Sniff)
    25  // to determine its offset.
    26  func NewFileWriter(file *os.File) (*Writer, error) {
    27  	info, err := file.Stat()
    28  	if err != nil {
    29  		return nil, err
    30  	}
    31  	goos, goarch, offset, err := Sniff(file, info.Size())
    32  	if err != nil {
    33  		return nil, err
    34  	}
    35  	if err := file.Truncate(offset); err != nil {
    36  		return nil, err
    37  	}
    38  	_, err = file.Seek(0, io.SeekEnd)
    39  	if err != nil {
    40  		return nil, err
    41  	}
    42  	return NewWriter(file, offset, goos, goarch), nil
    43  }
    44  
    45  // NewWriter returns a writer that may be used to append fatbin
    46  // images to the writer w. The writer should be positioned at the end
    47  // of the base binary image.
    48  func NewWriter(w io.Writer, offset int64, goos, goarch string) *Writer {
    49  	return &Writer{z: zip.NewWriter(w), w: w, off: offset, goos: goos, goarch: goarch}
    50  }
    51  
    52  // Create returns a Writer into which the image for the provided goos
    53  // and goarch should be written. The image's contents must be written
    54  // before the next call to Create or Close.
    55  func (w *Writer) Create(goos, goarch string) (io.Writer, error) {
    56  	return w.z.Create(goos + "/" + goarch)
    57  }
    58  
    59  // Flush flushes the unwritten data to the underlying file.
    60  func (w *Writer) Flush() error {
    61  	return w.z.Flush()
    62  }
    63  
    64  // Close should be called after all images have been written. No more
    65  // images can be written after a call to Close.
    66  func (w *Writer) Close() error {
    67  	if err := w.z.SetComment(w.goos + "/" + w.goarch); err != nil {
    68  		return err
    69  	}
    70  	if err := w.z.Close(); err != nil {
    71  		return err
    72  	}
    73  	_, err := writeFooter(w.w, w.off)
    74  	return err
    75  }