gtsocial-umbx

Unnamed repository; edit this file 'description' to name the repository.
Log | Files | Refs | README | LICENSE

terminator.go (4098B)

      1 /*
      2    exif-terminator
      3    Copyright (C) 2022 SuperSeriousBusiness admin@gotosocial.org
      4 
      5    This program is free software: you can redistribute it and/or modify
      6    it under the terms of the GNU Affero General Public License as published by
      7    the Free Software Foundation, either version 3 of the License, or
      8    (at your option) any later version.
      9 
     10    This program is distributed in the hope that it will be useful,
     11    but WITHOUT ANY WARRANTY; without even the implied warranty of
     12    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
     13    GNU Affero General Public License for more details.
     14 
     15    You should have received a copy of the GNU Affero General Public License
     16    along with this program.  If not, see <http://www.gnu.org/licenses/>.
     17 */
     18 
     19 package terminator
     20 
     21 import (
     22 	"bufio"
     23 	"bytes"
     24 	"errors"
     25 	"fmt"
     26 	"io"
     27 
     28 	pngstructure "github.com/dsoprea/go-png-image-structure/v2"
     29 	jpegstructure "github.com/superseriousbusiness/go-jpeg-image-structure/v2"
     30 )
     31 
     32 func Terminate(in io.Reader, fileSize int, mediaType string) (io.Reader, error) {
     33 	// to avoid keeping too much stuff in memory we want to pipe data directly
     34 	pipeReader, pipeWriter := io.Pipe()
     35 
     36 	// we don't know ahead of time how long segments might be: they could be as large as
     37 	// the file itself, so unfortunately we need to allocate a buffer here that'scanner as large
     38 	// as the file
     39 	scanner := bufio.NewScanner(in)
     40 	scanner.Buffer([]byte{}, fileSize)
     41 	var err error
     42 
     43 	switch mediaType {
     44 	case "image/jpeg", "jpeg", "jpg":
     45 		err = terminateJpeg(scanner, pipeWriter, fileSize)
     46 	case "image/webp", "webp":
     47 		err = terminateWebp(scanner, pipeWriter)
     48 	case "image/png", "png":
     49 		// for pngs we need to skip the header bytes, so read them in
     50 		// and check we're really dealing with a png here
     51 		header := make([]byte, len(pngstructure.PngSignature))
     52 		if _, headerError := in.Read(header); headerError != nil {
     53 			err = headerError
     54 			break
     55 		}
     56 
     57 		if !bytes.Equal(header, pngstructure.PngSignature[:]) {
     58 			err = errors.New("could not decode png: invalid header")
     59 			break
     60 		}
     61 
     62 		err = terminatePng(scanner, pipeWriter)
     63 	default:
     64 		err = fmt.Errorf("mediaType %s cannot be processed", mediaType)
     65 	}
     66 
     67 	return pipeReader, err
     68 }
     69 
     70 func terminateJpeg(scanner *bufio.Scanner, writer io.WriteCloser, expectedFileSize int) error {
     71 	// jpeg visitor is where the spicy hack of streaming the de-exifed data is contained
     72 	v := &jpegVisitor{
     73 		writer:           writer,
     74 		expectedFileSize: expectedFileSize,
     75 	}
     76 
     77 	// provide the visitor to the splitter so that it triggers on every section scan
     78 	js := jpegstructure.NewJpegSplitter(v)
     79 
     80 	// the visitor also needs to read back the list of segments: for this it needs
     81 	// to know what jpeg splitter it's attached to, so give it a pointer to the splitter
     82 	v.js = js
     83 
     84 	// use the jpeg splitters 'split' function, which satisfies the bufio.SplitFunc interface
     85 	scanner.Split(js.Split)
     86 
     87 	scanAndClose(scanner, writer)
     88 	return nil
     89 }
     90 
     91 func terminateWebp(scanner *bufio.Scanner, writer io.WriteCloser) error {
     92 	v := &webpVisitor{
     93 		writer: writer,
     94 	}
     95 
     96 	// use the webp visitor's 'split' function, which satisfies the bufio.SplitFunc interface
     97 	scanner.Split(v.split)
     98 
     99 	scanAndClose(scanner, writer)
    100 	return nil
    101 }
    102 
    103 func terminatePng(scanner *bufio.Scanner, writer io.WriteCloser) error {
    104 	ps := pngstructure.NewPngSplitter()
    105 
    106 	v := &pngVisitor{
    107 		ps:               ps,
    108 		writer:           writer,
    109 		lastWrittenChunk: -1,
    110 	}
    111 
    112 	// use the png visitor's 'split' function, which satisfies the bufio.SplitFunc interface
    113 	scanner.Split(v.split)
    114 
    115 	scanAndClose(scanner, writer)
    116 	return nil
    117 }
    118 
    119 func scanAndClose(scanner *bufio.Scanner, writer io.WriteCloser) {
    120 	// scan asynchronously until there's nothing left to scan, and then close the writer
    121 	// so that the reader on the other side knows that we're done
    122 	//
    123 	// due to the nature of io.Pipe, writing won't actually work
    124 	// until the pipeReader starts being read by the caller, which
    125 	// is why we do this asynchronously
    126 	go func() {
    127 		defer writer.Close()
    128 		for scanner.Scan() {
    129 		}
    130 		if scanner.Err() != nil {
    131 			logger.Error(scanner.Err())
    132 		}
    133 	}()
    134 }