ssl, dtls: Refactor so that DTLS records are handled correctly together with AEAD handling

4 files changed, 75 insertions, 75 deletions

diff --git a/lib/ssl/src/dtls_record.erl b/lib/ssl/src/dtls_record.erl
index c5992a72fa..049f83e49e 100644
--- a/lib/ssl/src/dtls_record.erl
+++ b/lib/ssl/src/dtls_record.erl
@@ -439,43 +439,59 @@ encode_dtls_cipher_text(Type, {MajVer, MinVer}, Fragment,
 encode_plain_text(Type, Version, Data, #{compression_state := CompS0,
 					 epoch := Epoch,
 					 sequence_number := Seq,
+                                         cipher_state := CipherS0,
 					 security_parameters :=
 					     #security_parameters{
 						cipher_type = ?AEAD,
+                                                  bulk_cipher_algorithm =
+                                                    BulkCipherAlgo,
 						compression_algorithm = CompAlg}
 					} = WriteState0) ->
     {Comp, CompS1} = ssl_record:compress(CompAlg, Data, CompS0),
-    WriteState1 = WriteState0#{compression_state => CompS1},
     AAD = calc_aad(Type, Version, Epoch, Seq),
-    ssl_record:cipher_aead(dtls_v1:corresponding_tls_version(Version), Comp, WriteState1, AAD);
-encode_plain_text(Type, Version, Data, #{compression_state := CompS0,
+    TLSVersion = dtls_v1:corresponding_tls_version(Version),
+    {CipherFragment, CipherS1} =
+	ssl_cipher:cipher_aead(BulkCipherAlgo, CipherS0, Seq, AAD, Comp, TLSVersion),
+    {CipherFragment,  WriteState0#{compression_state => CompS1,
+                                   cipher_state => CipherS1}};
+encode_plain_text(Type, Version, Fragment, #{compression_state := CompS0,
 					 epoch := Epoch,
 					 sequence_number := Seq,
+                                         cipher_state := CipherS0,
 					 security_parameters :=
-					     #security_parameters{compression_algorithm = CompAlg}
+					     #security_parameters{compression_algorithm = CompAlg,
+                                                                  bulk_cipher_algorithm =
+                                                                      BulkCipherAlgo}
 					}= WriteState0) ->
-    {Comp, CompS1} = ssl_record:compress(CompAlg, Data, CompS0),
+    {Comp, CompS1} = ssl_record:compress(CompAlg, Fragment, CompS0),
     WriteState1 = WriteState0#{compression_state => CompS1},
-    MacHash = calc_mac_hash(Type, Version, WriteState1, Epoch, Seq, Comp),
-    ssl_record:cipher(dtls_v1:corresponding_tls_version(Version), Comp, WriteState1, MacHash).
+    MAC = calc_mac_hash(Type, Version, WriteState1, Epoch, Seq, Comp),
+    TLSVersion = dtls_v1:corresponding_tls_version(Version),
+    {CipherFragment, CipherS1} =
+	ssl_cipher:cipher(BulkCipherAlgo, CipherS0, MAC, Fragment, TLSVersion),
+    {CipherFragment,  WriteState0#{cipher_state => CipherS1}}.
 
 decode_cipher_text(#ssl_tls{type = Type, version = Version,
 			    epoch = Epoch,
 			    sequence_number = Seq,
 			    fragment = CipherFragment} = CipherText,
 		   #{compression_state := CompressionS0,
+                     cipher_state := CipherS0,
 		     security_parameters :=
 			 #security_parameters{
 			    cipher_type = ?AEAD,
+                            bulk_cipher_algorithm =
+                                BulkCipherAlgo,
 			    compression_algorithm = CompAlg}} = ReadState0, 
 		   ConnnectionStates0) ->
     AAD = calc_aad(Type, Version, Epoch, Seq),
