diff options
Diffstat (limited to 'pkg')
| -rw-r--r-- | pkg/server/socket/socket.go | 29 | ||||
| -rw-r--r-- | pkg/server/tunnel.go | 32 |
2 files changed, 39 insertions, 22 deletions
diff --git a/pkg/server/socket/socket.go b/pkg/server/socket/socket.go index 0060ce6..1bb7549 100644 --- a/pkg/server/socket/socket.go +++ b/pkg/server/socket/socket.go @@ -14,8 +14,12 @@ import ( var errAlreadyClosed = errors.New("already closed") +type exportChannel struct { + info string + Channel +} + type Channel interface { - Origin() string Send(wq queue.Q) error Recv(rq queue.Q) error Close() error @@ -36,7 +40,6 @@ type dialSocket struct { } type connChannel struct { - origin string conn net.Conn once sync.Once } @@ -47,8 +50,12 @@ type loopChannel struct { q queue.Q } -func newConnChannel(origin string, conn net.Conn) Channel { - return &connChannel{origin: origin, conn: conn} +func (c exportChannel) String() string { + return c.info +} + +func newConnChannel(conn net.Conn) Channel { + return &connChannel{conn: conn} } func (c *connChannel) final(f func() error, err error) error { @@ -63,10 +70,6 @@ func (c *connChannel) final(f func() error, err error) error { return err } -func (c *connChannel) Origin() string { - return c.origin -} - func (c *connChannel) Send(wq queue.Q) error { err := queue.IoCopy(c.conn, wq.Writer()) return c.final(c.Close, err) @@ -121,9 +124,9 @@ func (s *listenSocket) Open(env env.Env) (Channel, error) { } addr := conn.RemoteAddr() - origin := fmt.Sprintf("%s/%s", addr.Network(), addr) + info := fmt.Sprintf("%s/%s", addr.Network(), addr) - return newConnChannel(origin, conn), nil + return exportChannel{info, newConnChannel(conn)}, nil } func (s *listenSocket) String() string { @@ -154,16 +157,12 @@ func (s *dialSocket) Open(env env.Env) (Channel, error) { if err != nil { return nil, err } - return newConnChannel("-", conn), nil + return exportChannel{"-", newConnChannel(conn)}, nil } func (s *dialSocket) Close() { } -func (c *loopChannel) Origin() string { - return "loop" -} - func (c *loopChannel) Send(wq queue.Q) error { return queue.Copy(c.q, wq) } diff --git a/pkg/server/tunnel.go b/pkg/server/tunnel.go index 01d7c22..72d4c13 100644 --- a/pkg/server/tunnel.go +++ b/pkg/server/tunnel.go @@ -56,6 +56,8 @@ type tunnel struct { quit chan struct{} done chan struct{} + mono bool + in, out socket.S hooks []hook.H @@ -123,6 +125,11 @@ func (t *tunnel) serve() { t.handle(in) wg.Done() }() + + if t.mono { + wg.Wait() + t.wg.Wait() + } } } @@ -240,7 +247,7 @@ func (s *stream) channel(c socket.Channel, m *metric, rq, wq queue.Q) { q := queue.New() go counter(&m.rx, rq, q) watch(q, c.Recv) - rq.Dry() + q.Dry() }() } @@ -311,7 +318,7 @@ func parseHooks(args []string, env env.Env) ([]hook.H, error) { return hooks, nil } -func newTunnel(args []string, env env.Env) (*tunnel, error) { +func newTunnel(mono bool, args []string, env env.Env) (*tunnel, error) { var in, out socket.S var hooks []hook.H var err error @@ -337,6 +344,7 @@ func newTunnel(args []string, env env.Env) (*tunnel, error) { args: strings.Join(args, " "), quit: make(chan struct{}), done: make(chan struct{}), + mono: mono, hooks: hooks, in: in, out: out, @@ -356,9 +364,10 @@ func isOkTunnelName(s string) bool { func tunnelAdd(r *request) { args := r.args name := "" + mono := false - if len(args) >= 2 { - if args[0] == "name" { + for len(args) > 1 { + if args[0] == "name" && len(args) > 1 { name = args[1] if !isOkTunnelName(name) { r.Fatal("bad name") @@ -369,14 +378,23 @@ func tunnelAdd(r *request) { } args = args[2:] + continue } + + if args[0] == "mono" { + mono = true + args = args[1:] + continue + } + + break } if len(args) < 2 { r.Fatal("not enough args") } - t, err := newTunnel(args, r.c.s.env) + t, err := newTunnel(mono, args, r.c.s.env) if err != nil { r.Fatal(err) } @@ -459,7 +477,7 @@ func showActive(r *request) { defer t.mu.Unlock() foreachStream(t.streams, func(s *stream) { - r.Println(t.id, s.id, s.in.Origin(), s.out.Origin(), s.info()) + r.Println(t.id, s.id, s.in, s.out, s.info()) }) }) } @@ -479,7 +497,7 @@ func showRecent(r *request) { for _, s := range streams { when := s.until.Format(config.TimeFormat) - r.Println(when, s.t.id, s.id, s.in.Origin(), s.out.Origin(), s.info()) + r.Println(when, s.t.id, s.id, s.in, s.out, s.info()) } } |