Fix user registration

Author: antoniomika

pkg/apps/pastes/ssh.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"github.com/picosh/pico/pkg/db/postgres"
 	"github.com/picosh/pico/pkg/filehandlers"
@@ -60,7 +61,7 @@ func StartSshServer() {
 			scp.Middleware(handler),
 			rsync.Middleware(handler),
 			auth.Middleware(handler),
-			pssh.PtyMdw(pssh.DeprecatedNotice()),
+			pssh.PtyMdw(pssh.DeprecatedNotice(), 200*time.Millisecond),
 			pssh.LogMiddleware(handler, dbh),
 		},
 		[]pssh.SSHServerMiddleware{

pkg/apps/pico/ssh.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"git.sr.ht/~rockorager/vaxis"
 	"github.com/picosh/pico/pkg/db/postgres"
@@ -98,7 +99,7 @@ func StartSshServer() {
 						Dbpool:  handler.DBPool,
 						Logger:  cfg.Logger,
 					}
-					return pssh.PtyMdw(createTui(shrd))(next)(sesh)
+					return pssh.PtyMdw(createTui(shrd), 200*time.Millisecond)(next)(sesh)
 				}
 			},
 			Middleware(cliHandler),

pkg/apps/prose/ssh.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"github.com/picosh/pico/pkg/db/postgres"
 	"github.com/picosh/pico/pkg/filehandlers"
@@ -79,7 +80,7 @@ func StartSshServer() {
 			scp.Middleware(handler),
 			rsync.Middleware(handler),
 			auth.Middleware(handler),
-			pssh.PtyMdw(pssh.DeprecatedNotice()),
+			pssh.PtyMdw(pssh.DeprecatedNotice(), 200*time.Millisecond),
 			pssh.LogMiddleware(handler, dbh),
 		},
 		[]pssh.SSHServerMiddleware{

pkg/pobj/storage/minio.go

@@ -82,6 +82,7 @@ func (s *StorageMinio) GetBucket(name string) (Bucket, error) {
 		s.Logger.Info("bucket found in lru cache", "name", name)
 		return cachedBucket.Bucket, cachedBucket.Error
 	}
+
 	s.Logger.Info("bucket not found in lru cache", "name", name)
 
 	bucket := Bucket{
@@ -114,6 +115,8 @@ func (s *StorageMinio) UpsertBucket(name string) (Bucket, error) {
 		return bucket, err
 	}
 
+	s.BucketCache.Remove(name)
+
 	return bucket, nil
 }
 

pkg/pssh/logger.go

@@ -1,7 +1,6 @@
 package pssh
 
 import (
-	"fmt"
 	"log/slog"
 	"time"
 
@@ -31,18 +30,20 @@ func LogMiddleware(getLogger GetLoggerInterface, db FindUserInterface) SSHServer
 				user := GetUser(s)
 				if user == nil {
 					userID, ok := s.Permissions().Extensions["user_id"]
-					if !ok {
-						return fmt.Errorf("`user_id` not set in permissions")
-					}
-					user, err := db.FindUser(userID)
-					if err == nil && user != nil {
-						logger = logger.With(
-							"user", user.Name,
-							"userId", user.ID,
-							"ip", s.RemoteAddr().String(),
-						)
-						s.SetValue(ctxUserKey{}, user)
+					if ok {
+						user, err := db.FindUser(userID)
+						if err == nil && user != nil {
+							logger = logger.With(
+								"user", user.Name,
+								"userId", user.ID,
+								"ip", s.RemoteAddr().String(),
+							)
+							s.SetValue(ctxUserKey{}, user)
+						}
+					} else {
+						logger.Error("`user_id` not set in permissions")
 					}
+
 				}
 
 				s.SetValue(ctxLoggerKey{}, logger)

pkg/pssh/pty.go

@@ -3,6 +3,7 @@ package pssh
 import (
 	"encoding/binary"
 	"fmt"
+	"time"
 )
 
 func SessionMessage(sesh *SSHServerConnSession, msg string) {
@@ -23,9 +24,26 @@ func DeprecatedNotice() SSHServerMiddleware {
 	}
 }
 
-func PtyMdw(mdw SSHServerMiddleware) SSHServerMiddleware {
+func PtyMdw(mdw SSHServerMiddleware, waitTimeout time.Duration) SSHServerMiddleware {
 	return func(next SSHServerHandler) SSHServerHandler {
 		return func(sesh *SSHServerConnSession) error {
+			if waitTimeout > 0 {
+				afterTime := time.Now().Add(waitTimeout)
+
+				for {
+					if time.Now().After(afterTime) {
+						break
+					}
+
+					_, _, ok := sesh.Pty()
+					if ok {
+						break
+					}
+
+					time.Sleep(50 * time.Millisecond)
+				}
+			}
+
 			_, _, ok := sesh.Pty()
 			if !ok {
 				return next(sesh)

pkg/pssh/server.go

@@ -614,7 +614,8 @@ func NewSSHServerWithConfig(
 	app, host, port, promPort, hostKey string,
 	pubKeyAuthHandler PubKeyAuthHandler,
 	middleware, subsystemMiddleware []SSHServerMiddleware,
-	channelMiddleware map[string]SSHServerChannelMiddleware) (*SSHServer, error) {
+	channelMiddleware map[string]SSHServerChannelMiddleware,
+) (*SSHServer, error) {
 	server := NewSSHServer(ctx, logger, &SSHServerConfig{
 		App:        app,
 		ListenAddr: fmt.Sprintf("%s:%s", host, port),