From fe8b2d78efb0d632be9ddc4e8c3ac152913ee4af Mon Sep 17 00:00:00 2001
From: Andy Wang <cbeuw.andy@gmail.com>
Date: Fri, 24 Jan 2020 14:38:41 +0000
Subject: [PATCH] Change authenticationInfo structure

---
 internal/server/TLS.go       | 5 ++---
 internal/server/auth.go      | 6 +++---
 internal/server/websocket.go | 6 +++---
 3 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/internal/server/TLS.go b/internal/server/TLS.go
index e9fe36d..219ab8c 100644
--- a/internal/server/TLS.go
+++ b/internal/server/TLS.go
@@ -54,14 +54,13 @@ func (TLS) handshake(clientHello []byte, privateKey crypto.PrivateKey, originalC
 }
 
 func unmarshalClientHello(ch *ClientHello, staticPv crypto.PrivateKey) (ai authenticationInfo, err error) {
-	ephPub, ok := ecdh.Unmarshal(ch.random)
+	ai.randPubKey = ch.random
+	ephPub, ok := ecdh.Unmarshal(ai.randPubKey)
 	if !ok {
 		err = ErrInvalidPubKey
 		return
 	}
 
-	ai.nonce = ch.random[:12]
-
 	ai.sharedSecret = ecdh.GenerateSharedSecret(staticPv, ephPub)
 	var keyShare []byte
 	keyShare, err = parseKeyShare(ch.extensions[[2]byte{0x00, 0x33}])
diff --git a/internal/server/auth.go b/internal/server/auth.go
index 8054f37..6bdc7c5 100644
--- a/internal/server/auth.go
+++ b/internal/server/auth.go
@@ -23,7 +23,7 @@ type ClientInfo struct {
 
 type authenticationInfo struct {
 	sharedSecret      []byte
-	nonce             []byte
+	randPubKey        []byte
 	ciphertextWithTag []byte
 }
 
@@ -37,7 +37,7 @@ var ErrUnreconisedProtocol = errors.New("unreconised protocol")
 // touchStone checks if a the authenticationInfo are valid. It doesn't check if the UID is authorised
 func touchStone(ai authenticationInfo, now func() time.Time) (info ClientInfo, err error) {
 	var plaintext []byte
-	plaintext, err = util.AESGCMDecrypt(ai.nonce, ai.sharedSecret, ai.ciphertextWithTag)
+	plaintext, err = util.AESGCMDecrypt(ai.randPubKey[0:12], ai.sharedSecret, ai.ciphertextWithTag)
 	if err != nil {
 		return
 	}
@@ -87,7 +87,7 @@ func PrepareConnection(firstPacket []byte, sta *State, conn net.Conn) (info Clie
 		return
 	}
 
-	if sta.registerRandom(ai.nonce) {
+	if sta.registerRandom(ai.randPubKey) {
 		err = ErrReplay
 		return
 	}
diff --git a/internal/server/websocket.go b/internal/server/websocket.go
index b3b832d..7d33c82 100644
--- a/internal/server/websocket.go
+++ b/internal/server/websocket.go
@@ -73,14 +73,14 @@ func unmarshalHidden(hidden []byte, staticPv crypto.PrivateKey) (ai authenticati
 		err = ErrBadGET
 		return
 	}
-	ephPub, ok := ecdh.Unmarshal(hidden[0:32])
+
+	ai.randPubKey = hidden[0:32]
+	ephPub, ok := ecdh.Unmarshal(ai.randPubKey)
 	if !ok {
 		err = ErrInvalidPubKey
 		return
 	}
 
-	ai.nonce = hidden[:12]
-
 	ai.sharedSecret = ecdh.GenerateSharedSecret(staticPv, ephPub)
 
 	ai.ciphertextWithTag = hidden[32:]