diff --git a/src/headers/tomcrypt_pk.h b/src/headers/tomcrypt_pk.h
index 89e9611..00a6c13 100644
--- a/src/headers/tomcrypt_pk.h
+++ b/src/headers/tomcrypt_pk.h
@@ -224,9 +224,7 @@ int dh_set_pg(const unsigned char *p, unsigned long plen,
 int dh_set_pg_dhparam(const unsigned char *dhparam, unsigned long dhparamlen, dh_key *key);
 int dh_set_pg_groupsize(int groupsize, dh_key *key);
 
-int dh_set_key(const unsigned char *pub, unsigned long publen,
-               const unsigned char *priv, unsigned long privlen,
-               dh_key *key);
+int dh_set_key(const unsigned char *in, unsigned long inlen, int type, dh_key *key);
 int dh_generate_key(prng_state *prng, int wprng, dh_key *key);
 
 int dh_shared_secret(dh_key        *private_key, dh_key        *public_key,
@@ -234,8 +232,7 @@ int dh_shared_secret(dh_key        *private_key, dh_key        *public_key,
 
 void dh_free(dh_key *key);
 
-int dh_export_key(void *out, unsigned long *outlen,
-                  int type, dh_key *key);
+int dh_export_key(void *out, unsigned long *outlen, int type, dh_key *key);
 
 #ifdef LTC_SOURCE
 /* internal helper functions */
@@ -449,9 +446,7 @@ int dsa_set_pqg(const unsigned char *p,  unsigned long plen,
 int dsa_set_pqg_dsaparam(const unsigned char *dsaparam, unsigned long dsaparamlen, dsa_key *key);
 int dsa_generate_pqg(prng_state *prng, int wprng, int group_size, int modulus_size, dsa_key *key);
 
-int dsa_set_key(const unsigned char *pub, unsigned long publen,
-                const unsigned char *priv, unsigned long privlen,
-                dsa_key *key);
+int dsa_set_key(const unsigned char *in, unsigned long inlen, int type, dsa_key *key);
 int dsa_generate_key(prng_state *prng, int wprng, dsa_key *key);
 
 void dsa_free(dsa_key *key);
diff --git a/src/pk/dh/dh_set.c b/src/pk/dh/dh_set.c
index edd28bb..8d0af7d 100644
--- a/src/pk/dh/dh_set.c
+++ b/src/pk/dh/dh_set.c
@@ -78,40 +78,31 @@ LBL_ERR:
 }
 
 /**
-  Import DH key parts pub and priv from raw numbers
+  Import DH public or private key part from raw numbers
 
-  @param pub     DH's pub (public key) (can be NULL if priv is valid)
-  @param publen  DH's pub's length
-  @param priv    DH's priv (private key) (can be NULL if pub is valid)
-  @param privlen DH's priv's length
+     NB: The p & g parts must be set beforehand
+
+  @param in      The key-part to import, either public or private.
+  @param inlen   The key-part's length
+  @param type    Which type of key (PK_PRIVATE or PK_PUBLIC)
   @param key     [out] the destination for the imported key
   @return CRYPT_OK if successful
 */
-int dh_set_key(const unsigned char *pub, unsigned long publen,
-               const unsigned char *priv, unsigned long privlen,
-               dh_key *key)
+int dh_set_key(const unsigned char *in, unsigned long inlen, int type, dh_key *key)
 {
    int err;
 
    LTC_ARGCHK(key         != NULL);
    LTC_ARGCHK(ltc_mp.name != NULL);
 
-   if(priv == NULL) {
-      if ((err = mp_read_unsigned_bin(key->y, (unsigned char*)pub, publen)) != CRYPT_OK)         { goto LBL_ERR; }
-      key->type = PK_PUBLIC;
-      mp_clear(key->x);
-      key->x = NULL;
+   if (type == PK_PRIVATE) {
+      key->type = PK_PRIVATE;
+      if ((err = mp_read_unsigned_bin(key->x, (unsigned char*)in, inlen)) != CRYPT_OK) { goto LBL_ERR; }
+      if ((err = mp_exptmod(key->base, key->x, key->prime, key->y)) != CRYPT_OK)       { goto LBL_ERR; }
    }
    else {
-      if ((err = mp_read_unsigned_bin(key->x, (unsigned char*)priv, privlen)) != CRYPT_OK)         { goto LBL_ERR; }
-      if (pub != NULL) {
-         if ((err = mp_read_unsigned_bin(key->y, (unsigned char*)pub, publen)) != CRYPT_OK)         { goto LBL_ERR; }
-      }
-      else {
-         /* compute y value */
-         if ((err = mp_exptmod(key->base, key->x, key->prime, key->y)) != CRYPT_OK)  { goto LBL_ERR; }
-      }
-      key->type = PK_PRIVATE;
+      key->type = PK_PUBLIC;
+      if ((err = mp_read_unsigned_bin(key->y, (unsigned char*)in, inlen)) != CRYPT_OK) { goto LBL_ERR; }
    }
 
    /* check public key */
diff --git a/src/pk/dsa/dsa_set.c b/src/pk/dsa/dsa_set.c
index d25de01..a630974 100644
--- a/src/pk/dsa/dsa_set.c
+++ b/src/pk/dsa/dsa_set.c
@@ -58,17 +58,17 @@ LBL_ERR:
 }
 
 /**
-  Import DSA public or private key from raw numbers
-  @param pub     DSA's y (public key) in binary representation
-  @param publen  The length of pub
-  @param priv    DSA's x (private key) in binary representation (can be NULL when importing public key)
-  @param privlen The length of priv
+  Import DSA public or private key-part from raw numbers
+
+     NB: The p, q & g parts must be set beforehand
+
+  @param in      The key-part to import, either public or private.
+  @param inlen   The key-part's length
+  @param type    Which type of key (PK_PRIVATE or PK_PUBLIC)
   @param key     [out] the destination for the imported key
   @return CRYPT_OK if successful.
 */
-int dsa_set_key(const unsigned char *pub, unsigned long publen,
-                const unsigned char *priv, unsigned long privlen,
-                dsa_key *key)
+int dsa_set_key(const unsigned char *in, unsigned long inlen, int type, dsa_key *key)
 {
    int err;
 
@@ -80,13 +80,14 @@ int dsa_set_key(const unsigned char *pub, unsigned long publen,
    LTC_ARGCHK(key->q      != NULL);
    LTC_ARGCHK(ltc_mp.name != NULL);
 
-   if ((err = mp_read_unsigned_bin(key->y, (unsigned char *)pub , publen)) != CRYPT_OK) { goto LBL_ERR; }
-   if (priv != NULL) {
+   if (type == PK_PRIVATE) {
       key->type = PK_PRIVATE;
-      if ((err = mp_read_unsigned_bin(key->x, (unsigned char *)priv , privlen)) != CRYPT_OK) { goto LBL_ERR; }
+      if ((err = mp_read_unsigned_bin(key->x, (unsigned char *)in, inlen)) != CRYPT_OK) { goto LBL_ERR; }
+      if ((err = mp_exptmod(key->g, key->x, key->p, key->y)) != CRYPT_OK)               { goto LBL_ERR; }
    }
    else {
       key->type = PK_PUBLIC;
+      if ((err = mp_read_unsigned_bin(key->y, (unsigned char *)in, inlen)) != CRYPT_OK) { goto LBL_ERR; }
    }
 
    return CRYPT_OK;
diff --git a/tests/dh_test.c b/tests/dh_test.c
index e3111c6..4ef15a9 100644
--- a/tests/dh_test.c
+++ b/tests/dh_test.c
@@ -268,7 +268,7 @@ static int _set_test(void)
 
    for (i = 0; i < 1; i++) {
       DO(dh_set_pg(test[i].p, test[i].plen, test[i].g, test[i].glen, &k1));
-      DO(dh_set_key(NULL, 0, test[i].x, test[i].xlen, &k1));
+      DO(dh_set_key(test[i].x, test[i].xlen, PK_PRIVATE, &k1));
 
       len = sizeof(buf);
       DO(dh_export(buf, &len, PK_PRIVATE, &k1));
@@ -301,7 +301,7 @@ static int _set_test(void)
       dh_free(&k1);
 
       DO(dh_set_pg(test[i].p, test[i].plen, test[i].g, test[i].glen, &k1));
-      DO(dh_set_key(test[i].y, test[i].ylen, test[i].x, test[i].xlen, &k1));
+      DO(dh_set_key(test[i].x, test[i].xlen, PK_PRIVATE, &k1));
 
       len = sizeof(buf);
       DO(dh_export(buf, &len, PK_PRIVATE, &k1));
@@ -320,7 +320,7 @@ static int _set_test(void)
       dh_free(&k1);
 
       DO(dh_set_pg(test[i].p, test[i].plen, test[i].g, test[i].glen, &k2));
-      DO(dh_set_key(test[i].y, test[i].ylen, NULL, 0, &k2));
+      DO(dh_set_key(test[i].y, test[i].ylen, PK_PUBLIC, &k2));
 
       len = sizeof(buf);
       DO(dh_export(buf, &len, PK_PUBLIC, &k2));
diff --git a/tests/dsa_test.c b/tests/dsa_test.c
index d871a29..f1d8ecf 100644
--- a/tests/dsa_test.c
+++ b/tests/dsa_test.c
@@ -179,8 +179,8 @@ static int _dsa_compat_test(void)
                  key_parts[1], key_lens[1],
                  key_parts[2], key_lens[2],
                  &key));
-  DO(dsa_set_key(key_parts[3], key_lens[3],
-                 key_parts[4], key_lens[4],
+  DO(dsa_set_key(key_parts[4], key_lens[4],
+                 PK_PRIVATE,
                  &key));
   len = sizeof(buf);
   DO(dsa_export(buf, &len, PK_PRIVATE | PK_STD, &key));
@@ -196,7 +196,7 @@ static int _dsa_compat_test(void)
                  key_parts[2], key_lens[2],
                  &key));
   DO(dsa_set_key(key_parts[3], key_lens[3],
-                 NULL, 0,
+                 PK_PUBLIC,
                  &key));
   len = sizeof(buf);
   DO(dsa_export(buf, &len, PK_PUBLIC | PK_STD, &key));
@@ -220,7 +220,7 @@ static int _dsa_compat_test(void)
   /* try import dsaparam - our public key */
   DO(dsa_set_pqg_dsaparam(dsaparam_der, sizeof(dsaparam_der), &key));
   DO(dsa_set_key(key_parts[3], key_lens[3],
-                 NULL, 0,
+                 PK_PUBLIC,
                  &key));
   len = sizeof(buf);
   DO(dsa_export(buf, &len, PK_PUBLIC | PK_STD, &key));
@@ -232,8 +232,8 @@ static int _dsa_compat_test(void)
 
   /* try import dsaparam - our private key */
   DO(dsa_set_pqg_dsaparam(dsaparam_der, sizeof(dsaparam_der), &key));
-  DO(dsa_set_key(key_parts[3], key_lens[3],
-                 key_parts[4], key_lens[4],
+  DO(dsa_set_key(key_parts[4], key_lens[4],
+                 PK_PRIVATE,
                  &key));
   len = sizeof(buf);
   DO(dsa_export(buf, &len, PK_PRIVATE | PK_STD, &key));