diff --git a/library/dhm.c b/library/dhm.c
index 6b4e29f..979fd07 100644
--- a/library/dhm.c
+++ b/library/dhm.c
@@ -344,14 +344,14 @@
  * Derive and export the shared secret (G^Y)^X mod P
  */
 int mbedtls_dhm_calc_secret( mbedtls_dhm_context *ctx,
-                     unsigned char *output, size_t *olen,
+                     unsigned char *output, size_t output_size, size_t *olen,
                      int (*f_rng)(void *, unsigned char *, size_t),
                      void *p_rng )
 {
     int ret;
     mbedtls_mpi GYb;
 
-    if( ctx == NULL || *olen < ctx->len )
+    if( ctx == NULL || output_size < ctx->len )
         return( MBEDTLS_ERR_DHM_BAD_INPUT_DATA );
 
     if( ( ret = dhm_check_range( &ctx->GY, &ctx->P ) ) != 0 )
diff --git a/library/ssl_cli.c b/library/ssl_cli.c
index cff8871..6eb190c 100644
--- a/library/ssl_cli.c
+++ b/library/ssl_cli.c
@@ -2477,10 +2477,9 @@
         MBEDTLS_SSL_DEBUG_MPI( 3, "DHM: X ", &ssl->handshake->dhm_ctx.X  );
         MBEDTLS_SSL_DEBUG_MPI( 3, "DHM: GX", &ssl->handshake->dhm_ctx.GX );
 
-        ssl->handshake->pmslen = MBEDTLS_PREMASTER_SIZE;
-
         if( ( ret = mbedtls_dhm_calc_secret( &ssl->handshake->dhm_ctx,
                                       ssl->handshake->premaster,
+                                      MBEDTLS_PREMASTER_SIZE,
                                      &ssl->handshake->pmslen,
                                       ssl->conf->f_rng, ssl->conf->p_rng ) ) != 0 )
         {
diff --git a/library/ssl_srv.c b/library/ssl_srv.c
index 72f9eee..7db5a3c 100644
--- a/library/ssl_srv.c
+++ b/library/ssl_srv.c
@@ -3145,10 +3145,9 @@
             return( MBEDTLS_ERR_SSL_BAD_HS_CLIENT_KEY_EXCHANGE );
         }
 
-        ssl->handshake->pmslen = MBEDTLS_PREMASTER_SIZE;
-
         if( ( ret = mbedtls_dhm_calc_secret( &ssl->handshake->dhm_ctx,
                                       ssl->handshake->premaster,
+                                      MBEDTLS_PREMASTER_SIZE,
                                      &ssl->handshake->pmslen,
                                       ssl->conf->f_rng, ssl->conf->p_rng ) ) != 0 )
         {
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 443b421..ee32502 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -1114,11 +1114,11 @@
     if( key_ex == MBEDTLS_KEY_EXCHANGE_DHE_PSK )
     {
         int ret;
-        size_t len = end - ( p + 2 );
+        size_t len;
 
         /* Write length only when we know the actual value */
         if( ( ret = mbedtls_dhm_calc_secret( &ssl->handshake->dhm_ctx,
-                                      p + 2, &len,
+                                      p + 2, end - ( p + 2 ), &len,
                                       ssl->conf->f_rng, ssl->conf->p_rng ) ) != 0 )
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_dhm_calc_secret", ret );