1 files changed, 44 insertions, 23 deletions

diff --git a/ckcapi.c b/ckcapi.c
index b0967a6..e5ce629 100644
--- a/ckcapi.c
+++ b/ckcapi.c
@@ -173,7 +173,7 @@ ckcapi_return_data(CK_ATTRIBUTE_PTR attr, CK_VOID_PTR input, DWORD n_input)
 CK_RV
 ckcapi_return_string(CK_ATTRIBUTE_PTR attr, WCHAR* string)
 {
-	DWORD error;
+	CK_UTF8CHAR_PTR buffer;
 	int result;
 
 	SetLastError(0);
@@ -184,39 +184,60 @@ ckcapi_return_string(CK_ATTRIBUTE_PTR attr, WCHAR* string)
 	 * case this part.
 	 */
 	if(!string[0])
-		return ckcapi_return_data(attr, "\0", 1);
+		return ckcapi_return_data(attr, "", 0);
 
+	/* The length of the string, including null termination */
 	result = WideCharToMultiByte(CP_UTF8, 0, string, -1, 
-	                             attr->pValue, attr->pValue ? attr->ulValueLen : 0, 
-	                             NULL, NULL);
+	                             NULL, 0, NULL, NULL);
 
-	/* An error result somehow */
-	if(!result)
+	if(result)
 	{
-		error = GetLastError();
-		switch(error)
+		/* Did the caller just want the length? */
+		if(!attr->pValue)
 		{
-		/* If buffer was too short, calculate good buffer length */
-		case ERROR_INSUFFICIENT_BUFFER:
-			result = WideCharToMultiByte(CP_UTF8, 0, string, -1, 
-				                         NULL, 0, NULL, NULL);
-			if(!result)
-				return CKR_GENERAL_ERROR;
-			attr->ulValueLen = result;
+			attr->ulValueLen = result - 1;
+			return CKR_OK;
+		}
+
+		/* Is the callers buffer too short? */
+		if((int)attr->ulValueLen < result - 1)
+		{
+			attr->ulValueLen = result - 1;
 			return CKR_BUFFER_TOO_SMALL;
+		}
+
+		/* 
+		 * Allocate a buffer for the conversion. We have to 
+		 * do this because strings in PKCS#11 are not null
+		 * terminated and strings returned from 
+		 * WideCharToMultiByte are always null terminated.
+		 */
+		buffer = malloc(result);
+		if(!buffer)
+			return CKR_HOST_MEMORY;
 
-		/* A strange zero length success */
-		case ERROR_SUCCESS:
-			break;
+		/* Do the actual conversion */
+		result = WideCharToMultiByte(CP_UTF8, 0, string, -1, 
+		                             buffer, result, NULL, NULL);
 
-		/* All other errors are not handleable */
-		default:
-			return CKR_GENERAL_ERROR;
+		if(result)
+		{
+			attr->ulValueLen = result - 1;
+			memcpy(attr->pValue, buffer, attr->ulValueLen);
+
+			free(buffer);
+			return CKR_OK;
 		}
+
+		free(buffer);
 	}
 
-	attr->ulValueLen = result;
-	return CKR_OK;
+	/*
+	 * We should never have too little buffer, or 
+	 * get a zero length success code. It's a very 
+	 * strange error that arrives here.
+	 */
+	return CKR_GENERAL_ERROR;
 }
 
 CK_RV