package hook import ( "crypto/md5" "crypto/rand" "errors" "io" "sync" "time" "tunnel/pkg/pack" "tunnel/pkg/server/env" "tunnel/pkg/server/opts" "tunnel/pkg/server/queue" ) const authTimeout = 5 * time.Second const challengeLen = 16 type auth struct { h *authHook secret string challenge struct { self string peer string } hash string recvChallenge chan struct{} recvHash chan struct{} tmr *time.Timer fail chan struct{} ok chan struct{} } var errDupChallenge = errors.New("peer repeats challenge") var errAuthFail = errors.New("peer auth fail") var errTimeout = errors.New("timeout") type authHook struct { m sync.Map } func (a *auth) generateChallenge() error { b := make([]byte, challengeLen) if _, err := rand.Read(b); err != nil { return err } 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, c) io.WriteString(h, a.secret) return string(h.Sum(nil)) } func (a *auth) wait(c chan struct{}) error { select { case <-a.tmr.C: return errTimeout case <-a.fail: return io.EOF case <-c: return nil } } func (a *auth) Send(rq, wq queue.Q) error { e := pack.NewEncoder(wq.Writer()) if err := a.generateChallenge(); err != nil { return err } defer a.deleteChallenge() e.Lps([]byte(a.challenge.self)) if err := a.wait(a.recvChallenge); err != nil { return err } if _, ok := a.h.m.Load(a.challenge.peer); ok { return errDupChallenge } e.Lps([]byte(a.getHash(a.challenge.peer))) if err := a.wait(a.recvHash); err != nil { return err } 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 { r := rq.Reader() d := pack.NewDecoder(r) if b, err := d.Lps(); err != nil { close(a.fail) return err } else { a.challenge.peer = string(b) } close(a.recvChallenge) 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 err } a.tmr.Stop() return queue.IoCopy(r, wq.Writer()) } func (h *authHook) Open(env env.Env) (interface{}, error) { a := &auth{ h: h, secret: env.GetLocal("secret"), recvChallenge: make(chan struct{}), recvHash: make(chan struct{}), fail: make(chan struct{}), ok: make(chan struct{}), tmr: time.NewTimer(authTimeout), } return a, nil } func newAuthHook(opts.Opts) (hook, error) { return &authHook{}, nil } func init() { register("auth", newAuthHook) }