fix: restore previous session on coder server --dev (#1821)

f0ssel · web-flow · commit 8a5277e29105 · 2022-05-27T17:02:02.000Z
diff --git a/cli/server.go b/cli/server.go
@@ -336,10 +336,11 @@ func server() *cobra.Command {
 						return xerrors.Errorf("generate random admin password for dev: %w", err)
 					}
 				}
-				err = createFirstUser(cmd, client, config, devUserEmail, devUserPassword)
+				restorePreviousSession, err := createFirstUser(logger, cmd, client, config, devUserEmail, devUserPassword)
 				if err != nil {
 					return xerrors.Errorf("create first user: %w", err)
 				}
+				defer restorePreviousSession()
 				_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "email: %s\n", devUserEmail)
 				_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "password: %s\n", devUserPassword)
 				_, _ = fmt.Fprintln(cmd.ErrOrStderr())
@@ -518,12 +519,14 @@ func server() *cobra.Command {
 	return root
 }
 
-func createFirstUser(cmd *cobra.Command, client *codersdk.Client, cfg config.Root, email, password string) error {
+// createFirstUser creates the first user and sets a valid session.
+// Caller must call restorePreviousSession on server exit.
+func createFirstUser(logger slog.Logger, cmd *cobra.Command, client *codersdk.Client, cfg config.Root, email, password string) (func(), error) {
 	if email == "" {
-		return xerrors.New("email is empty")
+		return nil, xerrors.New("email is empty")
 	}
 	if password == "" {
-		return xerrors.New("password is empty")
+		return nil, xerrors.New("password is empty")
 	}
 	_, err := client.CreateFirstUser(cmd.Context(), codersdk.CreateFirstUserRequest{
 		Email:            email,
@@ -532,26 +535,63 @@ func createFirstUser(cmd *cobra.Command, client *codersdk.Client, cfg config.Roo
 		OrganizationName: "acme-corp",
 	})
 	if err != nil {
-		return xerrors.Errorf("create first user: %w", err)
+		return nil, xerrors.Errorf("create first user: %w", err)
 	}
 	token, err := client.LoginWithPassword(cmd.Context(), codersdk.LoginWithPasswordRequest{
 		Email:    email,
 		Password: password,
 	})
 	if err != nil {
-		return xerrors.Errorf("login with first user: %w", err)
+		return nil, xerrors.Errorf("login with first user: %w", err)
 	}
 	client.SessionToken = token.SessionToken
 
+	// capture the current session and if exists recover session on server exit
+	restorePreviousSession := func() {}
+	oldURL, _ := cfg.URL().Read()
+	oldSession, _ := cfg.Session().Read()
+	if oldURL != "" && oldSession != "" {
+		restorePreviousSession = func() {
+			currentURL, err := cfg.URL().Read()
+			if err != nil {
+				logger.Error(cmd.Context(), "failed to read current session url", slog.Error(err))
+				return
+			}
+			currentSession, err := cfg.Session().Read()
+			if err != nil {
+				logger.Error(cmd.Context(), "failed to read current session token", slog.Error(err))
+				return
+			}
+
+			// if it's changed since we wrote to it don't restore session
+			if currentURL != client.URL.String() ||
+				currentSession != token.SessionToken {
+				return
+			}
+
+			err = cfg.URL().Write(oldURL)
+			if err != nil {
+				logger.Error(cmd.Context(), "failed to recover previous session url", slog.Error(err))
+				return
+			}
+			err = cfg.Session().Write(oldSession)
+			if err != nil {
+				logger.Error(cmd.Context(), "failed to recover previous session token", slog.Error(err))
+				return
+			}
+		}
+	}
+
 	err = cfg.URL().Write(client.URL.String())
 	if err != nil {
-		return xerrors.Errorf("write local url: %w", err)
+		return nil, xerrors.Errorf("write local url: %w", err)
 	}
 	err = cfg.Session().Write(token.SessionToken)
 	if err != nil {
-		return xerrors.Errorf("write session token: %w", err)
+		return nil, xerrors.Errorf("write session token: %w", err)
 	}
-	return nil
+
+	return restorePreviousSession, nil
 }
 
 // nolint:revive
