From 6a57e8c3fcd6c84dcfb542bcc2027685c3eed184 Mon Sep 17 00:00:00 2001
From: "Dolf Schimmel (Freeaqingme)" <dolf@dolfschimmel.nl>
Date: Tue, 24 Nov 2015 17:39:31 +0100
Subject: [PATCH] Make gobencode() optional in SetPriv(), return errors instead
 of integers

---
 gomilter.go | 53 +++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 39 insertions(+), 14 deletions(-)

diff --git a/gomilter.go b/gomilter.go
index 47a9e32..2061119 100644
--- a/gomilter.go
+++ b/gomilter.go
@@ -28,6 +28,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"encoding/gob"
+	"errors"
 	"fmt"
 	"net"
 	"os"
@@ -349,14 +350,24 @@ func GetSymVal(ctx uintptr, symname string) string {
 	return C.GoString(cval)
 }
 
-// See also: http://bit.ly/1HVWA9I
-func SetPriv(ctx uintptr, privatedata interface{}) int {
+// Beware that if your struct is too large it will be discarded
+// without storing. You may want to simply provide a reference
+// to something you keep in a map or similar structure yourself.
+func SetPriv(ctx uintptr, privatedata interface{}) error {
 	// privatedata seems to work for any data type
 	// Structs must have exported fields
 
 	// Serialize Go privatedata into a byte slice
-	bytedata, _ := GobEncode(privatedata)
+	bytedata, err := GobEncode(privatedata)
+	if err != nil {
+		return err
+	}
 
+	return SetPrivBytes(ctx, bytedata)
+}
+
+// See also: http://bit.ly/1HVWA9I
+func SetPrivBytes(ctx uintptr, bytedata []byte) error {
 	// length and size
 	// length is a uint32 (usually 4 bytes)
 	// the length will be stored in front of the byte sequence
@@ -365,7 +376,7 @@ func SetPriv(ctx uintptr, privatedata interface{}) int {
 	buf := new(bytes.Buffer)
 	err := binary.Write(buf, binary.BigEndian, length)
 	if err != nil {
-		return -1
+		return errors.New("Could not write binary data into buffer: " + err.Error())
 	}
 	lengthbytes := buf.Bytes()
 
@@ -391,10 +402,29 @@ func SetPriv(ctx uintptr, privatedata interface{}) int {
 
 	// Call libmilter smfi_setpriv
 	type CtxPtr *C.struct_smfi_str
-	return int(C.smfi_setpriv(int2ctx(ctx), unsafe.Pointer(lenStart)))
+	res := int(C.smfi_setpriv(int2ctx(ctx), unsafe.Pointer(lenStart)))
+	if res != int(C.MI_SUCCESS) {
+		return errors.New("smfi_setpriv() returned a failure")
+	}
+
+	return nil
 }
 
-func GetPriv(ctx uintptr, privatedata interface{}) int {
+func GetPriv(ctx uintptr, privatedata interface{}) error {
+	databytes, err := GetPrivBytes(ctx)
+	if err != nil {
+		return err
+	}
+
+	err = GobDecode(databytes, privatedata)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func GetPrivBytes(ctx uintptr) ([]byte, error) {
 	/*  Retrieve the private data stored by the milter
 	    Retrieving the data will release the memory allocated for it
 	    Don't try to retrieve it again unless you call SetPriv first
@@ -405,7 +435,7 @@ func GetPriv(ctx uintptr, privatedata interface{}) int {
 
 	// Make sure data has been set with a previous call to SetPriv
 	if CArray == nil {
-		return -1
+		return nil, errors.New("smfi_getpriv() call failed")
 	}
 
 	// Read uint32 size bytes from the start of the pointer
@@ -425,7 +455,7 @@ func GetPriv(ctx uintptr, privatedata interface{}) int {
 	buf := bytes.NewBuffer(lengthbytes)
 	err := binary.Read(buf, binary.BigEndian, &length)
 	if err != nil {
-		return -1
+		return nil, errors.New("Could not parse binary data")
 	}
 
 	// Read byte sequence of data
@@ -440,12 +470,7 @@ func GetPriv(ctx uintptr, privatedata interface{}) int {
 	C.smfi_setpriv(int2ctx(ctx), nil)
 
 	// Unserialize the data bytes back into a data structure
-	err = GobDecode(databytes, privatedata)
-	if err != nil {
-		return -1
-	}
-
-	return 0
+	return databytes, nil
 }
 
 func SetReply(ctx uintptr, rcode, xcode, message string) int {