updated ssh-broker to be more defensive

backend/internal/handlers/ssh-broker.go

@@ -170,7 +170,23 @@ func (api *API) logSSHSessionConnectStart(sessionItem models.SSHSession, profile
 
 func (api *API) logSSHSessionPrepare(sessionID string, remoteAddr string, user models.User, profile models.SSHAccessProfile) {
 	api.Logger.Write(logIDSSHBroker, codit_logger.LOG_INFO,
-		"connect prepare session=%s requester=%s user=%s profile=%s profile_name=%q server=%s server_name=%q target=%s:%d remote_user=%s auth=%s",
+		"connect prepare build session=%s requester=%s user=%s profile=%s profile_name=%q server=%s server_name=%q target=%s:%d remote_user=%s auth=%s",
+		sessionID,
+		remoteAddr,
+		user.Username,
+		profile.ID,
+		profile.Name,
+		profile.ServerID,
+		profile.Server.Name,
+		profile.Server.Host,
+		profile.Server.Port,
+		profile.RemoteUsername,
+		profile.AuthMethod)
+}
+
+func (api *API) logSSHSessionPreparedUse(sessionID string, remoteAddr string, user models.User, profile models.SSHAccessProfile) {
+	api.Logger.Write(logIDSSHBroker, codit_logger.LOG_INFO,
+		"connect prepare use session=%s requester=%s user=%s profile=%s profile_name=%q server=%s server_name=%q target=%s:%d remote_user=%s auth=%s",
 		sessionID,
 		remoteAddr,
 		user.Username,
@@ -1852,8 +1868,6 @@ func (api *API) StreamSSHSessionForSelf(w http.ResponseWriter, r *http.Request,
 	var user models.User
 	var ok bool
 	var item models.SSHSession
-	var prepared sshSessionPrepared
-	var preparedOK bool
 	var store *db.Store
 	var err error
 
@@ -1864,19 +1878,14 @@ func (api *API) StreamSSHSessionForSelf(w http.ResponseWriter, r *http.Request,
 		return
 	}
 	store = api.store(r)
-	prepared, preparedOK = api.SSHPreparedSessionStore.Get(params["id"])
-	if preparedOK {
-		item = prepared.Session
-	} else {
-		item, err = store.GetSSHSession(params["id"])
-		if err != nil {
-			if err == sql.ErrNoRows {
-				w.WriteHeader(http.StatusNotFound)
-				return
-			}
-			w.WriteHeader(http.StatusInternalServerError)
+	item, err = store.GetSSHSession(params["id"])
+	if err != nil {
+		if err == sql.ErrNoRows {
+			w.WriteHeader(http.StatusNotFound)
 			return
 		}
+		w.WriteHeader(http.StatusInternalServerError)
+		return
 	}
 	if item.UserID != user.ID {
 		w.WriteHeader(http.StatusForbidden)
@@ -2396,6 +2405,81 @@ func (api *API) prepareSSHSession(store *db.Store, user models.User, sessionItem
 	return prepared, "", nil
 }
 
+func (api *API) reconcileSSHPreparedSessionForUse(store *db.Store, sessionItem models.SSHSession, prepared sshSessionPrepared) (sshSessionPrepared, error) {
+	var err error
+	var i int
+	var reloadServer bool
+	var reloadHostKeys bool
+
+	if strings.TrimSpace(prepared.Session.ID) == "" {
+		return prepared, errors.New("prepared ssh session is missing session id")
+	}
+	if prepared.Session.ID != sessionItem.ID {
+		return prepared, fmt.Errorf("prepared ssh session id mismatch: prepared=%s session=%s", prepared.Session.ID, sessionItem.ID)
+	}
+	if strings.TrimSpace(prepared.Session.ServerID) != "" && prepared.Session.ServerID != sessionItem.ServerID {
+		return prepared, fmt.Errorf("prepared ssh session server mismatch: prepared=%s session=%s", prepared.Session.ServerID, sessionItem.ServerID)
+	}
+	if strings.TrimSpace(prepared.Profile.ID) != "" && prepared.Profile.ID != sessionItem.ProfileID {
+		return prepared, fmt.Errorf("prepared ssh profile mismatch: prepared=%s session=%s", prepared.Profile.ID, sessionItem.ProfileID)
+	}
+	if strings.TrimSpace(sessionItem.ServerID) == "" {
+		return prepared, errors.New("ssh session server_id is empty")
+	}
+
+	reloadServer = prepared.Profile.ServerID != sessionItem.ServerID || prepared.Profile.Server.ID != sessionItem.ServerID
+	if !reloadServer && strings.TrimSpace(sessionItem.Host) != "" {
+		reloadServer = prepared.Profile.Server.Host != sessionItem.Host || prepared.Profile.Server.Port != sessionItem.Port
+	}
+	if reloadServer {
+		api.Logger.Write(logIDSSHBroker, codit_logger.LOG_WARN,
+			"prepared session target mismatch session=%s profile=%s prepared_server=%s prepared_target=%s:%d session_server=%s session_target=%s:%d action=reload_selected_server",
+			sessionItem.ID,
+			sessionItem.ProfileID,
+			prepared.Profile.ServerID,
+			prepared.Profile.Server.Host,
+			prepared.Profile.Server.Port,
+			sessionItem.ServerID,
+			sessionItem.Host,
+			sessionItem.Port)
+		prepared.Profile.Server, err = store.GetSSHServer(sessionItem.ServerID)
+		if err != nil {
+			return prepared, err
+		}
+		prepared.Profile.ServerID = prepared.Profile.Server.ID
+		if strings.TrimSpace(sessionItem.Host) != "" {
+			prepared.Profile.Server.Host = sessionItem.Host
+		}
+		if sessionItem.Port > 0 {
+			prepared.Profile.Server.Port = sessionItem.Port
+		}
+		reloadHostKeys = true
+	}
+
+	for i = 0; i < len(prepared.HostKeys); i++ {
+		if prepared.HostKeys[i].ServerID != sessionItem.ServerID {
+			reloadHostKeys = true
+			break
+		}
+	}
+	if reloadHostKeys {
+		prepared.HostKeys, err = store.ListSSHServerHostKeys(sessionItem.ServerID)
+		if err != nil {
+			return prepared, err
+		}
+		prepared.PinHostKeyOnFirstUse = false
+		if len(prepared.HostKeys) == 0 {
+			if normalizeSSHServerHostKeyPolicy(prepared.Profile.Server.HostKeyPolicy) != "trust_on_first_use" {
+				return prepared, errors.New("no pinned host key for " + prepared.Profile.Server.Name)
+			}
+			prepared.PinHostKeyOnFirstUse = true
+		}
+	}
+
+	prepared.Session = sessionItem
+	return prepared, nil
+}
+
 func (api *API) serveSSHSessionStream(ws *websocket.Conn, store *db.Store, user models.User, sessionItem models.SSHSession) {
 	var sendMu sync.Mutex
 	var inputCh chan sshSessionStreamMessage
@@ -2517,20 +2601,15 @@ event_loop:
 				}
 				if attachments[msg.SessionID] != nil { continue } // session exists
 
-				prepared, preparedOK = api.SSHPreparedSessionStore.Get(msg.SessionID)
-				if preparedOK {
-					item = prepared.Session
-				} else {
-					item, err = store.GetSSHSession(msg.SessionID)
-					if err != nil {
-						api.sendSSHSessionWS(&sendMu, ws, sshSessionStreamMessage{
-							SessionID: msg.SessionID,
-							Type:      "status",
-							Status:    "error",
-							Message:   "ssh session not found",
-						})
-						continue
-					}
+				item, err = store.GetSSHSession(msg.SessionID)
+				if err != nil {
+					api.sendSSHSessionWS(&sendMu, ws, sshSessionStreamMessage{
+						SessionID: msg.SessionID,
+						Type:      "status",
+						Status:    "error",
+						Message:   "ssh session not found",
+					})
+					continue
 				}
 				if item.UserID != user.ID {
 					api.sendSSHSessionWS(&sendMu, ws, sshSessionStreamMessage{
@@ -2541,6 +2620,7 @@ event_loop:
 					})
 					continue
 				}
+				prepared, preparedOK = api.SSHPreparedSessionStore.Get(msg.SessionID)
 				if !preparedOK {
 					attachedStatus = "closed"
 					attachedMessage = item.Error
@@ -2728,10 +2808,25 @@ func (api *API) runSSHSessionStream(store *db.Store, user models.User, sessionIt
 		promptedOTPCode = promptedAuthInput.OTPCode
 	}
 	if prepared != nil {
+		*prepared, err = api.reconcileSSHPreparedSessionForUse(store, sessionItem, *prepared)
+		if err != nil {
+			api.Logger.Write(logIDSSHBroker, codit_logger.LOG_WARN,
+				"prepared session rejected session=%s requester=%s user=%s err=%v",
+				sessionItem.ID,
+				sessionItem.RemoteAddr,
+				user.Username,
+				err)
+			if sendStatus != nil {
+				sendStatus(sshSessionStreamMessage{Type: "status", Status: "error", Message: err.Error()})
+			}
+			now = time.Now().UTC().Unix()
+			_ = store.UpdateSSHSessionStatus(sessionItem.ID, "error", sessionItem.HostKeyFingerprint, sessionItem.ConnectedAt, now, err.Error())
+			return
+		}
 		profile = prepared.Profile
 		hostKeys = prepared.HostKeys
 		authMethods = prepared.AuthMethods
-		api.logSSHSessionPrepare(sessionItem.ID, sessionItem.RemoteAddr, user, profile)
+		api.logSSHSessionPreparedUse(sessionItem.ID, sessionItem.RemoteAddr, user, profile)
 	} else {
 		prepared = &sshSessionPrepared{}
 		*prepared, prepareStage, err = api.prepareSSHSession(store, user, sessionItem, promptedPassword, promptedOTPCode)