package pngstructure import ( "bytes" "errors" "fmt" "io" "encoding/binary" "hash/crc32" "github.com/dsoprea/go-exif/v2" "github.com/dsoprea/go-logging" "github.com/dsoprea/go-utility/image" ) var ( PngSignature = [8]byte{137, 'P', 'N', 'G', '\r', '\n', 26, '\n'} EXifChunkType = "eXIf" IHDRChunkType = "IHDR" ) var ( ErrNotPng = errors.New("not png data") ErrNoExif = errors.New("file does not have EXIF") ErrCrcFailure = errors.New("crc failure") ) // ChunkSlice encapsulates a slice of chunks. type ChunkSlice struct { chunks []*Chunk } func NewChunkSlice(chunks []*Chunk) *ChunkSlice { if len(chunks) == 0 { log.Panicf("ChunkSlice must be initialized with at least one chunk (IHDR)") } else if chunks[0].Type != IHDRChunkType { log.Panicf("first chunk in any ChunkSlice must be an IHDR") } return &ChunkSlice{ chunks: chunks, } } func NewPngChunkSlice() *ChunkSlice { ihdrChunk := &Chunk{ Type: IHDRChunkType, } ihdrChunk.UpdateCrc32() return NewChunkSlice([]*Chunk{ihdrChunk}) } func (cs *ChunkSlice) String() string { return fmt.Sprintf("ChunkSlize", len(cs.chunks)) } // Chunks exposes the actual slice. func (cs *ChunkSlice) Chunks() []*Chunk { return cs.chunks } // Write encodes and writes all chunks. func (cs *ChunkSlice) WriteTo(w io.Writer) (err error) { defer func() { if state := recover(); state != nil { err = log.Wrap(state.(error)) } }() _, err = w.Write(PngSignature[:]) log.PanicIf(err) // TODO(dustin): !! This should respect the safe-to-copy characteristic. for _, c := range cs.chunks { _, err := c.WriteTo(w) log.PanicIf(err) } return nil } // Index returns a map of chunk types to chunk slices, grouping all like chunks. func (cs *ChunkSlice) Index() (index map[string][]*Chunk) { index = make(map[string][]*Chunk) for _, c := range cs.chunks { if grouped, found := index[c.Type]; found == true { index[c.Type] = append(grouped, c) } else { index[c.Type] = []*Chunk{c} } } return index } // FindExif returns the the segment that hosts the EXIF data. func (cs *ChunkSlice) FindExif() (chunk *Chunk, err error) { defer func() { if state := recover(); state != nil { err = log.Wrap(state.(error)) } }() index := cs.Index() if chunks, found := index[EXifChunkType]; found == true { return chunks[0], nil } log.Panic(ErrNoExif) // Never called. return nil, nil } // Exif returns an `exif.Ifd` instance with the existing tags. func (cs *ChunkSlice) Exif() (rootIfd *exif.Ifd, data []byte, err error) { defer func() { if state := recover(); state != nil { err = log.Wrap(state.(error)) } }() chunk, err := cs.FindExif() log.PanicIf(err) im := exif.NewIfdMappingWithStandard() ti := exif.NewTagIndex() // TODO(dustin): Refactor and support `exif.GetExifData()`. _, index, err := exif.Collect(im, ti, chunk.Data) log.PanicIf(err) return index.RootIfd, chunk.Data, nil } // ConstructExifBuilder returns an `exif.IfdBuilder` instance (needed for // modifying) preloaded with all existing tags. func (cs *ChunkSlice) ConstructExifBuilder() (rootIb *exif.IfdBuilder, err error) { defer func() { if state := recover(); state != nil { err = log.Wrap(state.(error)) } }() rootIfd, _, err := cs.Exif() log.PanicIf(err) ib := exif.NewIfdBuilderFromExistingChain(rootIfd) return ib, nil } // SetExif encodes and sets EXIF data into this segment. func (cs *ChunkSlice) SetExif(ib *exif.IfdBuilder) (err error) { defer func() { if state := recover(); state != nil { err = log.Wrap(state.(error)) } }() // Encode. ibe := exif.NewIfdByteEncoder() exifData, err := ibe.EncodeToExif(ib) log.PanicIf(err) // Set. exifChunk, err := cs.FindExif() if err == nil { // EXIF chunk already exists. exifChunk.Data = exifData exifChunk.Length = uint32(len(exifData)) } else { if log.Is(err, ErrNoExif) != true { log.Panic(err) } // Add a EXIF chunk for the first time. exifChunk = &Chunk{ Type: EXifChunkType, Data: exifData, Length: uint32(len(exifData)), } // Insert it after the IHDR chunk (it's a reliably appropriate place to // put it). cs.chunks = append(cs.chunks[:1], append([]*Chunk{exifChunk}, cs.chunks[1:]...)...) } exifChunk.UpdateCrc32() return nil } // PngSplitter hosts the princpal `Split()` method uses by `bufio.Scanner`. type PngSplitter struct { chunks []*Chunk currentOffset int doCheckCrc bool crcErrors []string } func (ps *PngSplitter) Chunks() *ChunkSlice { return NewChunkSlice(ps.chunks) } func (ps *PngSplitter) DoCheckCrc(doCheck bool) { ps.doCheckCrc = doCheck } func (ps *PngSplitter) CrcErrors() []string { return ps.crcErrors } func NewPngSplitter() *PngSplitter { return &PngSplitter{ chunks: make([]*Chunk, 0), doCheckCrc: true, crcErrors: make([]string, 0), } } // Chunk describes a single chunk. type Chunk struct { Offset int Length uint32 Type string Data []byte Crc uint32 } func (c *Chunk) String() string { return fmt.Sprintf("Chunk", c.Offset, c.Length, c.Type, c.Crc) } func calculateCrc32(chunk *Chunk) uint32 { c := crc32.NewIEEE() c.Write([]byte(chunk.Type)) c.Write(chunk.Data) return c.Sum32() } func (c *Chunk) UpdateCrc32() { c.Crc = calculateCrc32(c) } func (c *Chunk) CheckCrc32() bool { expected := calculateCrc32(c) return c.Crc == expected } // Bytes encodes and returns the bytes for this chunk. func (c *Chunk) Bytes() []byte { defer func() { if state := recover(); state != nil { err := log.Wrap(state.(error)) log.Panic(err) } }() if len(c.Data) != int(c.Length) { log.Panicf("length of data not correct") } preallocated := make([]byte, 0, 4+4+c.Length+4) b := bytes.NewBuffer(preallocated) err := binary.Write(b, binary.BigEndian, c.Length) log.PanicIf(err) _, err = b.Write([]byte(c.Type)) log.PanicIf(err) if c.Data != nil { _, err = b.Write(c.Data) log.PanicIf(err) } err = binary.Write(b, binary.BigEndian, c.Crc) log.PanicIf(err) return b.Bytes() } // Write encodes and writes the bytes for this chunk. func (c *Chunk) WriteTo(w io.Writer) (count int, err error) { defer func() { if state := recover(); state != nil { err = log.Wrap(state.(error)) } }() if len(c.Data) != int(c.Length) { log.Panicf("length of data not correct") } err = binary.Write(w, binary.BigEndian, c.Length) log.PanicIf(err) _, err = w.Write([]byte(c.Type)) log.PanicIf(err) _, err = w.Write(c.Data) log.PanicIf(err) err = binary.Write(w, binary.BigEndian, c.Crc) log.PanicIf(err) return 4 + len(c.Type) + len(c.Data) + 4, nil } // readHeader verifies that the PNG header bytes appear next. func (ps *PngSplitter) readHeader(r io.Reader) (err error) { defer func() { if state := recover(); state != nil { err = log.Wrap(state.(error)) } }() len_ := len(PngSignature) header := make([]byte, len_) _, err = r.Read(header) log.PanicIf(err) ps.currentOffset += len_ if bytes.Compare(header, PngSignature[:]) != 0 { log.Panic(ErrNotPng) } return nil } // Split fulfills the `bufio.SplitFunc` function definition for // `bufio.Scanner`. func (ps *PngSplitter) Split(data []byte, atEOF bool) (advance int, token []byte, err error) { defer func() { if state := recover(); state != nil { err = log.Wrap(state.(error)) } }() // We might have more than one chunk's worth, and, if `atEOF` is true, we // won't be called again. We'll repeatedly try to read additional chunks, // but, when we run out of the data we were given then we'll return the // number of bytes fo rthe chunks we've already completely read. Then, // we'll be called again from theend ofthose bytes, at which point we'll // indicate that we don't yet have enough for another chunk, and we should // be then called with more. for { len_ := len(data) if len_ < 8 { return advance, nil, nil } length := binary.BigEndian.Uint32(data[:4]) type_ := string(data[4:8]) chunkSize := (8 + int(length) + 4) if len_ < chunkSize { return advance, nil, nil } crcIndex := 8 + length crc := binary.BigEndian.Uint32(data[crcIndex : crcIndex+4]) content := make([]byte, length) copy(content, data[8:8+length]) c := &Chunk{ Length: length, Type: type_, Data: content, Crc: crc, Offset: ps.currentOffset, } ps.chunks = append(ps.chunks, c) if c.CheckCrc32() == false { ps.crcErrors = append(ps.crcErrors, type_) if ps.doCheckCrc == true { log.Panic(ErrCrcFailure) } } advance += chunkSize ps.currentOffset += chunkSize data = data[chunkSize:] } return advance, nil, nil } var ( // Enforce interface conformance. _ riimage.MediaContext = new(ChunkSlice) )