From f49c17dfdee0d5401ab6fc66ad4c6c3f32754bef Mon Sep 17 00:00:00 2001
From: Juan Lang <juan.lang@gmail.com>
Date: Tue, 25 Mar 2008 11:15:08 -0700
Subject: [PATCH] Keep track of state in each message type's update function rather than in
CryptMsgUpdate
---
 dlls/crypt32/msg.c |   40 ++++++++++++++++++++++++----------------
 1 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/dlls/crypt32/msg.c b/dlls/crypt32/msg.c
index f06b8ea..5607f18 100644
--- a/dlls/crypt32/msg.c
+++ b/dlls/crypt32/msg.c
@@ -188,7 +188,9 @@ static BOOL CDataEncodeMsg_Update(HCRYPT
     CDataEncodeMsg *msg = (CDataEncodeMsg *)hCryptMsg;
     BOOL ret = FALSE;
 
-    if (msg->base.streamed)
+    if (msg->base.state == MsgStateFinalized)
+        SetLastError(CRYPT_E_MSG_ERROR);
+    else if (msg->base.streamed)
     {
         __TRY
         {
@@ -225,11 +227,15 @@ static BOOL CDataEncodeMsg_Update(HCRYPT
                 }
             }
             if (!fFinal)
+            {
                 ret = msg->base.stream_info.pfnStreamOutput(
                  msg->base.stream_info.pvArg, (BYTE *)pbData, cbData,
                  FALSE);
+                msg->base.state = MsgStateUpdated;
+            }
             else
             {
+                msg->base.state = MsgStateFinalized;
                 if (msg->base.stream_info.cbContent == 0xffffffff)
                 {
                     BYTE indefinite_trailer[6] = { 0 };
@@ -265,6 +271,7 @@ static BOOL CDataEncodeMsg_Update(HCRYPT
         }
         else
         {
+            msg->base.state = MsgStateFinalized;
             if (!cbData)
                 SetLastError(E_INVALIDARG);
             else
@@ -504,12 +511,15 @@ static BOOL CHashEncodeMsg_Update(HCRYPT
 
     TRACE("(%p, %p, %d, %d)\n", hCryptMsg, pbData, cbData, fFinal);
 
-    if (msg->base.streamed || (msg->base.open_flags & CMSG_DETACHED_FLAG))
+    if (msg->base.state == MsgStateFinalized)
+        SetLastError(CRYPT_E_MSG_ERROR);
+    else if (msg->base.streamed || (msg->base.open_flags & CMSG_DETACHED_FLAG))
     {
         /* Doesn't do much, as stream output is never called, and you
          * can't get the content.
          */
         ret = CryptHashData(msg->hash, pbData, cbData, 0);
+        msg->base.state = fFinal ? MsgStateFinalized : MsgStateUpdated;
     }
     else
     {
@@ -529,6 +539,7 @@ static BOOL CHashEncodeMsg_Update(HCRYPT
                 else
                     ret = FALSE;
             }
+            msg->base.state = MsgStateFinalized;
         }
     }
     return ret;
@@ -1183,12 +1194,15 @@ static BOOL CSignedEncodeMsg_Update(HCRY
     CSignedEncodeMsg *msg = (CSignedEncodeMsg *)hCryptMsg;
     BOOL ret = FALSE;
 
-    if (msg->base.streamed || (msg->base.open_flags & CMSG_DETACHED_FLAG))
+    if (msg->base.state == MsgStateFinalized)
+        SetLastError(CRYPT_E_MSG_ERROR);
+    else if (msg->base.streamed || (msg->base.open_flags & CMSG_DETACHED_FLAG))
     {
         ret = CSignedMsgData_Update(&msg->msg_data, pbData, cbData, fFinal,
          Sign);
         if (msg->base.streamed)
             FIXME("streamed partial stub\n");
+        msg->base.state = fFinal ? MsgStateFinalized : MsgStateUpdated;
     }
     else
     {
@@ -1211,6 +1225,7 @@ static BOOL CSignedEncodeMsg_Update(HCRY
             if (ret)
                 ret = CSignedMsgData_Update(&msg->msg_data, pbData, cbData,
                  fFinal, Sign);
+            msg->base.state = MsgStateFinalized;
         }
     }
     return ret;
@@ -1644,11 +1659,14 @@ static BOOL CDecodeMsg_Update(HCRYPTMSG 
 
     TRACE("(%p, %p, %d, %d)\n", hCryptMsg, pbData, cbData, fFinal);
 
-    if (msg->base.streamed)
+    if (msg->base.state == MsgStateFinalized)
+        SetLastError(CRYPT_E_MSG_ERROR);
+    else if (msg->base.streamed)
     {
         ret = CDecodeMsg_CopyData(msg, pbData, cbData);
         FIXME("(%p, %p, %d, %d): streamed update stub\n", hCryptMsg, pbData,
          cbData, fFinal);
+        msg->base.state = fFinal ? MsgStateFinalized : MsgStateUpdated;
     }
     else
     {
@@ -1659,7 +1677,7 @@ static BOOL CDecodeMsg_Update(HCRYPTMSG 
             ret = CDecodeMsg_CopyData(msg, pbData, cbData);
             if (ret)
                 ret = CDecodeMsg_DecodeContent(msg, &msg->msg_data, msg->type);
-
+            msg->base.state = MsgStateFinalized;
         }
     }
     return ret;
@@ -2360,20 +2378,10 @@ BOOL WINAPI CryptMsgUpdate(HCRYPTMSG hCr
  DWORD cbData, BOOL fFinal)
 {
     CryptMsgBase *msg = (CryptMsgBase *)hCryptMsg;
-    BOOL ret = FALSE;
 
     TRACE("(%p, %p, %d, %d)\n", hCryptMsg, pbData, cbData, fFinal);
 
-    if (msg->state == MsgStateFinalized)
-        SetLastError(CRYPT_E_MSG_ERROR);
-    else
-    {
-        ret = msg->update(hCryptMsg, pbData, cbData, fFinal);
-        msg->state = MsgStateUpdated;
-        if (fFinal)
-            msg->state = MsgStateFinalized;
-    }
-    return ret;
+    return msg->update(hCryptMsg, pbData, cbData, fFinal);
 }
 
 BOOL WINAPI CryptMsgGetParam(HCRYPTMSG hCryptMsg, DWORD dwParamType,
-- 
1.4.1