github.com/coreos/goproxy@v0.0.0-20190513173959-f8dc2d7ba04e/ext/image/image.go (about)

     1  package goproxy_image
     2  
     3  import (
     4  	"bytes"
     5  	"image"
     6  	_ "image/gif"
     7  	"image/jpeg"
     8  	"image/png"
     9  	"io/ioutil"
    10  	"net/http"
    11  	. "github.com/elazarl/goproxy"
    12  	"github.com/elazarl/goproxy/regretable"
    13  )
    14  
    15  var RespIsImage = ContentTypeIs("image/gif",
    16  	"image/jpeg",
    17  	"image/pjpeg",
    18  	"application/octet-stream",
    19  	"image/png")
    20  
    21  // "image/tiff" tiff support is in external package, and rarely used, so we omitted it
    22  
    23  func HandleImage(f func(img image.Image, ctx *ProxyCtx) image.Image) RespHandler {
    24  	return FuncRespHandler(func(resp *http.Response, ctx *ProxyCtx) *http.Response {
    25  		if !RespIsImage.HandleResp(resp, ctx) {
    26  			return resp
    27  		}
    28  		if resp.StatusCode != 200 {
    29  			// we might get 304 - not modified response without data
    30  			return resp
    31  		}
    32  		contentType := resp.Header.Get("Content-Type")
    33  
    34  		const kb = 1024
    35  		regret := regretable.NewRegretableReaderCloserSize(resp.Body, 16*kb)
    36  		resp.Body = regret
    37  		img, imgType, err := image.Decode(resp.Body)
    38  		if err != nil {
    39  			regret.Regret()
    40  			ctx.Warnf("%s: %s", ctx.Req.Method+" "+ctx.Req.URL.String()+" Image from "+ctx.Req.RequestURI+"content type"+
    41  				contentType+"cannot be decoded returning original image", err)
    42  			return resp
    43  		}
    44  		result := f(img, ctx)
    45  		buf := bytes.NewBuffer([]byte{})
    46  		switch contentType {
    47  		// No gif image encoder in go - convert to png
    48  		case "image/gif", "image/png":
    49  			if err := png.Encode(buf, result); err != nil {
    50  				ctx.Warnf("Cannot encode image, returning orig %v %v", ctx.Req.URL.String(), err)
    51  				return resp
    52  			}
    53  			resp.Header.Set("Content-Type", "image/png")
    54  		case "image/jpeg", "image/pjpeg":
    55  			if err := jpeg.Encode(buf, result, nil); err != nil {
    56  				ctx.Warnf("Cannot encode image, returning orig %v %v", ctx.Req.URL.String(), err)
    57  				return resp
    58  			}
    59  		case "application/octet-stream":
    60  			switch imgType {
    61  			case "jpeg":
    62  				if err := jpeg.Encode(buf, result, nil); err != nil {
    63  					ctx.Warnf("Cannot encode image as jpeg, returning orig %v %v", ctx.Req.URL.String(), err)
    64  					return resp
    65  				}
    66  			case "png", "gif":
    67  				if err := png.Encode(buf, result); err != nil {
    68  					ctx.Warnf("Cannot encode image as png, returning orig %v %v", ctx.Req.URL.String(), err)
    69  					return resp
    70  				}
    71  			}
    72  		default:
    73  			panic("unhandlable type" + contentType)
    74  		}
    75  		resp.Body = ioutil.NopCloser(buf)
    76  		return resp
    77  	})
    78  }