Merge pull request #129 from jumpserver/dev_bugfix

[Bugfix] fix bugs and add net pprof debug handlers

pkg/handler/session.go

@@ -116,10 +116,20 @@ func (h *interactiveHandler) watchWinSizeChange() {
 }
 
 func (h *interactiveHandler) pauseWatchWinSize() {
+	select {
+	case <-h.sess.Sess.Context().Done():
+		return
+	default:
+	}
 	h.winWatchChan <- false
 }
 
 func (h *interactiveHandler) resumeWatchWinSize() {
+	select {
+	case <-h.sess.Sess.Context().Done():
+		return
+	default:
+	}
 	h.winWatchChan <- true
 }
 

pkg/handler/wrappersession.go

@@ -16,10 +16,13 @@ type WrapperSession struct {
 	inWriter  io.WriteCloser
 	outReader io.ReadCloser
 	mux       *sync.RWMutex
+
+	closed chan struct{}
 }
 
 func (w *WrapperSession) initial() {
 	w.Uuid = uuid.NewV4().String()
+	w.closed = make(chan struct{})
 	w.initReadPip()
 	go w.readLoop()
 }
@@ -38,11 +41,19 @@ func (w *WrapperSession) readLoop() {
 			break
 		}
 	}
+	w.mux.RLock()
 	_ = w.inWriter.Close()
+	w.mux.RUnlock()
+	close(w.closed)
 	logger.Infof("Request %s: Read loop break", w.Uuid)
 }
 
 func (w *WrapperSession) Read(p []byte) (int, error) {
+	select {
+	case <-w.closed:
+		return 0, io.EOF
+	default:
+	}
 	w.mux.RLock()
 	defer w.mux.RUnlock()
 	return w.outReader.Read(p)

pkg/httpd/server.go

@@ -5,8 +5,11 @@ import (
 	"net"
 	"net/http"
 	"path/filepath"
+	"strings"
 	"time"
 
+	"net/http/pprof"
+
 	"github.com/gorilla/mux"
 	gorillaws "github.com/gorilla/websocket"
 	"github.com/jumpserver/koko/pkg/config"
@@ -105,7 +108,9 @@ func StartHTTPServer() {
 	//router.HandleFunc("/coco/elfinder/sftp/", AuthDecorator(sftpFinder))
 	//router.HandleFunc("/coco/elfinder/sftp/connector/{host}/",
 	//	AuthDecorator(sftpHostConnectorView)).Methods("GET", "POST")
-
+	if strings.ToUpper(conf.LogLevel) == "DEBUG" {
+		router.PathPrefix("/debug/pprof/").HandlerFunc(pprof.Index)
+	}
 	addr := net.JoinHostPort(conf.BindHost, conf.HTTPPort)
 	logger.Info("Start HTTP server at ", addr)
 	httpServer = &http.Server{Addr: addr, Handler: router}