Enable panic recovery in ssh and set logger to be based on debug

picosh · Nov 15, 2023 · 5a1f3d8 · 5a1f3d8
1 parent f0191d0
commit 5a1f3d8
Show file tree

Hide file tree

Showing 12 changed files with 50 additions and 12 deletions.
diff --git a/auth/auth.go b/auth/auth.go
@@ -289,7 +289,7 @@ func StartApiServer() {
 		Port:   shared.GetEnv("AUTH_WEB_PORT", "3000"),
 	}
 
-	logger := shared.CreateLogger()
+	logger := shared.CreateLogger(true)
 	db := postgres.NewDB(cfg.DbURL, logger)
 	defer db.Close()
 

diff --git a/feeds/config.go b/feeds/config.go
@@ -50,7 +50,7 @@ func NewConfigSite() *shared.ConfigSite {
 			Space:         "feeds",
 			AllowedExt:    []string{".txt"},
 			HiddenPosts:   []string{"_header.txt", "_readme.txt"},
-			Logger:        shared.CreateLogger(),
+			Logger:        shared.CreateLogger(debug == "1"),
 			AllowRegister: allowRegister == "1",
 		},
 	}

diff --git a/imgs/config.go b/imgs/config.go
@@ -66,7 +66,7 @@ func NewConfigSite() *shared.ConfigSite {
 			IntroText:     intro,
 			Space:         "imgs",
 			AllowedExt:    []string{".jpg", ".jpeg", ".png", ".gif", ".webp", ".svg"},
-			Logger:        shared.CreateLogger(),
+			Logger:        shared.CreateLogger(debug == "1"),
 			AllowRegister: allowRegister == "1",
 		},
 	}

diff --git a/lists/config.go b/lists/config.go
@@ -49,7 +49,7 @@ func NewConfigSite() *shared.ConfigSite {
 			Space:         "lists",
 			AllowedExt:    []string{".txt"},
 			HiddenPosts:   []string{"_header.txt", "_readme.txt"},
-			Logger:        shared.CreateLogger(),
+			Logger:        shared.CreateLogger(debug == "1"),
 			AllowRegister: allowRegister == "1",
 		},
 	}

diff --git a/pastes/config.go b/pastes/config.go
@@ -45,7 +45,7 @@ func NewConfigSite() *shared.ConfigSite {
 			Description:   "a pastebin for hackers.",
 			IntroText:     intro,
 			Space:         "pastes",
-			Logger:        shared.CreateLogger(),
+			Logger:        shared.CreateLogger(debug == "1"),
 			AllowRegister: allowRegister == "1",
 		},
 	}

diff --git a/pgs/config.go b/pgs/config.go
@@ -79,7 +79,7 @@ func NewConfigSite() *shared.ConfigSite {
 			},
 			MaxSize:       maxSize,
 			MaxAssetSize:  maxAssetSize,
-			Logger:        shared.CreateLogger(),
+			Logger:        shared.CreateLogger(debug == "1"),
 			AllowRegister: allowRegister == "1",
 		},
 	}

diff --git a/prose/config.go b/prose/config.go
@@ -49,7 +49,7 @@ func NewConfigSite() *shared.ConfigSite {
 			Space:         "prose",
 			AllowedExt:    []string{".md"},
 			HiddenPosts:   []string{"_readme.md", "_styles.css", "_footer.md"},
-			Logger:        shared.CreateLogger(),
+			Logger:        shared.CreateLogger(debug == "1"),
 			AllowRegister: allowRegister == "1",
 		},
 	}

diff --git a/shared/config.go b/shared/config.go
@@ -251,8 +251,18 @@ func (c *ConfigSite) AssetURL(username, projectName, fpath string) string {
 	)
 }
 
