fix hook variable check

4 files changed, 11 insertions, 15 deletions

diff --git a/pkg/server/hook/aes.go b/pkg/server/hook/aes.go
index b461a34..414f6d6 100644
--- a/pkg/server/hook/aes.go
+++ b/pkg/server/hook/aes.go
@@ -65,7 +65,7 @@ func (a *aesPipe) Recv(rq, wq queue.Q) error {
 }
 
 func newAes(env env.Env) *aesPipe {
-	s := getAuthSecret(env)
+	s := getHookVar(env, "secret")
 	h := md5.Sum([]byte(s))
 
 	a := &aesPipe{key: make([]byte, 16)}
diff --git a/pkg/server/hook/auth.go b/pkg/server/hook/auth.go
index fc19c2a..8d7d23e 100644
--- a/pkg/server/hook/auth.go
+++ b/pkg/server/hook/auth.go
@@ -122,17 +122,9 @@ func (a *auth) Recv(rq, wq queue.Q) (err error) {
 	return queue.IoCopy(r, wq.Writer())
 }
 
-func getAuthSecret(env env.Env) string {
-	if v := env.Eval("@{tunnel.@{tunnel}.secret}"); v != "" {
-		return v
-	}
-
-	return env.Get("secret")
-}
-
 func (h authHook) Open(env env.Env) (interface{}, error) {
 	a := &auth{
-		secret:        getAuthSecret(env),
+		secret:        getHookVar(env, "secret"),
 		recvChallenge: make(chan struct{}),
 		recvHash:      make(chan struct{}),
 		fail:          make(chan struct{}),
diff --git a/pkg/server/hook/hook.go b/pkg/server/hook/hook.go
index b2970ac..e4497ff 100644
--- a/pkg/server/hook/hook.go
+++ b/pkg/server/hook/hook.go
@@ -99,6 +99,14 @@ func New(desc string, env env.Env) (H, error) {
 	}
 }
 
+func getHookVar(env env.Env, s string) string {
+	if v := env.Eval("@{tunnel.@{tunnel}." + s + "}"); v != "" {
+		return v
+	}
+
+	return env.Get(s)
+}
+
 func register(name string, f hookInitFunc) {
 	if _, ok := hooks[name]; ok {
 		log.Panicf("duplicate hook name '%s'", name)
diff --git a/pkg/server/hook/tee.go b/pkg/server/hook/tee.go
index 6591e33..fd30a81 100644
--- a/pkg/server/hook/tee.go
+++ b/pkg/server/hook/tee.go
@@ -67,11 +67,7 @@ func (h *teeHook) where(env env.Env) string {
 		return h.file
 	}
 
-	if v := env.Eval("@{tunnel.@{tunnel}.tee.file}"); v != "" {
-		return v
-	}
-
-	if v, ok := env.Find("hook.tee.file"); ok {
+	if v := getHookVar(env, "tee.file"); v != "" {
 		return v
 	}