From c912ff669b409c4384776868ed40ca9718985662 Mon Sep 17 00:00:00 2001
From: Mike Gleason jr Couturier <mikegleasonjr@gmail.com>
Date: Fri, 29 Mar 2019 09:27:06 +0100
Subject: [PATCH] Can now specify a safe encoding which prevents segfaults in
 containers

---
 .gitignore           |  2 ++
 goheif.go            |  9 ++++++++-
 goheif_test.go       | 32 ++++++++++++++++++++++++++++++++
 libde265/libde265.go | 44 +++++++++++++++++++++++++++++++-------------
 4 files changed, 73 insertions(+), 14 deletions(-)

diff --git a/.gitignore b/.gitignore
index f1c181e..ba4f05c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,3 +10,5 @@
 
 # Output of the go coverage tool, specifically when used with LiteIDE
 *.out
+
+.vscode
diff --git a/goheif.go b/goheif.go
index 7a5d78d..c6b0011 100644
--- a/goheif.go
+++ b/goheif.go
@@ -12,6 +12,10 @@ import (
 	"github.com/jdeng/goheif/libde265"
 )
 
+// SafeEncoding uses more memory but seems to make
+// the library safer to use in containers.
+var SafeEncoding bool
+
 type gridBox struct {
 	columns, rows int
 	width, height int
@@ -100,7 +104,10 @@ func Decode(r io.Reader) (image.Image, error) {
 		return nil, fmt.Errorf("No item info")
 	}
 
-	dec := libde265.NewDecoder()
+	dec, err := libde265.NewDecoder(libde265.WithSafeEncoding(SafeEncoding))
+	if err != nil {
+		return nil, err
+	}
 	defer dec.Free()
 	if it.Info.ItemType == "hvc1" {
 		return decodeHevcItem(dec, hf, it)
diff --git a/goheif_test.go b/goheif_test.go
index 660fcc7..9ff1273 100644
--- a/goheif_test.go
+++ b/goheif_test.go
@@ -3,6 +3,7 @@ package goheif
 import (
 	"bytes"
 	"image"
+	"io"
 	"io/ioutil"
 	"testing"
 )
@@ -26,3 +27,34 @@ func TestFormatRegistered(t *testing.T) {
 		t.Errorf("unexpected decoded image size: got %dx%d, want 1596x1064", w, h)
 	}
 }
+
+func BenchmarkSafeEncoding(b *testing.B) {
+	benchEncoding(b, true)
+}
+
+func BenchmarkRegularEncoding(b *testing.B) {
+	benchEncoding(b, false)
+}
+
+func benchEncoding(b *testing.B, safe bool) {
+	b.Helper()
+
+	currentSetting := SafeEncoding
+	defer func() {
+		SafeEncoding = currentSetting
+	}()
+	SafeEncoding = safe
+
+	f, err := ioutil.ReadFile("testdata/camel.heic")
+	if err != nil {
+		b.Fatal(err)
+	}
+	r := bytes.NewReader(f)
+
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		Decode(r)
+		r.Seek(0, io.SeekStart)
+	}
+}
diff --git a/libde265/libde265.go b/libde265/libde265.go
index 1f0ae43..40f2efb 100644
--- a/libde265/libde265.go
+++ b/libde265/libde265.go
@@ -14,8 +14,9 @@ import (
 )
 
 type Decoder struct {
-	ctx      unsafe.Pointer
-	hasImage bool
+	ctx        unsafe.Pointer
+	hasImage   bool
+	safeEncode bool
 }
 
 func Init() {
@@ -26,13 +27,27 @@ func Fini() {
 	C.de265_free()
 }
 
-func NewDecoder() *Decoder {
-	if p := C.de265_new_decoder(); p != nil {
-		return &Decoder{ctx: p, hasImage: false}
+func NewDecoder(opts ...Option) (*Decoder, error) {
+	p := C.de265_new_decoder()
+	if p == nil {
+		return nil, fmt.Errorf("Unable to create decoder")
 	}
-	return nil
+
+	dec := &Decoder{ctx: p, hasImage: false}
+	for _, opt := range opts {
+		opt(dec)
+	}
+
+	return dec, nil
 }
 
+type Option func(*Decoder)
+
+func WithSafeEncoding(b bool) Option {
+	return func(dec *Decoder) {
+		dec.safeEncode = b
+	}
+}
 
 func (dec *Decoder) Free() {
 	dec.Reset()
@@ -113,7 +128,7 @@ func (dec *Decoder) DecodeImage(data []byte) (image.Image, error) {
 			//			crh := C.de265_get_image_height(img, 2)
 
 			// sanity check
-			if int(height) * int(ystride) >= int(1 << 30) {
+			if int(height)*int(ystride) >= int(1<<30) {
 				return nil, fmt.Errorf("image too big")
 			}
 
@@ -127,17 +142,20 @@ func (dec *Decoder) DecodeImage(data []byte) (image.Image, error) {
 				r = image.YCbCrSubsampleRatio444
 			}
 			ycc := &image.YCbCr{
-				Y:  (*[1 << 30]byte)(unsafe.Pointer(y))[:int(height)*int(ystride)],
-				Cb: (*[1 << 30]byte)(unsafe.Pointer(cb))[:int(cheight)*int(cstride)],
-				Cr: (*[1 << 30]byte)(unsafe.Pointer(cr))[:int(cheight)*int(cstride)],
-				//Y:              C.GoBytes(unsafe.Pointer(y), C.int(height*ystride)),
-				//Cb:             C.GoBytes(unsafe.Pointer(cb), C.int(cheight*cstride)),
-				//Cr:             C.GoBytes(unsafe.Pointer(cr), C.int(cheight*cstride)),
 				YStride:        int(ystride),
 				CStride:        int(cstride),
 				SubsampleRatio: r,
 				Rect:           image.Rectangle{Min: image.Point{0, 0}, Max: image.Point{int(width), int(height)}},
 			}
+			if dec.safeEncode {
+				ycc.Y = C.GoBytes(unsafe.Pointer(y), C.int(height*ystride))
+				ycc.Cb = C.GoBytes(unsafe.Pointer(cb), C.int(cheight*cstride))
+				ycc.Cr = C.GoBytes(unsafe.Pointer(cr), C.int(cheight*cstride))
+			} else {
+				ycc.Y = (*[1 << 30]byte)(unsafe.Pointer(y))[:int(height)*int(ystride)]
+				ycc.Cb = (*[1 << 30]byte)(unsafe.Pointer(cb))[:int(cheight)*int(cstride)]
+				ycc.Cr = (*[1 << 30]byte)(unsafe.Pointer(cr))[:int(cheight)*int(cstride)]
+			}
 
 			//C.de265_release_next_picture(dec.ctx)