1 files changed, 20 insertions, 6 deletions
diff --git a/pkg/server/hook/auth.go b/pkg/server/hook/auth.go
index 7f02816..f347b2a 100644
--- a/pkg/server/hook/auth.go
+++ b/pkg/server/hook/auth.go
@@ -5,6 +5,7 @@ import (
 	"crypto/rand"
 	"errors"
 	"io"
+	"sync"
 	"time"
 	"tunnel/pkg/netstring"
 	"tunnel/pkg/server/env"
@@ -16,6 +17,8 @@ const authTimeout = 5 * time.Second
 const challengeLen = 16
 
 type auth struct {
+	h *authHook
+
 	secret string
 
 	challenge struct {
@@ -34,11 +37,13 @@ type auth struct {
 	ok   chan struct{}
 }
 
-var errDupChallenge = errors.New("peer duplicates challenge")
+var errDupChallenge = errors.New("peer repeats challenge")
 var errAuthFail = errors.New("peer auth fail")
 var errTimeout = errors.New("timeout")
 
-type authHook struct{}
+type authHook struct {
+	m sync.Map
+}
 
 func (a *auth) generateChallenge() error {
 	b := make([]byte, challengeLen)
@@ -48,14 +53,20 @@ func (a *auth) generateChallenge() error {
 
 	a.challenge.self = string(b)
 
+	a.h.m.Store(a.challenge.self, struct{}{})
+
 	return nil
 }
 
+func (a *auth) deleteChallenge() {
+	a.h.m.Delete(a.challenge.self)
+}
+
 func (a *auth) getHash(c string) string {
 	h := md5.New()
 
-	io.WriteString(h, a.secret)
 	io.WriteString(h, c)
+	io.WriteString(h, a.secret)
 
 	return string(h.Sum(nil))
 }
@@ -78,13 +89,15 @@ func (a *auth) Send(rq, wq queue.Q) error {
 		return err
 	}
 
+	defer a.deleteChallenge()
+
 	e.Encode(a.challenge.self)
 
 	if err := a.wait(a.recvChallenge); err != nil {
 		return err
 	}
 
-	if a.challenge.self == a.challenge.peer {
+	if _, ok := a.h.m.Load(a.challenge.peer); ok {
 		return errDupChallenge
 	}
 
@@ -131,8 +144,9 @@ func (a *auth) Recv(rq, wq queue.Q) (err error) {
 	return queue.IoCopy(r, wq.Writer())
 }
 
-func (authHook) Open(env env.Env) (interface{}, error) {
+func (h *authHook) Open(env env.Env) (interface{}, error) {
 	a := &auth{
+		h:             h,
 		secret:        getHookVar(env, "secret"),
 		recvChallenge: make(chan struct{}),
 		recvHash:      make(chan struct{}),
@@ -144,7 +158,7 @@ func (authHook) Open(env env.Env) (interface{}, error) {
 }
 
 func newAuthHook(opts.Opts, env.Env) (hook, error) {
-	return authHook{}, nil
+	return &authHook{}, nil
 }
 
 func init() {