package hook import ( "bytes" "encoding/hex" "fmt" "os" "path" "time" "tunnel/pkg/config" "tunnel/pkg/server/env" "tunnel/pkg/server/opts" "tunnel/pkg/server/queue" ) const teeDefaultFile = "/tmp/tunnel/dump" type tee struct { f *os.File h *teeHook } type teeHook struct { file string time bool } func (t *tee) dump(s string, p []byte) error { var out bytes.Buffer if t.h.time { now := time.Now().Format(config.TimeMsFormat) fmt.Fprintln(&out, now, s, len(p)) } else { fmt.Fprintln(&out, s, len(p)) } w := hex.Dumper(&out) w.Write(p) w.Close() if _, err := t.f.Write(out.Bytes()); err != nil { return err } return nil } func (t *tee) Send(rq, wq queue.Q) error { for b := range rq { t.dump(">", b) wq <- b } return nil } func (t *tee) Recv(rq, wq queue.Q) error { for b := range rq { t.dump("<", b) wq <- b } return nil } func (t *tee) Close() { t.f.Close() } func (h *teeHook) where(env env.Env) string { if h.file != "" { return h.file } if v := env.GetLocal("tee.file"); v != "" { return v } return teeDefaultFile } func (h *teeHook) Open(env env.Env) (interface{}, error) { file := h.where(env) dir := path.Dir(file) if err := os.MkdirAll(dir, 0755); err != nil { return nil, err } tid, sid := env.Get("tunnel"), env.Get("stream") name := fmt.Sprintf("%s.%s.%s", file, tid, sid) t := &tee{h: h} if f, err := os.Create(name); err != nil { return nil, err } else { t.f = f } return t, nil } func newTeeHook(opts opts.Opts) (hook, error) { h := &teeHook{ file: opts["file"], time: opts.Bool("time"), } return h, nil } func init() { register("tee", newTeeHook) }