auth with pack

3 files changed, 18 insertions, 14 deletions

diff --git a/pkg/pack/pack.go b/pkg/pack/pack.go
index e186773..0511df8 100644
--- a/pkg/pack/pack.go
+++ b/pkg/pack/pack.go
@@ -29,7 +29,7 @@ func NewDecoder(r Reader) *Decoder {
 
 func (e *Encoder) Lps(b []byte) error {
 	buf := make([]byte, binary.MaxVarintLen64)
-	n := binary.PutVarint(buf, int64(len(b)))
+	n := binary.PutUvarint(buf, uint64(len(b)))
 
 	if _, err := e.w.Write(buf[:n]); err != nil {
 		return err
@@ -40,7 +40,7 @@ func (e *Encoder) Lps(b []byte) error {
 }
 
 func (d *Decoder) Lps() ([]byte, error) {
-	if n, err := binary.ReadVarint(d.r); err != nil {
+	if n, err := binary.ReadUvarint(d.r); err != nil {
 		return nil, fmt.Errorf("lps: %w", err)
 	} else {
 		buf := make([]byte, n)
diff --git a/pkg/server/hook/auth.go b/pkg/server/hook/auth.go
index f347b2a..b31855a 100644
--- a/pkg/server/hook/auth.go
+++ b/pkg/server/hook/auth.go
@@ -7,7 +7,7 @@ import (
 	"io"
 	"sync"
 	"time"
-	"tunnel/pkg/netstring"
+	"tunnel/pkg/pack"
 	"tunnel/pkg/server/env"
 	"tunnel/pkg/server/opts"
 	"tunnel/pkg/server/queue"
@@ -83,7 +83,7 @@ func (a *auth) wait(c chan struct{}) error {
 }
 
 func (a *auth) Send(rq, wq queue.Q) error {
-	e := netstring.NewEncoder(wq.Writer())
+	e := pack.NewEncoder(wq.Writer())
 
 	if err := a.generateChallenge(); err != nil {
 		return err
@@ -91,7 +91,7 @@ func (a *auth) Send(rq, wq queue.Q) error {
 
 	defer a.deleteChallenge()
 
-	e.Encode(a.challenge.self)
+	e.Lps([]byte(a.challenge.self))
 
 	if err := a.wait(a.recvChallenge); err != nil {
 		return err
@@ -101,7 +101,7 @@ func (a *auth) Send(rq, wq queue.Q) error {
 		return errDupChallenge
 	}
 
-	e.Encode(a.getHash(a.challenge.peer))
+	e.Lps([]byte(a.getHash(a.challenge.peer)))
 
 	if err := a.wait(a.recvHash); err != nil {
 		return err
@@ -117,26 +117,30 @@ func (a *auth) Send(rq, wq queue.Q) error {
 	return queue.Copy(rq, wq)
 }
 
-func (a *auth) Recv(rq, wq queue.Q) (err error) {
+func (a *auth) Recv(rq, wq queue.Q) error {
 	r := rq.Reader()
-	d := netstring.NewDecoder(r)
+	d := pack.NewDecoder(r)
 
-	if a.challenge.peer, err = d.Decode(); err != nil {
+	if b, err := d.Lps(); err != nil {
 		close(a.fail)
-		return
+		return err
+	} else {
+		a.challenge.peer = string(b)
 	}
 
 	close(a.recvChallenge)
 
-	if a.hash, err = d.Decode(); err != nil {
+	if b, err := d.Lps(); err != nil {
 		close(a.fail)
 		return err
+	} else {
+		a.hash = string(b)
 	}
 
 	close(a.recvHash)
 
-	if err = a.wait(a.ok); err != nil {
-		return
+	if err := a.wait(a.ok); err != nil {
+		return err
 	}
 
 	a.tmr.Stop()
diff --git a/pkg/server/queue/queue.go b/pkg/server/queue/queue.go
index 71f914c..704d884 100644
--- a/pkg/server/queue/queue.go
+++ b/pkg/server/queue/queue.go
@@ -19,7 +19,7 @@ func New() Q {
 	return make(Q)
 }
 
-func (q Q) Reader() io.Reader {
+func (q Q) Reader() *reader {
 	return &reader{q: q}
 }