From 1c04e859d9be1418a41f589418662333aa63d03c Mon Sep 17 00:00:00 2001
From: Kris Kwiatkowski <kris@amongbytes.com>
Date: Tue, 9 Apr 2019 23:47:53 +0100
Subject: [PATCH] CTR-DRBG: Use hardware acceleration on X86

benchmark              old ns/op     new ns/op     delta
BenchmarkInit-4        3403          397           -88.33%
BenchmarkRead-4        14535         1560          -89.27%
---
 drbg/ctr_drbg.go                | 8 ++++++--
 drbg/internal/aes/cipher.go     | 7 +++++++
 drbg/internal/aes/cipher_asm.go | 4 ++--
 3 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/drbg/ctr_drbg.go b/drbg/ctr_drbg.go
index 70aa9aa..2d3c387 100644
--- a/drbg/ctr_drbg.go
+++ b/drbg/ctr_drbg.go
@@ -14,6 +14,7 @@ package drbg
 
 import (
 	"github.com/henrydcase/nobs/drbg/internal/aes"
+	"github.com/henrydcase/nobs/utils"
 )
 
 // Constants below correspond to AES-256, which is currently
@@ -30,12 +31,15 @@ type CtrDrbg struct {
 	counter    uint
 	strength   uint
 	resistance bool
-	blockEnc   aes.AES
+	blockEnc   aes.IAES
 	tmpBlk     [3 * BlockLen]byte
 }
 
 func NewCtrDrbg() *CtrDrbg {
-	return new(CtrDrbg)
+	if utils.X86.HasAES {
+		return &CtrDrbg{blockEnc: &aes.AESAsm{}}
+	}
+	return &CtrDrbg{blockEnc: &aes.AES{}}
 }
 
 func (c *CtrDrbg) inc() {
diff --git a/drbg/internal/aes/cipher.go b/drbg/internal/aes/cipher.go
index 554bd03..db233df 100644
--- a/drbg/internal/aes/cipher.go
+++ b/drbg/internal/aes/cipher.go
@@ -18,6 +18,13 @@ type AES struct {
 	keyLen int
 }
 
+// AES interface
+type IAES interface {
+	SetKey(key []byte) error
+	Encrypt(dst, src []byte)
+	Decrypt(dst, src []byte)
+}
+
 type KeySizeError int
 
 func (k KeySizeError) Error() string {
diff --git a/drbg/internal/aes/cipher_asm.go b/drbg/internal/aes/cipher_asm.go
index 050ec85..e2738fd 100644
--- a/drbg/internal/aes/cipher_asm.go
+++ b/drbg/internal/aes/cipher_asm.go
@@ -21,8 +21,8 @@ func decryptBlockAsm(nr int, xk *uint32, dst, src *byte)
 func expandKeyAsm(nr int, key *byte, enc *uint32, dec *uint32)
 
 type AESAsm struct {
-	enc []uint32
-	dec []uint32
+	enc [32 + 28]uint32
+	dec [32 + 28]uint32
 }
 
 func (c *AESAsm) SetKey(key []byte) error {