-    case ssl_record:decipher_aead(dtls_v1:corresponding_tls_version(Version),
-				  CipherFragment, ReadState0, AAD) of
-	{PlainFragment, ReadState1} ->
+    TLSVersion = dtls_v1:corresponding_tls_version(Version),
+    case  ssl_cipher:decipher_aead(BulkCipherAlgo, CipherS0, Seq, AAD, CipherFragment, TLSVersion) of
+	{PlainFragment, CipherState} ->
 	    {Plain, CompressionS1} = ssl_record:uncompress(CompAlg,
 							   PlainFragment, CompressionS0),
-	    ReadState = ReadState1#{compression_state => CompressionS1},
+	    ReadState = ReadState0#{compression_state => CompressionS1,
+                                    cipher_state => CipherState},
 	    ConnnectionStates = set_connection_state_by_epoch(ReadState, Epoch, ConnnectionStates0, read),
 	    {CipherText#ssl_tls{fragment = Plain}, ConnnectionStates};
 	  #alert{} = Alert ->
diff --git a/lib/ssl/src/ssl_cipher.erl b/lib/ssl/src/ssl_cipher.erl
index 3919070e9a..d04f09efdc 100644
--- a/lib/ssl/src/ssl_cipher.erl
+++ b/lib/ssl/src/ssl_cipher.erl
@@ -40,7 +40,7 @@
 	 ec_keyed_suites/0, anonymous_suites/1, psk_suites/1, srp_suites/0,
 	 rc4_suites/1, des_suites/1, openssl_suite/1, openssl_suite_name/1, filter/2, filter_suites/1,
 	 hash_algorithm/1, sign_algorithm/1, is_acceptable_hash/2, is_fallback/1,
-	 random_bytes/1, calc_aad/3, calc_mac_hash/4,
+	 random_bytes/1, calc_mac_hash/4,
          is_stream_ciphersuite/1]).
 
 -export_type([cipher_suite/0,
@@ -1531,10 +1531,6 @@ is_fallback(CipherSuites)->
 random_bytes(N) ->
     crypto:strong_rand_bytes(N).
 
-calc_aad(Type, {MajVer, MinVer},
-	 #{sequence_number := SeqNo}) ->
-    <<?UINT64(SeqNo), ?BYTE(Type), ?BYTE(MajVer), ?BYTE(MinVer)>>.
-
 calc_mac_hash(Type, Version,
 	      PlainFragment, #{sequence_number := SeqNo,
 			       mac_secret := MacSecret,
diff --git a/lib/ssl/src/ssl_record.erl b/lib/ssl/src/ssl_record.erl
index 539e189c4f..24e52655b0 100644
--- a/lib/ssl/src/ssl_record.erl
+++ b/lib/ssl/src/ssl_record.erl
@@ -45,11 +45,7 @@
 -export([compress/3, uncompress/3, compressions/0]).
 
 %% Payload encryption/decryption
--export([cipher/4, decipher/4, is_correct_mac/2,
-	 cipher_aead/4, decipher_aead/4]).
-
-%% Encoding
--export([encode_plain_text/4]).
+-export([cipher/4, decipher/4, cipher_aead/4, is_correct_mac/2]).
 
 -export_type([ssl_version/0, ssl_atom_version/0, connection_states/0, connection_state/0]).
 
@@ -271,26 +267,6 @@ set_pending_cipher_state(#{pending_read := Read,
       pending_read => Read#{cipher_state => ServerState},
       pending_write => Write#{cipher_state => ClientState}}.
 
-encode_plain_text(Type, Version, Data, #{compression_state := CompS0,
-					 security_parameters :=
-					     #security_parameters{
-						cipher_type = ?AEAD,
-						compression_algorithm = CompAlg}
-					} = WriteState0) ->
-    {Comp, CompS1} = ssl_record:compress(CompAlg, Data, CompS0),
-    WriteState1 = WriteState0#{compression_state => CompS1},
-    AAD = ssl_cipher:calc_aad(Type, Version, WriteState1),
-    ssl_record:cipher_aead(Version, Comp, WriteState1, AAD);
-encode_plain_text(Type, Version, Data, #{compression_state := CompS0,
-					 security_parameters :=
-					     #security_parameters{compression_algorithm = CompAlg}
-					}= WriteState0) ->
-    {Comp, CompS1} = ssl_record:compress(CompAlg, Data, CompS0),
-    WriteState1 = WriteState0#{compression_state => CompS1},
-    MacHash = ssl_cipher:calc_mac_hash(Type, Version, Comp, WriteState1),
-    ssl_record:cipher(Version, Comp, WriteState1, MacHash);
-encode_plain_text(_,_,_,CS) ->
-    exit({cs, CS}).
 
 uncompress(?NULL, Data, CS) ->
     {Data, CS}.
@@ -322,12 +298,12 @@ cipher(Version, Fragment,
     {CipherFragment, CipherS1} =
 	ssl_cipher:cipher(BulkCipherAlgo, CipherS0, MacHash, Fragment, Version),
     {CipherFragment,  WriteState0#{cipher_state => CipherS1}}.
-%%--------------------------------------------------------------------
--spec cipher_aead(ssl_version(), iodata(), connection_state(), MacHash::binary()) ->
-			 {CipherFragment::binary(), connection_state()}.
-%%
-%% Description: Payload encryption
-%%--------------------------------------------------------------------
+%% %%--------------------------------------------------------------------
+%% -spec cipher_aead(ssl_version(), iodata(), connection_state(), MacHash::binary()) ->
+%% 			 {CipherFragment::binary(), connection_state()}.
+%% %%
+%% %% Description: Payload encryption
+%% %%--------------------------------------------------------------------
 cipher_aead(Version, Fragment,
 	    #{cipher_state := CipherS0,
 	      sequence_number := SeqNo,
@@ -341,7 +317,8 @@ cipher_aead(Version, Fragment,
     {CipherFragment,  WriteState0#{cipher_state => CipherS1}}.
 
 %%--------------------------------------------------------------------
--spec decipher(ssl_version(), binary(), connection_state(), boolean()) -> {binary(), binary(), connection_state} | #alert{}.
+-spec decipher(ssl_version(), binary(), connection_state(), boolean()) ->
+                      {binary(), binary(), connection_state} | #alert{}.
 %%
 %% Description: Payload decryption
 %%--------------------------------------------------------------------
@@ -359,26 +336,7 @@ decipher(Version, CipherFragment,
 	#alert{} = Alert ->
 	    Alert
     end.
-%%--------------------------------------------------------------------
--spec decipher_aead(ssl_version(), binary(), connection_state(), binary()) -> 
-			   {binary(), binary(), connection_state()} | #alert{}.
-%%
-%% Description: Payload decryption
-%%--------------------------------------------------------------------
-decipher_aead(Version, CipherFragment,
-	      #{sequence_number := SeqNo,
-		security_parameters :=
-		    #security_parameters{bulk_cipher_algorithm =
-					     BulkCipherAlgo},
-		cipher_state := CipherS0
-	       } = ReadState, AAD) ->
-    case ssl_cipher:decipher_aead(BulkCipherAlgo, CipherS0, SeqNo, AAD, CipherFragment, Version) of
-	{PlainFragment, CipherS1} ->
-	    CS1 = ReadState#{cipher_state => CipherS1},
-	    {PlainFragment, CS1};
-	#alert{} = Alert ->
-	    Alert
-    end.
+
 %%--------------------------------------------------------------------
 %%% Internal functions
 %%--------------------------------------------------------------------
diff --git a/lib/ssl/src/tls_record.erl b/lib/ssl/src/tls_record.erl
index 993a1622fe..065c6dc8a7 100644
--- a/lib/ssl/src/tls_record.erl
+++ b/lib/ssl/src/tls_record.erl
@@ -372,7 +372,7 @@ get_tls_records_aux(Data, Acc) ->
 	end.
 
 encode_plain_text(Type, Version, Data, #{current_write := Write0} = ConnectionStates) ->
-    {CipherFragment, Write1} = ssl_record:encode_plain_text(Type, Version, Data, Write0),
+    {CipherFragment, Write1} = do_encode_plain_text(Type, Version, Data, Write0),
     {CipherText, Write} = encode_tls_cipher_text(Type, Version, CipherFragment, Write1),
     {CipherText, ConnectionStates#{current_write => Write}}.
 
@@ -446,19 +446,24 @@ decode_cipher_text(#ssl_tls{type = Type, version = Version,
 		   #{current_read :=
 			 #{compression_state := CompressionS0,
 			   sequence_number := Seq,
+                           cipher_state := CipherS0,
 			   security_parameters :=
 			       #security_parameters{
 				  cipher_type = ?AEAD,
+                                  bulk_cipher_algorithm =
+                                      BulkCipherAlgo,
 				  compression_algorithm = CompAlg}
 			  } = ReadState0} = ConnnectionStates0, _) ->
-    AAD = ssl_cipher:calc_aad(Type, Version, ReadState0),
-    case ssl_record:decipher_aead(Version, CipherFragment, ReadState0, AAD) of
-	{PlainFragment, ReadState1} ->
+    AAD = calc_aad(Type, Version, ReadState0),
+    case ssl_cipher:decipher_aead(BulkCipherAlgo, CipherS0, Seq, AAD, CipherFragment, Version) of
+	{PlainFragment, CipherS1} ->
 	    {Plain, CompressionS1} = ssl_record:uncompress(CompAlg,
 							   PlainFragment, CompressionS0),
 	    ConnnectionStates = ConnnectionStates0#{
-				  current_read => ReadState1#{sequence_number => Seq + 1,
-							      compression_state => CompressionS1}},
+				  current_read => ReadState0#{
+                                                    cipher_state => CipherS1,
+                                                    sequence_number => Seq + 1,
+                                                    compression_state => CompressionS1}},
 	    {CipherText#ssl_tls{fragment = Plain}, ConnnectionStates};
 	#alert{} = Alert ->
 	    Alert
@@ -489,4 +494,29 @@ decode_cipher_text(#ssl_tls{type = Type, version = Version,
 	    end;
 	    #alert{} = Alert ->
 	    Alert
-    end. 
+    end.
+
+do_encode_plain_text(Type, Version, Data, #{compression_state := CompS0,
+					 security_parameters :=
+					     #security_parameters{
+						cipher_type = ?AEAD,
+						compression_algorithm = CompAlg}
+					} = WriteState0) ->
+    {Comp, CompS1} = ssl_record:compress(CompAlg, Data, CompS0),
+    WriteState1 = WriteState0#{compression_state => CompS1},
+    AAD = calc_aad(Type, Version, WriteState1),
+    ssl_record:cipher_aead(Version, Comp, WriteState1, AAD);
+do_encode_plain_text(Type, Version, Data, #{compression_state := CompS0,
+					 security_parameters :=
+					     #security_parameters{compression_algorithm = CompAlg}
+					}= WriteState0) ->
+    {Comp, CompS1} = ssl_record:compress(CompAlg, Data, CompS0),
+    WriteState1 = WriteState0#{compression_state => CompS1},
+    MacHash = ssl_cipher:calc_mac_hash(Type, Version, Comp, WriteState1),
+    ssl_record:cipher(Version, Comp, WriteState1, MacHash);
+do_encode_plain_text(_,_,_,CS) ->
+    exit({cs, CS}).
+
+calc_aad(Type, {MajVer, MinVer},
+	 #{sequence_number := SeqNo}) ->
+    <<?UINT64(SeqNo), ?BYTE(Type), ?BYTE(MajVer), ?BYTE(MinVer)>>.