diff options
| -rw-r--r-- | pkg/pack/pack.go | 4 | ||||
| -rw-r--r-- | pkg/server/hook/auth.go | 26 | ||||
| -rw-r--r-- | pkg/server/queue/queue.go | 2 |
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} } |