png.go

Documentation: github.com/dsoprea/go-png-image-structure/v2

     1  package pngstructure
     2  
     3  import (
     4  	"bytes"
     5  	"errors"
     6  	"fmt"
     7  	"io"
     8  
     9  	"encoding/binary"
    10  	"hash/crc32"
    11  
    12  	"github.com/dsoprea/go-exif/v3"
    13  	"github.com/dsoprea/go-exif/v3/common"
    14  	"github.com/dsoprea/go-logging"
    15  	"github.com/dsoprea/go-utility/v2/image"
    16  )
    17  
    18  var (
    19  	PngSignature  = [8]byte{137, 'P', 'N', 'G', '\r', '\n', 26, '\n'}
    20  	EXifChunkType = "eXIf"
    21  	IHDRChunkType = "IHDR"
    22  )
    23  
    24  var (
    25  	ErrNotPng     = errors.New("not png data")
    26  	ErrCrcFailure = errors.New("crc failure")
    27  )
    28  
    29  // ChunkSlice encapsulates a slice of chunks.
    30  type ChunkSlice struct {
    31  	chunks []*Chunk
    32  }
    33  
    34  func NewChunkSlice(chunks []*Chunk) *ChunkSlice {
    35  	if len(chunks) == 0 {
    36  		log.Panicf("ChunkSlice must be initialized with at least one chunk (IHDR)")
    37  	} else if chunks[0].Type != IHDRChunkType {
    38  		log.Panicf("first chunk in any ChunkSlice must be an IHDR")
    39  	}
    40  
    41  	return &ChunkSlice{
    42  		chunks: chunks,
    43  	}
    44  }
    45  
    46  func NewPngChunkSlice() *ChunkSlice {
    47  
    48  	ihdrChunk := &Chunk{
    49  		Type: IHDRChunkType,
    50  	}
    51  
    52  	ihdrChunk.UpdateCrc32()
    53  
    54  	return NewChunkSlice([]*Chunk{ihdrChunk})
    55  }
    56  
    57  func (cs *ChunkSlice) String() string {
    58  	return fmt.Sprintf("ChunkSlize<LEN=(%d)>", len(cs.chunks))
    59  }
    60  
    61  // Chunks exposes the actual slice.
    62  func (cs *ChunkSlice) Chunks() []*Chunk {
    63  	return cs.chunks
    64  }
    65  
    66  // Write encodes and writes all chunks.
    67  func (cs *ChunkSlice) WriteTo(w io.Writer) (err error) {
    68  	defer func() {
    69  		if state := recover(); state != nil {
    70  			err = log.Wrap(state.(error))
    71  		}
    72  	}()
    73  
    74  	_, err = w.Write(PngSignature[:])
    75  	log.PanicIf(err)
    76  
    77  	// TODO(dustin): !! This should respect the safe-to-copy characteristic.
    78  	for _, c := range cs.chunks {
    79  		_, err := c.WriteTo(w)
    80  		log.PanicIf(err)
    81  	}
    82  
    83  	return nil
    84  }
    85  
    86  // Index returns a map of chunk types to chunk slices, grouping all like chunks.
    87  func (cs *ChunkSlice) Index() (index map[string][]*Chunk) {
    88  	index = make(map[string][]*Chunk)
    89  	for _, c := range cs.chunks {
    90  		if grouped, found := index[c.Type]; found == true {
    91  			index[c.Type] = append(grouped, c)
    92  		} else {
    93  			index[c.Type] = []*Chunk{c}
    94  		}
    95  	}
    96  
    97  	return index
    98  }
    99  
   100  // FindExif returns the the segment that hosts the EXIF data.
   101  func (cs *ChunkSlice) FindExif() (chunk *Chunk, err error) {
   102  	defer func() {
   103  		if state := recover(); state != nil {
   104  			err = log.Wrap(state.(error))
   105  		}
   106  	}()
   107  
   108  	index := cs.Index()
   109  
   110  	if chunks, found := index[EXifChunkType]; found == true {
   111  		return chunks[0], nil
   112  	}
   113  
   114  	log.Panic(exif.ErrNoExif)
   115  
   116  	// Never called.
   117  	return nil, nil
   118  }
   119  
   120  // Exif returns an `exif.Ifd` instance with the existing tags.
   121  func (cs *ChunkSlice) Exif() (rootIfd *exif.Ifd, data []byte, err error) {
   122  	defer func() {
   123  		if state := recover(); state != nil {
   124  			err = log.Wrap(state.(error))
   125  		}
   126  	}()
   127  
   128  	chunk, err := cs.FindExif()
   129  	log.PanicIf(err)
   130  
   131  	im, err := exifcommon.NewIfdMappingWithStandard()
   132  	log.PanicIf(err)
   133  
   134  	ti := exif.NewTagIndex()
   135  
   136  	// TODO(dustin): Refactor and support `exif.GetExifData()`.
   137  
   138  	_, index, err := exif.Collect(im, ti, chunk.Data)
   139  	log.PanicIf(err)
   140  
   141  	return index.RootIfd, chunk.Data, nil
   142  }
   143  
   144  // ConstructExifBuilder returns an `exif.IfdBuilder` instance (needed for
   145  // modifying) preloaded with all existing tags.
   146  func (cs *ChunkSlice) ConstructExifBuilder() (rootIb *exif.IfdBuilder, err error) {
   147  	defer func() {
   148  		if state := recover(); state != nil {
   149  			err = log.Wrap(state.(error))
   150  		}
   151  	}()
   152  
   153  	rootIfd, _, err := cs.Exif()
   154  	log.PanicIf(err)
   155  
   156  	ib := exif.NewIfdBuilderFromExistingChain(rootIfd)
   157  
   158  	return ib, nil
   159  }
   160  
   161  // SetExif encodes and sets EXIF data into this segment.
   162  func (cs *ChunkSlice) SetExif(ib *exif.IfdBuilder) (err error) {
   163  	defer func() {
   164  		if state := recover(); state != nil {
   165  			err = log.Wrap(state.(error))
   166  		}
   167  	}()
   168  
   169  	// Encode.
   170  
   171  	ibe := exif.NewIfdByteEncoder()
   172  
   173  	exifData, err := ibe.EncodeToExif(ib)
   174  	log.PanicIf(err)
   175  
   176  	// Set.
   177  
   178  	exifChunk, err := cs.FindExif()
   179  	if err == nil {
   180  		// EXIF chunk already exists.
   181  
   182  		exifChunk.Data = exifData
   183  		exifChunk.Length = uint32(len(exifData))
   184  	} else {
   185  		if log.Is(err, exif.ErrNoExif) != true {
   186  			log.Panic(err)
   187  		}
   188  
   189  		// Add a EXIF chunk for the first time.
   190  
   191  		exifChunk = &Chunk{
   192  			Type:   EXifChunkType,
   193  			Data:   exifData,
   194  			Length: uint32(len(exifData)),
   195  		}
   196  
   197  		// Insert it after the IHDR chunk (it's a reliably appropriate place to
   198  		// put it).
   199  		cs.chunks = append(cs.chunks[:1], append([]*Chunk{exifChunk}, cs.chunks[1:]...)...)
   200  	}
   201  
   202  	exifChunk.UpdateCrc32()
   203  
   204  	return nil
   205  }
   206  
   207  // PngSplitter hosts the princpal `Split()` method uses by `bufio.Scanner`.
   208  type PngSplitter struct {
   209  	chunks        []*Chunk
   210  	currentOffset int
   211  
   212  	doCheckCrc bool
   213  	crcErrors  []string
   214  }
   215  
   216  func (ps *PngSplitter) Chunks() *ChunkSlice {
   217  	return NewChunkSlice(ps.chunks)
   218  }
   219  
   220  func (ps *PngSplitter) DoCheckCrc(doCheck bool) {
   221  	ps.doCheckCrc = doCheck
   222  }
   223  
   224  func (ps *PngSplitter) CrcErrors() []string {
   225  	return ps.crcErrors
   226  }
   227  
   228  func NewPngSplitter() *PngSplitter {
   229  	return &PngSplitter{
   230  		chunks:     make([]*Chunk, 0),
   231  		doCheckCrc: true,
   232  		crcErrors:  make([]string, 0),
   233  	}
   234  }
   235  
   236  // Chunk describes a single chunk.
   237  type Chunk struct {
   238  	Offset int
   239  	Length uint32
   240  	Type   string
   241  	Data   []byte
   242  	Crc    uint32
   243  }
   244  
   245  func (c *Chunk) String() string {
   246  	return fmt.Sprintf("Chunk<OFFSET=(%d) LENGTH=(%d) TYPE=[%s] CRC=(%d)>", c.Offset, c.Length, c.Type, c.Crc)
   247  }
   248  
   249  func calculateCrc32(chunk *Chunk) uint32 {
   250  	c := crc32.NewIEEE()
   251  
   252  	c.Write([]byte(chunk.Type))
   253  	c.Write(chunk.Data)
   254  
   255  	return c.Sum32()
   256  }
   257  
   258  func (c *Chunk) UpdateCrc32() {
   259  	c.Crc = calculateCrc32(c)
   260  }
   261  
   262  func (c *Chunk) CheckCrc32() bool {
   263  	expected := calculateCrc32(c)
   264  	return c.Crc == expected
   265  }
   266  
   267  // Bytes encodes and returns the bytes for this chunk.
   268  func (c *Chunk) Bytes() []byte {
   269  	defer func() {
   270  		if state := recover(); state != nil {
   271  			err := log.Wrap(state.(error))
   272  			log.Panic(err)
   273  		}
   274  	}()
   275  
   276  	if len(c.Data) != int(c.Length) {
   277  		log.Panicf("length of data not correct")
   278  	}
   279  
   280  	preallocated := make([]byte, 0, 4+4+c.Length+4)
   281  	b := bytes.NewBuffer(preallocated)
   282  
   283  	err := binary.Write(b, binary.BigEndian, c.Length)
   284  	log.PanicIf(err)
   285  
   286  	_, err = b.Write([]byte(c.Type))
   287  	log.PanicIf(err)
   288  
   289  	if c.Data != nil {
   290  		_, err = b.Write(c.Data)
   291  		log.PanicIf(err)
   292  	}
   293  
   294  	err = binary.Write(b, binary.BigEndian, c.Crc)
   295  	log.PanicIf(err)
   296  
   297  	return b.Bytes()
   298  }
   299  
   300  // Write encodes and writes the bytes for this chunk.
   301  func (c *Chunk) WriteTo(w io.Writer) (count int, err error) {
   302  	defer func() {
   303  		if state := recover(); state != nil {
   304  			err = log.Wrap(state.(error))
   305  		}
   306  	}()
   307  
   308  	if len(c.Data) != int(c.Length) {
   309  		log.Panicf("length of data not correct")
   310  	}
   311  
   312  	err = binary.Write(w, binary.BigEndian, c.Length)
   313  	log.PanicIf(err)
   314  
   315  	_, err = w.Write([]byte(c.Type))
   316  	log.PanicIf(err)
   317  
   318  	_, err = w.Write(c.Data)
   319  	log.PanicIf(err)
   320  
   321  	err = binary.Write(w, binary.BigEndian, c.Crc)
   322  	log.PanicIf(err)
   323  
   324  	return 4 + len(c.Type) + len(c.Data) + 4, nil
   325  }
   326  
   327  // readHeader verifies that the PNG header bytes appear next.
   328  func (ps *PngSplitter) readHeader(r io.Reader) (err error) {
   329  	defer func() {
   330  		if state := recover(); state != nil {
   331  			err = log.Wrap(state.(error))
   332  		}
   333  	}()
   334  
   335  	len_ := len(PngSignature)
   336  	header := make([]byte, len_)
   337  
   338  	_, err = r.Read(header)
   339  	log.PanicIf(err)
   340  
   341  	ps.currentOffset += len_
   342  
   343  	if bytes.Compare(header, PngSignature[:]) != 0 {
   344  		log.Panic(ErrNotPng)
   345  	}
   346  
   347  	return nil
   348  }
   349  
   350  // Split fulfills the `bufio.SplitFunc` function definition for
   351  // `bufio.Scanner`.
   352  func (ps *PngSplitter) Split(data []byte, atEOF bool) (advance int, token []byte, err error) {
   353  	defer func() {
   354  		if state := recover(); state != nil {
   355  			err = log.Wrap(state.(error))
   356  		}
   357  	}()
   358  
   359  	// We might have more than one chunk's worth, and, if `atEOF` is true, we
   360  	// won't be called again. We'll repeatedly try to read additional chunks,
   361  	// but, when we run out of the data we were given then we'll return the
   362  	// number of bytes fo rthe chunks we've already completely read. Then,
   363  	// we'll be called again from theend ofthose bytes, at which point we'll
   364  	// indicate that we don't yet have enough for another chunk, and we should
   365  	// be then called with more.
   366  	for {
   367  		len_ := len(data)
   368  		if len_ < 8 {
   369  			return advance, nil, nil
   370  		}
   371  
   372  		length := binary.BigEndian.Uint32(data[:4])
   373  		type_ := string(data[4:8])
   374  		chunkSize := (8 + int(length) + 4)
   375  
   376  		if len_ < chunkSize {
   377  			return advance, nil, nil
   378  		}
   379  
   380  		crcIndex := 8 + length
   381  		crc := binary.BigEndian.Uint32(data[crcIndex : crcIndex+4])
   382  
   383  		content := make([]byte, length)
   384  		copy(content, data[8:8+length])
   385  
   386  		c := &Chunk{
   387  			Length: length,
   388  			Type:   type_,
   389  			Data:   content,
   390  			Crc:    crc,
   391  			Offset: ps.currentOffset,
   392  		}
   393  
   394  		ps.chunks = append(ps.chunks, c)
   395  
   396  		if c.CheckCrc32() == false {
   397  			ps.crcErrors = append(ps.crcErrors, type_)
   398  
   399  			if ps.doCheckCrc == true {
   400  				log.Panic(ErrCrcFailure)
   401  			}
   402  		}
   403  
   404  		advance += chunkSize
   405  		ps.currentOffset += chunkSize
   406  
   407  		data = data[chunkSize:]
   408  	}
   409  
   410  	return advance, nil, nil
   411  }
   412  
   413  var (
   414  	// Enforce interface conformance.
   415  	_ riimage.MediaContext = new(ChunkSlice)
   416  )
   417
View as plain text