[ssh-gateway] Reject ssh connection when workspace not found

Author: iQQBot

components/ws-proxy/pkg/sshproxy/server.go

@@ -43,18 +43,19 @@ var (
 )
 
 var (
-	ErrWorkspaceNotFound  = NewSSHError("WS_NOTFOUND", "not found workspace")
-	ErrWorkspaceIDInvalid = NewSSHError("WS_ID_INVALID", "workspace id invalid")
+	ErrWorkspaceNotFound  = NewSSHErrorWithError("WS_NOTFOUND", "not found workspace", ssh.ErrDenied)
+	ErrWorkspaceIDInvalid = NewSSHErrorWithError("WS_ID_INVALID", "workspace id invalid", ssh.ErrDenied)
 	ErrAuthFailed         = NewSSHError("AUTH_FAILED", "auth failed")
-	ErrUsernameFormat     = NewSSHError("USER_FORMAT", "username format is not correct")
-	ErrMissPrivateKey     = NewSSHError("MISS_KEY", "missing privateKey")
+	ErrUsernameFormat     = NewSSHErrorWithError("USER_FORMAT", "username format is not correct", ssh.ErrDenied)
+	ErrMissPrivateKey     = NewSSHErrorWithError("MISS_KEY", "missing privateKey", ssh.ErrDenied)
 	ErrConnFailed         = NewSSHError("CONN_FAILED", "cannot to connect with workspace")
 	ErrCreateSSHKey       = NewSSHError("CREATE_KEY_FAILED", "cannot create private pair in workspace")
 )
 
 type SSHError struct {
 	shortName   string
 	description string
+	err         error
 }
 
 func (e SSHError) Error() string {
@@ -64,11 +65,18 @@ func (e SSHError) Error() string {
 func (e SSHError) ShortName() string {
 	return e.shortName
 }
+func (e SSHError) Unwrap() error {
+	return e.err
+}
 
 func NewSSHError(shortName string, description string) SSHError {
 	return SSHError{shortName: shortName, description: description}
 }
 
+func NewSSHErrorWithError(shortName string, description string, err error) SSHError {
+	return SSHError{shortName: shortName, description: description, err: err}
+}
+
 type Session struct {
 	Conn *ssh.ServerConn
 
@@ -110,15 +118,15 @@ func New(signers []ssh.Signer, workspaceInfoProvider p.WorkspaceInfoProvider, he
 		NoClientAuth:  true,
 		NoClientAuthCallback: func(conn ssh.ConnMetadata) (*ssh.Permissions, error) {
 			args := strings.Split(conn.User(), "#")
-			// NoClientAuthCallback only support workspaceId#ownerToken
-			if len(args) != 2 {
-				return nil, ssh.ErrNoAuth
-			}
 			workspaceId := args[0]
 			wsInfo, err := server.GetWorkspaceInfo(workspaceId)
 			if err != nil {
 				return nil, err
 			}
+			// NoClientAuthCallback only support workspaceId#ownerToken
+			if len(args) != 2 {
+				return nil, ssh.ErrNoAuth
+			}
 			if wsInfo.Auth.OwnerToken != args[1] {
 				return nil, ssh.ErrNoAuth
 			}