1 files changed, 6 insertions, 3 deletions

diff --git a/lib/basic_rp.cc b/lib/basic_rp.cc
index 763a391..2da8416 100644
--- a/lib/basic_rp.cc
+++ b/lib/basic_rp.cc
@@ -10,25 +10,28 @@
 namespace opkele {
 
     static void dh_get_secret(
             secret_t& secret, const basic_openid_message& om,
             const char *exp_assoc, const char *exp_sess,
             util::dh_t& dh,
-            size_t d_len, unsigned char *(*d_fun)(const unsigned char*,size_t,unsigned char*) ) try {
+            size_t d_len, unsigned char *(*d_fun)(const unsigned char*,size_t,unsigned char*),
+            size_t exp_s_len) try {
         if(om.get_field("assoc_type")!=exp_assoc || om.get_field("session_type")!=exp_sess)
             throw bad_input(OPKELE_CP_ "Unexpected associate response");
         util::bignum_t s_pub = util::base64_to_bignum(om.get_field("dh_server_public"));
         vector<unsigned char> ck(DH_size(dh)+1);
         unsigned char *ckptr = &(ck.front())+1;
         int cklen = DH_compute_key(ckptr,s_pub,dh);
         if(cklen<0)
             throw exception_openssl(OPKELE_CP_ "failed to DH_compute_key()");
         if(cklen && (*ckptr)&0x80) {
             (*(--ckptr))=0; ++cklen; }
         unsigned char key_digest[d_len];
         secret.enxor_from_base64((*d_fun)(ckptr,cklen,key_digest),om.get_field("enc_mac_key"));
+        if(secret.size()!=exp_s_len)
+            throw bad_input(OPKELE_CP_ "Secret length isn't consistent with association type");
     }catch(opkele::failed_lookup& ofl) {
         throw bad_input(OPKELE_CP_ "Incoherent response from OP");
     } OPKELE_RETHROW
 
     static void direct_request(basic_openid_message& oum,const basic_openid_message& inm,const string& OP) {
         util::curl_pick_t curl = util::curl_pick_t::easy_init();
@@ -70,22 +73,22 @@ namespace opkele {
         secret_t secret;
         int expires_in;
         try {
             direct_request(res,req,OP);
             dh_get_secret( secret, res,
                     "HMAC-SHA256", "DH-SHA256",
-                    dh, SHA256_DIGEST_LENGTH, SHA256 );
+                    dh, SHA256_DIGEST_LENGTH, SHA256, SHA256_DIGEST_LENGTH );
             expires_in = util::string_to_long(res.get_field("expires_in"));
         }catch(exception& e) {
             try {
                 req.set_field("assoc_type","HMAC-SHA1");
                 req.set_field("session_type","DH-SHA1");
                 direct_request(res,req,OP);
                 dh_get_secret( secret, res,
                         "HMAC-SHA1", "DH-SHA1",
-                        dh, SHA_DIGEST_LENGTH, SHA1 );
+                        dh, SHA_DIGEST_LENGTH, SHA1, SHA_DIGEST_LENGTH );
                 expires_in = util::string_to_long(res.get_field("expires_in"));
             }catch(bad_input& e) {
                 throw dumb_RP(OPKELE_CP_ "OP failed to supply an association");
             }
         }
         return store_assoc(