From 514f8647eac04939916eb68e894da95dff2e5402 Mon Sep 17 00:00:00 2001
From: yuhan6665 <1588741+yuhan6665@users.noreply.github.com>
Date: Tue, 13 May 2025 08:52:09 -0400
Subject: [PATCH] More fixes for AES block

---
 cipher_suites.go | 8 ++++----
 tls.go           | 1 -
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/cipher_suites.go b/cipher_suites.go
index 57ce64c..c443a82 100644
--- a/cipher_suites.go
+++ b/cipher_suites.go
@@ -512,11 +512,11 @@ func aeadAESGCM(key, noncePrefix []byte) aead {
 	if len(noncePrefix) != noncePrefixLength {
 		panic("tls: internal error: wrong nonce length")
 	}
-	aes, err := aes.NewCipher(key)
+	aes, err := fipsaes.New(key)
 	if err != nil {
 		panic(err)
 	}
-	aead, err := gcm.NewGCMForTLS12(aes.(*fipsaes.Block))
+	aead, err := gcm.NewGCMForTLS12(aes)
 	if err != nil {
 		panic(err)
 	}
@@ -540,11 +540,11 @@ func aeadAESGCMTLS13(key, nonceMask []byte) aead {
 	if len(nonceMask) != aeadNonceLength {
 		panic("tls: internal error: wrong nonce length")
 	}
-	aes, err := aes.NewCipher(key)
+	aes, err := fipsaes.New(key)
 	if err != nil {
 		panic(err)
 	}
-	aead, err := gcm.NewGCMForTLS13(aes.(*fipsaes.Block))
+	aead, err := gcm.NewGCMForTLS13(aes)
 	if err != nil {
 		panic(err)
 	}
diff --git a/tls.go b/tls.go
index de585ed..1261b04 100644
--- a/tls.go
+++ b/tls.go
@@ -30,7 +30,6 @@ import (
 	"bytes"
 	"context"
 	"crypto"
-	"crypto/aes"
 	"crypto/cipher"
 	"crypto/ecdsa"
 	"crypto/ed25519"