make auth (chap)

4 files changed, 164 insertions, 9 deletions

diff --git a/pkg/server/env/env.go b/pkg/server/env/env.go
index ab47ae8..affb4ef 100644
--- a/pkg/server/env/env.go
+++ b/pkg/server/env/env.go
@@ -15,7 +15,7 @@ type Env struct {
 	*env
 }
 
-const namePattern = "[a-zA-Z][a-zA-Z0-9]*"
+const namePattern = "[a-zA-Z][a-zA-Z0-9.]*"
 var isNamePattern = regexp.MustCompile("^" + namePattern + "$").MatchString
 var namePatternRe = regexp.MustCompile("@" + namePattern)
 
@@ -25,6 +25,26 @@ func New() Env {
 	return Env{new(env)}
 }
 
+func (e *env) init() {
+	if e.m == nil {
+		e.m = make(map[string]string)
+	}
+}
+
+func (e *env) Copy() Env {
+	c := New()
+
+	if len(e.m) > 0 {
+		c.init()
+
+		for k, v := range e.m {
+			c.m[k] = v
+		}
+	}
+
+	return c
+}
+
 func (e *env) Find(key string) (string, bool) {
 	e.Lock()
 	defer e.Unlock()
@@ -47,9 +67,7 @@ func (e *env) Set(key string, value string) error {
 	e.Lock()
 	defer e.Unlock()
 
-	if e.m == nil {
-		e.m = make(map[string]string)
-	}
+	e.init()
 
 	e.m[key] = value
 
diff --git a/pkg/server/module/auth.go b/pkg/server/module/auth.go
index 05761ed..5dbdccb 100644
--- a/pkg/server/module/auth.go
+++ b/pkg/server/module/auth.go
@@ -4,24 +4,158 @@ import (
 	"tunnel/pkg/server/queue"
 	"tunnel/pkg/server/opts"
 	"tunnel/pkg/server/env"
+	"tunnel/pkg/netstring"
+	"encoding/hex"
+	"crypto/rand"
+	"crypto/md5"
+	"errors"
+	"fmt"
+	"io"
 )
 
+const ChallengeLen = 16
+
 type auth struct {
 	secret string
+
+	challenge struct {
+		self string
+		peer string
+	}
+
+	hash string
+
+	recvChallenge chan struct{}
+	recvHash chan struct{}
+
+	fail chan struct{}
+	ok chan struct{}
 }
 
+var errDupChallenge = errors.New("peer duplicates challenge")
+var errAuthFail = errors.New("peer auth fail")
+
 type authModule struct{}
 
+func (a *auth) generateChallenge() error {
+	b := make([]byte, ChallengeLen)
+	if _, err := rand.Read(b); err != nil {
+		return err
+	}
+
+	c := make([]byte, hex.EncodedLen(len(b)))
+	hex.Encode(c, b)
+
+	a.challenge.self = string(c)
+
+	return nil
+}
+
+func (a *auth) sendChallenge(q queue.Q) bool {
+	enc := netstring.NewEncoder(q.Writer())
+	enc.Encode(a.challenge.self)
+	return a.wait(a.recvChallenge)
+}
+
+func (a *auth) getHash(c string) string {
+	h := md5.New()
+	io.WriteString(h, a.secret)
+	io.WriteString(h, c)
+
+	b := make([]byte, hex.EncodedLen(h.Size()))
+	hex.Encode(b, h.Sum(nil))
+
+	return string(b)
+}
+
+func (a *auth) sendHash(q queue.Q) bool {
+	enc := netstring.NewEncoder(q.Writer())
+	enc.Encode(a.getHash(a.challenge.peer))
+	return a.wait(a.recvHash)
+}
+
+func (a *auth) wait(c chan struct{}) bool {
+	select {
+	case <-a.fail:
+		return false
+	case <-c:
+		return true
+	}
+}
+
 func (a *auth) Send(rq, wq queue.Q) error {
+	if err := a.generateChallenge(); err != nil {
+		return err
+	}
+
+	if !a.sendChallenge(wq) {
+		return nil
+	}
+
+	if a.challenge.self == a.challenge.peer {
+		return errDupChallenge
+	}
+
+	if !a.sendHash(wq) {
+		return nil
+	}
+
+	if a.hash != a.getHash(a.challenge.self) {
+		close(a.fail)
+		return errAuthFail
+	}
+
+	close(a.ok)
+
 	return queue.Copy(rq, wq)
 }
 
 func (a *auth) Recv(rq, wq queue.Q) error {
+	dec := netstring.NewDecoder(rq.Reader())
+
+	if c, err := dec.Decode(); err != nil {
+		close(a.fail)
+		return err
+	} else {
+		a.challenge.peer = c
+		close(a.recvChallenge)
+	}
+
+	if h, err := dec.Decode(); err != nil {
+		close(a.fail)
+		return err
+	} else {
+		a.hash = h
+		close(a.recvHash)
+	}
+
+	if !a.wait(a.ok) {
+		return nil
+	}
+
 	return queue.Copy(rq, wq)
 }
 
+func getAuthSecret(env env.Env) string {
+	if id, ok := env.Find("tunnel"); ok {
+		k := fmt.Sprintf("tunnel.%s.secret", id)
+
+		if v, ok := env.Find(k); ok {
+			return v
+		}
+	}
+
+	return env.Get("secret")
+}
+
 func (m authModule) Open(env env.Env) (Pipe, Pipe) {
-	a := &auth{env.Get("secret")}
+	a := &auth{
+		secret: getAuthSecret(env),
+		recvChallenge: make(chan struct{}),
+		recvHash: make(chan struct{}),
+		fail: make(chan struct{}),
+		ok: make(chan struct{}),
+	}
 	return a.Send, a.Recv
 }
 
diff --git a/pkg/server/server.go b/pkg/server/server.go
index d380fb4..a7b9b06 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -238,11 +238,11 @@ func (r *request) decode(query string) []string {
 		if s, err := dec.Decode(); err == nil {
 			t = append(t, s)
 		} else {
-			if err == io.EOF {
-				break
+			if !errors.Is(err, io.EOF) {
+				r.Fatal("request parse failed")
 			}
 
-			r.Fatal("failed to parse request")
+			break
 		}
 	}
 
diff --git a/pkg/server/tunnel.go b/pkg/server/tunnel.go
index a7854bc..32c81c3 100644
--- a/pkg/server/tunnel.go
+++ b/pkg/server/tunnel.go
@@ -195,6 +195,9 @@ func (s *stream) pipe(m module.M, p module.Pipe, rq, wq queue.Q) {
 }
 
 func (s *stream) run() {
+	env := s.t.env.Copy()
+	env.Set("tunnel", s.t.id)
+
 	s.t.wg.Add(1)
 
 	rq, wq := queue.New(), queue.New()
@@ -202,7 +205,7 @@ func (s *stream) run() {
 	s.channel(s.in, rq, wq)
 
 	for _, m := range s.t.m {
-		send, recv := m.Open(s.t.env)
+		send, recv := m.Open(env)
 		if send != nil {
 			q := queue.New()
 			s.pipe(m, send, wq, q)