s3:libsmb: add a smb_signing_init_ex() function

Make it possible to overload memory handling functions.

metze

2 files changed, 56 insertions, 8 deletions

diff --git a/source3/include/smb_signing.h b/source3/include/smb_signing.h
index 770c40cb35..d2eda9b3e6 100644
--- a/source3/include/smb_signing.h
+++ b/source3/include/smb_signing.h
@@ -27,6 +27,11 @@ struct smb_signing_state;
 struct smb_signing_state *smb_signing_init(TALLOC_CTX *mem_ctx,
 					   bool allowed,
 					   bool mandatory);
+struct smb_signing_state *smb_signing_init_ex(TALLOC_CTX *mem_ctx,
+					      bool allowed,
+					      bool mandatory,
+					      void *(*alloc_fn)(TALLOC_CTX *, size_t),
+					      void (*free_fn)(TALLOC_CTX *, void *));
 uint32_t smb_signing_next_seqnum(struct smb_signing_state *si, bool oneway);
 void smb_signing_cancel_reply(struct smb_signing_state *si, bool oneway);
 void smb_signing_sign_pdu(struct smb_signing_state *si,
diff --git a/source3/libsmb/smb_signing.c b/source3/libsmb/smb_signing.c
index 32d2883965..104cf76160 100644
--- a/source3/libsmb/smb_signing.c
+++ b/source3/libsmb/smb_signing.c
@@ -43,25 +43,50 @@ struct smb_signing_state {
 
 	/* the next expected seqnum */
 	uint32_t seqnum;
+
+	TALLOC_CTX *mem_ctx;
+	void *(*alloc_fn)(TALLOC_CTX *mem_ctx, size_t len);
+	void (*free_fn)(TALLOC_CTX *mem_ctx, void *ptr);
 };
 
 static void smb_signing_reset_info(struct smb_signing_state *si)
 {
 	si->active = false;
 	si->bsrspyl = false;
-	data_blob_free(&si->mac_key);
 	si->seqnum = 0;
+
+	if (si->free_fn) {
+		si->free_fn(si->mem_ctx, si->mac_key.data);
+	} else {
+		talloc_free(si->mac_key.data);
+	}
+	si->mac_key.data = NULL;
+	si->mac_key.length = 0;
 }
 
-struct smb_signing_state *smb_signing_init(TALLOC_CTX *mem_ctx,
-					   bool allowed,
-					   bool mandatory)
+struct smb_signing_state *smb_signing_init_ex(TALLOC_CTX *mem_ctx,
+					      bool allowed,
+					      bool mandatory,
+					      void *(*alloc_fn)(TALLOC_CTX *, size_t),
+					      void (*free_fn)(TALLOC_CTX *, void *))
 {
 	struct smb_signing_state *si;
 
-	si = talloc_zero(mem_ctx, struct smb_signing_state);
-	if (si == NULL) {
-		return NULL;
+	if (alloc_fn) {
+		void *p = alloc_fn(mem_ctx, sizeof(struct smb_signing_state));
+		if (p == NULL) {
+			return NULL;
+		}
+		memset(p, 0, sizeof(struct smb_signing_state));
+		si = (struct smb_signing_state *)p;
+		si->mem_ctx = mem_ctx;
+		si->alloc_fn = alloc_fn;
+		si->free_fn = free_fn;
+	} else {
+		si = talloc_zero(mem_ctx, struct smb_signing_state);
+		if (si == NULL) {
+			return NULL;
+		}
 	}
 
 	if (mandatory) {
@@ -74,6 +99,13 @@ struct smb_signing_state *smb_signing_init(TALLOC_CTX *mem_ctx,
 	return si;
 }
 
+struct smb_signing_state *smb_signing_init(TALLOC_CTX *mem_ctx,
+					   bool allowed,
+					   bool mandatory)
+{
+	return smb_signing_init_ex(mem_ctx, allowed, mandatory, NULL, NULL);
+}
+
 static bool smb_signing_good(struct smb_signing_state *si,
 			     bool good, uint32_t seq)
 {
@@ -312,7 +344,18 @@ bool smb_signing_activate(struct smb_signing_state *si,
 	smb_signing_reset_info(si);
 
 	len = response.length + user_session_key.length;
-	si->mac_key = data_blob_talloc(si, NULL, len);
+	if (si->alloc_fn) {
+		si->mac_key.data = (uint8_t *)si->alloc_fn(si->mem_ctx, len);
+		if (si->mac_key.data == NULL) {
+			return false;
+		}
+	} else {
+		si->mac_key.data = (uint8_t *)talloc_size(si, len);
+		if (si->mac_key.data == NULL) {
+			return false;
+		}
+	}
+	si->mac_key.length = len;
 
 	ofs = 0;
 	memcpy(&si->mac_key.data[ofs], user_session_key.data, user_session_key.length);