log tcpdump

dumper/dumper.go

@@ -0,0 +1,62 @@
+package dumper
+
+import (
+	"bytes"
+	"encoding/hex"
+	"strings"
+	"sync"
+
+	"go.uber.org/zap"
+)
+
+// Dumper dumper struct
+type Dumper struct {
+	direction uint
+	logger    *zap.Logger
+	buf       *bytes.Buffer
+	mu        *sync.RWMutex
+}
+
+// New new handler
+func New(direction uint, logger *zap.Logger) *Dumper {
+	d := &Dumper{
+		direction: direction,
+		logger:    logger,
+		buf:       new(bytes.Buffer),
+		mu:        new(sync.RWMutex),
+	}
+	return d
+}
+
+// Write to dump
+func (d *Dumper) Write(p []byte) (n int, err error) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	d.buf.Write(p)
+	return len(p), nil
+}
+
+// Flush  flush buffer
+func (d *Dumper) Flush() {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	if d.buf.Len() == 0 {
+		return
+	}
+	hexdump := strings.Split(hex.Dump(d.buf.Bytes()), "\n")
+	d.buf.Truncate(0)
+	byteString := []string{}
+	ascii := []string{}
+	for _, hd := range hexdump {
+		if hd == "" {
+			continue
+		}
+		byteString = append(byteString, strings.TrimRight(strings.Replace(hd[10:58], "  ", " ", 1), " "))
+		ascii = append(ascii, hd[61:len(hd)-1])
+	}
+	d.logger.Info("dump",
+		zap.Uint("direction", d.direction),
+		zap.String("hex", strings.Join(byteString, " ")),
+		zap.String("ascii", strings.Join(ascii, "")),
+	)
+}

handler/handler.go

@@ -9,11 +9,18 @@ import (
 
 	"github.com/gorilla/mux"
 	"github.com/gorilla/websocket"
+	"github.com/kazeburo/wsgate-server/dumper"
 	"github.com/kazeburo/wsgate-server/mapping"
 	"github.com/kazeburo/wsgate-server/publickey"
 	"go.uber.org/zap"
 )
 
+var (
+	websocketUpstream   uint          = 1
+	upstreamWebsocket   uint          = 2
+	flushDumperInterval time.Duration = 300
+)
+
 // Handler handlers
 type Handler struct {
 	logger       *zap.Logger
@@ -22,6 +29,7 @@ type Handler struct {
 	writeTimeout time.Duration
 	mp           *mapping.Mapping
 	pk           *publickey.Publickey
+	dumpTCP      uint
 }
 
 // New new handler
@@ -31,6 +39,7 @@ func New(
 	writeTimeout time.Duration,
 	mp *mapping.Mapping,
 	pk *publickey.Publickey,
+	dumpTCP uint,
 	logger *zap.Logger) (*Handler, error) {
 
 	upgrader := websocket.Upgrader{
@@ -49,6 +58,7 @@ func New(
 		writeTimeout: writeTimeout,
 		mp:           mp,
 		pk:           pk,
+		dumpTCP:      dumpTCP,
 	}, nil
 }
 
@@ -114,8 +124,12 @@ func (h *Handler) Proxy() func(w http.ResponseWriter, r *http.Request) {
 		}
 
 		logger.Info("log", zap.String("status", "Connected"))
+		dr := dumper.New(websocketUpstream, logger)
+		ds := dumper.New(upstreamWebsocket, logger)
 
 		defer func() {
+			dr.Flush()
+			ds.Flush()
 			status := "Suceeded"
 			if hasError {
 				status = "Failed"
@@ -128,6 +142,20 @@ func (h *Handler) Proxy() func(w http.ResponseWriter, r *http.Request) {
 			)
 		}()
 
+		ticker := time.NewTicker(flushDumperInterval * time.Millisecond)
+		defer ticker.Stop()
+		go func() {
+			for {
+				select {
+				case <-r.Context().Done():
+					return
+				case _ = <-ticker.C:
+					dr.Flush()
+					ds.Flush()
+				}
+			}
+		}()
+
 		doneCh := make(chan bool)
 		goClose := false
 
@@ -157,6 +185,9 @@ func (h *Handler) Proxy() func(w http.ResponseWriter, r *http.Request) {
 					hasError = true
 					return
 				}
+				if h.dumpTCP > 0 {
+					r = io.TeeReader(r, dr)
+				}
 				n, err := io.Copy(s, r)
 				if err != nil {
 					if !goClose {
@@ -190,7 +221,9 @@ func (h *Handler) Proxy() func(w http.ResponseWriter, r *http.Request) {
 				}
 
 				b = b[:n]
-
+				if h.dumpTCP > 1 {
+					ds.Write(b)
+				}
 				if err := conn.WriteMessage(websocket.BinaryMessage, b); err != nil {
 					if !goClose {
 						logger.Warn("WriteMessage", zap.Error(err))

wsgate-server.go

@@ -26,6 +26,7 @@ var (
 	writeTimeout     = flag.Duration("write_timeout", 10*time.Second, "Write timeout.")
 	mapFile          = flag.String("map", "", "path and proxy host mapping file")
 	publicKeyFile    = flag.String("public-key", "", "public key for signing auth header")
+	dumpTCP          = flag.Uint("dump-tcp", 0, "Dump TCP. 0 = disable, 1 = src to dest, 2 = both")
 )
 
 func printVersion() {
@@ -62,6 +63,7 @@ func main() {
 		*writeTimeout,
 		mp,
 		pk,
+		*dumpTCP,
 		logger,
 	)
 	if err != nil {