-func CreateLogger() *zap.SugaredLogger {
-	logger, err := zap.NewProduction()
+func CreateLogger(debug bool) *zap.SugaredLogger {
+	var (
+		err    error
+		logger *zap.Logger
+	)
+
+	if debug {
+		logger, err = zap.NewDevelopment()
+	} else {
+		logger, err = zap.NewProduction()
+	}
+
 	if err != nil {
 		log.Fatal(err)
 	}

diff --git a/wish/send/auth/auth.go b/wish/send/auth/auth.go
@@ -9,6 +9,13 @@ import (
 func Middleware(writeHandler utils.CopyFromClientHandler) wish.Middleware {
 	return func(sshHandler ssh.Handler) ssh.Handler {
 		return func(session ssh.Session) {
+			defer func() {
+				if r := recover(); r != nil {
+					writeHandler.GetLogger().Error("error running auth middleware: ", r)
+					_, _ = session.Stderr().Write([]byte("error running auth middleware\r\n"))
+				}
+			}()
+
 			err := writeHandler.Validate(session)
 			if err != nil {
 				utils.ErrorHandler(session, err)

diff --git a/wish/send/rsync/rsync.go b/wish/send/rsync/rsync.go
@@ -152,6 +152,13 @@ func (h *handler) Put(file *rsyncutils.ReceiverFile) (int64, error) {
 func Middleware(writeHandler utils.CopyFromClientHandler) wish.Middleware {
 	return func(sshHandler ssh.Handler) ssh.Handler {
 		return func(session ssh.Session) {
+			defer func() {
+				if r := recover(); r != nil {
+					writeHandler.GetLogger().Error("error running rsync middleware: ", r)
+					_, _ = session.Stderr().Write([]byte("error running rsync middleware, check the flags you are using\r\n"))
+				}
+			}()
+
 			cmd := session.Command()
 			if len(cmd) == 0 || cmd[0] != "rsync" {
 				sshHandler(session)
@@ -193,7 +200,7 @@ func Middleware(writeHandler utils.CopyFromClientHandler) wish.Middleware {
 					}
 
 					if err := rsyncsender.ClientRun(opts, session, fileHandler, fileHandler.root, true); err != nil {
-						writeHandler.GetLogger().Error("error running rsync sender:", err)
+						writeHandler.GetLogger().Error("error running rsync sender: ", err)
 					}
 					return
 				}
@@ -224,7 +231,7 @@ func Middleware(writeHandler utils.CopyFromClientHandler) wish.Middleware {
 			}
 
 			if _, err := rsyncreceiver.ClientRun(opts, session, fileHandler, true); err != nil {
-				writeHandler.GetLogger().Error("error running rsync receiver:", err)
+				writeHandler.GetLogger().Error("error running rsync receiver: ", err)
 			}
 		}
 	}

diff --git a/wish/send/scp/scp.go b/wish/send/scp/scp.go
@@ -11,6 +11,13 @@ import (
 func Middleware(writeHandler utils.CopyFromClientHandler) wish.Middleware {
 	return func(sshHandler ssh.Handler) ssh.Handler {
 		return func(session ssh.Session) {
+			defer func() {
+				if r := recover(); r != nil {
+					writeHandler.GetLogger().Error("error running scp middleware: ", r)
+					_, _ = session.Stderr().Write([]byte("error running scp middleware, check the flags you are using\r\n"))
+				}
+			}()
+
 			cmd := session.Command()
 			if len(cmd) == 0 || cmd[0] != "scp" {
 				sshHandler(session)

diff --git a/wish/send/sftp/sftp.go b/wish/send/sftp/sftp.go
@@ -22,6 +22,13 @@ func SSHOption(writeHandler utils.CopyFromClientHandler) ssh.Option {
 
 func SubsystemHandler(writeHandler utils.CopyFromClientHandler) ssh.SubsystemHandler {
 	return func(session ssh.Session) {
+		defer func() {
+			if r := recover(); r != nil {
+				writeHandler.GetLogger().Error("error running sftp middleware: ", r)
+				_, _ = session.Stderr().Write([]byte("error running sftp middleware, check the flags you are using\r\n"))
+			}
+		}()
+
 		err := writeHandler.Validate(session)
 		if err != nil {
 			utils.ErrorHandler(session, err)
@@ -44,7 +51,7 @@ func SubsystemHandler(writeHandler utils.CopyFromClientHandler) ssh.SubsystemHan
 
 		err = requestServer.Serve()
 		if err != nil && !errors.Is(err, io.EOF) {
-			writeHandler.GetLogger().Error("Error serving sftp subsystem:", err)
+			writeHandler.GetLogger().Error("Error serving sftp subsystem: ", err)
 		}
 	}
 }