From b5c6d7fedcd5b7f36b1df60ef1c5f0430611bb8c Mon Sep 17 00:00:00 2001
From: Andy Wang <cbeuw.andy@gmail.com>
Date: Fri, 16 Aug 2019 23:56:46 +0100
Subject: [PATCH] Change touchUp function signature

---
 internal/server/TLS.go  |  8 +++++++-
 internal/server/auth.go | 14 ++++----------
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/internal/server/TLS.go b/internal/server/TLS.go
index 04c24bc..d00f466 100644
--- a/internal/server/TLS.go
+++ b/internal/server/TLS.go
@@ -216,6 +216,7 @@ func composeReply(ch *ClientHello, sharedSecret []byte, sessionKey []byte) ([]by
 
 var ErrBadClientHello = errors.New("non (or malformed) ClientHello")
 var ErrNotCloak = errors.New("TLS but non-Cloak ClientHello")
+var ErrReplay = errors.New("duplicate random")
 var ErrBadProxyMethod = errors.New("invalid proxy method")
 
 func PrepareConnection(firstPacket []byte, sta *State, conn net.Conn) (info ClientInfo, finisher func([]byte) error, err error) {
@@ -226,8 +227,13 @@ func PrepareConnection(firstPacket []byte, sta *State, conn net.Conn) (info Clie
 		return
 	}
 
+	if sta.registerRandom(ch.random) {
+		err = ErrReplay
+		return
+	}
+
 	var sharedSecret []byte
-	info, sharedSecret, err = TouchStone(ch, sta)
+	info, sharedSecret, err = touchStone(ch, sta.staticPv, sta.Now)
 	if err != nil {
 		log.Debug(err)
 		err = ErrNotCloak
diff --git a/internal/server/auth.go b/internal/server/auth.go
index e7dbb3d..82be4dc 100644
--- a/internal/server/auth.go
+++ b/internal/server/auth.go
@@ -2,6 +2,7 @@ package server
 
 import (
 	"bytes"
+	"crypto"
 	"encoding/binary"
 	"errors"
 	"fmt"
@@ -22,25 +23,18 @@ const (
 	UNORDERED_FLAG = 0x01 // 0000 0001
 )
 
-var ErrReplay = errors.New("duplicate random")
 var ErrInvalidPubKey = errors.New("public key has invalid format")
 var ErrCiphertextLength = errors.New("ciphertext has the wrong length")
 var ErrTimestampOutOfWindow = errors.New("timestamp is outside of the accepting window")
 
-func TouchStone(ch *ClientHello, sta *State) (info ClientInfo, sharedSecret []byte, err error) {
-
-	if sta.registerRandom(ch.random) {
-		err = ErrReplay
-		return
-	}
-
+func touchStone(ch *ClientHello, staticPv crypto.PrivateKey, now func() time.Time) (info ClientInfo, sharedSecret []byte, err error) {
 	ephPub, ok := ecdh.Unmarshal(ch.random)
 	if !ok {
 		err = ErrInvalidPubKey
 		return
 	}
 
-	sharedSecret = ecdh.GenerateSharedSecret(sta.staticPv, ephPub)
+	sharedSecret = ecdh.GenerateSharedSecret(staticPv, ephPub)
 	var keyShare []byte
 	keyShare, err = parseKeyShare(ch.extensions[[2]byte{0x00, 0x33}])
 	if err != nil {
@@ -69,7 +63,7 @@ func TouchStone(ch *ClientHello, sta *State) (info ClientInfo, sharedSecret []by
 
 	timestamp := int64(binary.BigEndian.Uint64(plaintext[29:37]))
 	clientTime := time.Unix(timestamp, 0)
-	serverTime := sta.Now()
+	serverTime := now()
 	if !(clientTime.After(serverTime.Truncate(TIMESTAMP_TOLERANCE)) && clientTime.Before(serverTime.Add(TIMESTAMP_TOLERANCE))) {
 		err = fmt.Errorf("%v: received timestamp %v", ErrTimestampOutOfWindow, timestamp)
 		return