Show service banner before MOTD

coder · code-asher · Jun 30, 2023 · Jun 23, 2023 · Jun 23, 2023 · Jun 27, 2023
commit a40e38728c344fe7a085f8ab8feab8f2a889fb8d
diff --git a/agent/agent.go b/agent/agent.go
@@ -192,6 +192,7 @@ func (a *agent) init(ctx context.Context) {
 	sshSrv.Env = a.envVars
 	sshSrv.AgentToken = func() string { return *a.sessionToken.Load() }
 	sshSrv.Manifest = &a.manifest
+	sshSrv.GetServiceBanner = a.client.GetServiceBanner
 	a.sshServer = sshSrv
 
 	go a.runLoop(ctx)

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -191,7 +191,7 @@ func TestAgent_Stats_Magic(t *testing.T) {
 
 func TestAgent_SessionExec(t *testing.T) {
 	t.Parallel()
-	session := setupSSHSession(t, agentsdk.Manifest{})
+	session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{})
 
 	command := "echo test"
 	if runtime.GOOS == "windows" {
@@ -204,7 +204,7 @@ func TestAgent_SessionExec(t *testing.T) {
 
 func TestAgent_GitSSH(t *testing.T) {
 	t.Parallel()
-	session := setupSSHSession(t, agentsdk.Manifest{})
+	session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{})
 	command := "sh -c 'echo $GIT_SSH_COMMAND'"
 	if runtime.GOOS == "windows" {
 		command = "cmd.exe /c echo %GIT_SSH_COMMAND%"
@@ -224,7 +224,7 @@ func TestAgent_SessionTTYShell(t *testing.T) {
 		// it seems like it could be either.
 		t.Skip("ConPTY appears to be inconsistent on Windows.")
 	}
-	session := setupSSHSession(t, agentsdk.Manifest{})
+	session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{})
 	command := "sh"
 	if runtime.GOOS == "windows" {
 		command = "cmd.exe"
@@ -247,7 +247,7 @@ func TestAgent_SessionTTYShell(t *testing.T) {
 
 func TestAgent_SessionTTYExitCode(t *testing.T) {
 	t.Parallel()
-	session := setupSSHSession(t, agentsdk.Manifest{})
+	session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{})
 	command := "areallynotrealcommand"
 	err := session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
 	require.NoError(t, err)
@@ -277,6 +277,7 @@ func TestAgent_Session_TTY_MOTD(t *testing.T) {
 	}
 
 	wantMOTD := "Welcome to your Coder workspace!"
+	wantServiceBanner := "Service banner text goes here"
 
 	tmpdir := t.TempDir()
 	name := filepath.Join(tmpdir, "motd")
@@ -286,25 +287,84 @@ func TestAgent_Session_TTY_MOTD(t *testing.T) {
 	// Set HOME so we can ensure no ~/.hushlogin is present.
 	t.Setenv("HOME", tmpdir)
 
-	session := setupSSHSession(t, agentsdk.Manifest{
-		MOTDFile: name,
-	})
-	err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
-	require.NoError(t, err)
+	tests := []struct {
+		name       string
+		manifest   agentsdk.Manifest
+		banner     codersdk.ServiceBannerConfig
+		expected   []string
+		unexpected []string
+	}{
+		{
+			name:       "WithoutServiceBanner",
+			manifest:   agentsdk.Manifest{MOTDFile: name},
+			banner:     codersdk.ServiceBannerConfig{},
+			expected:   []string{wantMOTD},
+			unexpected: []string{wantServiceBanner},
+		},
+		{
+			name:     "WithServiceBanner",
+			manifest: agentsdk.Manifest{MOTDFile: name},
+			banner: codersdk.ServiceBannerConfig{
+				Enabled: true,
+				Message: wantServiceBanner,
+			},
+			expected: []string{wantMOTD, wantServiceBanner},
+		},
+		{
+			name:     "ServiceBannerDisabled",
+			manifest: agentsdk.Manifest{MOTDFile: name},
+			banner: codersdk.ServiceBannerConfig{
+				Enabled: false,
+				Message: wantServiceBanner,
+			},
+			expected:   []string{wantMOTD},
+			unexpected: []string{wantServiceBanner},
+		},
+		{
+			name:     "ServiceBannerOnly",
+			manifest: agentsdk.Manifest{},
+			banner: codersdk.ServiceBannerConfig{
+				Enabled: true,
+				Message: wantServiceBanner,
+			},
+			expected:   []string{wantServiceBanner},
+			unexpected: []string{wantMOTD},
+		},
+		{
+			name:       "None",
+			manifest:   agentsdk.Manifest{},
+			banner:     codersdk.ServiceBannerConfig{},
+			unexpected: []string{wantServiceBanner, wantMOTD},
+		},
+	}
 
-	ptty := ptytest.New(t)
-	var stdout bytes.Buffer
-	session.Stdout = &stdout
-	session.Stderr = ptty.Output()
-	session.Stdin = ptty.Input()
-	err = session.Shell()
-	require.NoError(t, err)
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			session := setupSSHSession(t, test.manifest, test.banner)
+			err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
+			require.NoError(t, err)
 
-	ptty.WriteLine("exit 0")
-	err = session.Wait()
-	require.NoError(t, err)
+			ptty := ptytest.New(t)
+			var stdout bytes.Buffer
+			session.Stdout = &stdout
+			session.Stderr = ptty.Output()
+			session.Stdin = ptty.Input()
+			err = session.Shell()
+			require.NoError(t, err)
+
+			ptty.WriteLine("exit 0")
+			err = session.Wait()
+			require.NoError(t, err)
 
-	require.Contains(t, stdout.String(), wantMOTD, "should show motd")
+			for _, unexpected := range test.unexpected {
+				require.NotContains(t, stdout.String(), unexpected, "should not show motd")
+			}
+			for _, expect := range test.expected {
+				require.Contains(t, stdout.String(), expect, "should show motd")
+			}
+		})
+	}
 }
 
 //nolint:paralleltest // This test sets an environment variable.
@@ -317,6 +377,7 @@ func TestAgent_Session_TTY_Hushlogin(t *testing.T) {
 	}
 
 	wantNotMOTD := "Welcome to your Coder workspace!"
+	wantNotServiceBanner := "Service banner text goes here"
 
 	tmpdir := t.TempDir()
 	name := filepath.Join(tmpdir, "motd")
@@ -334,6 +395,9 @@ func TestAgent_Session_TTY_Hushlogin(t *testing.T) {
 
 	session := setupSSHSession(t, agentsdk.Manifest{
 		MOTDFile: name,
+	}, codersdk.ServiceBannerConfig{
+		Enabled: true,
+		Message: wantNotServiceBanner,
 	})
 	err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
 	require.NoError(t, err)
@@ -351,6 +415,7 @@ func TestAgent_Session_TTY_Hushlogin(t *testing.T) {
 	require.NoError(t, err)
 
 	require.NotContains(t, stdout.String(), wantNotMOTD, "should not show motd")
+	require.NotContains(t, stdout.String(), wantNotServiceBanner, "should not show service banner")
 }
 
 func TestAgent_Session_TTY_FastCommandHasOutput(t *testing.T) {
@@ -797,7 +862,7 @@ func TestAgent_EnvironmentVariables(t *testing.T) {
 		EnvironmentVariables: map[string]string{
 			key: value,
 		},
-	})
+	}, codersdk.ServiceBannerConfig{})
 	command := "sh -c 'echo $" + key + "'"
 	if runtime.GOOS == "windows" {
 		command = "cmd.exe /c echo %" + key + "%"
@@ -814,7 +879,7 @@ func TestAgent_EnvironmentVariableExpansion(t *testing.T) {
 		EnvironmentVariables: map[string]string{
 			key: "$SOMETHINGNOTSET",
 		},
-	})
+	}, codersdk.ServiceBannerConfig{})
 	command := "sh -c 'echo $" + key + "'"
 	if runtime.GOOS == "windows" {
 		command = "cmd.exe /c echo %" + key + "%"
@@ -837,7 +902,7 @@ func TestAgent_CoderEnvVars(t *testing.T) {
 		t.Run(key, func(t *testing.T) {
 			t.Parallel()
 
-			session := setupSSHSession(t, agentsdk.Manifest{})
+			session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{})
 			command := "sh -c 'echo $" + key + "'"
 			if runtime.GOOS == "windows" {
 				command = "cmd.exe /c echo %" + key + "%"
@@ -860,7 +925,7 @@ func TestAgent_SSHConnectionEnvVars(t *testing.T) {
 		t.Run(key, func(t *testing.T) {
 			t.Parallel()
 
-			session := setupSSHSession(t, agentsdk.Manifest{})
+			session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{})
 			command := "sh -c 'echo $" + key + "'"
 			if runtime.GOOS == "windows" {
 				command = "cmd.exe /c echo %" + key + "%"
@@ -1666,11 +1731,18 @@ func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) (*pt
 	return ptytest.Start(t, cmd)
 }
 
-func setupSSHSession(t *testing.T, options agentsdk.Manifest) *ssh.Session {
+func setupSSHSession(
+	t *testing.T,
+	options agentsdk.Manifest,
+	serviceBanner codersdk.ServiceBannerConfig,
+) *ssh.Session {
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 	//nolint:dogsled
-	conn, _, _, _, _ := setupAgent(t, options, 0)
+	conn, client, _, _, _ := setupAgent(t, options, 0)
+	client.getServiceBanner = func() (codersdk.ServiceBannerConfig, error) {
+		return serviceBanner, nil
+	}
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
 	t.Cleanup(func() {
@@ -1809,6 +1881,7 @@ type client struct {
 	coordinator        tailnet.Coordinator
 	lastWorkspaceAgent func()
 	patchWorkspaceLogs func() error
+	getServiceBanner   func() (codersdk.ServiceBannerConfig, error)
 
 	mu              sync.Mutex // Protects following.
 	lifecycleStates []codersdk.WorkspaceAgentLifecycle
@@ -1930,6 +2003,15 @@ func (c *client) PatchStartupLogs(_ context.Context, logs agentsdk.PatchStartupL
 	return nil
 }
 
+func (c *client) GetServiceBanner(_ context.Context) (codersdk.ServiceBannerConfig, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.getServiceBanner != nil {
+		return c.getServiceBanner()
+	}
+	return codersdk.ServiceBannerConfig{}, nil
+}
+
 // tempDirUnixSocket returns a temporary directory that can safely hold unix
 // sockets (probably).
 //

diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
@@ -29,6 +29,7 @@ import (
 	"cdr.dev/slog"
 
 	"github.com/coder/coder/agent/usershell"
+	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/codersdk/agentsdk"
 	"github.com/coder/coder/pty"
 )
@@ -63,9 +64,10 @@ type Server struct {
 	srv          *ssh.Server
 	x11SocketDir string
 
-	Env        map[string]string
-	AgentToken func() string
-	Manifest   *atomic.Pointer[agentsdk.Manifest]
+	Env              map[string]string
+	AgentToken       func() string
+	Manifest         *atomic.Pointer[agentsdk.Manifest]
+	GetServiceBanner func(ctx context.Context) (codersdk.ServiceBannerConfig, error)
 
 	connCountVSCode     atomic.Int64
 	connCountJetBrains  atomic.Int64
@@ -346,6 +348,11 @@ func (s *Server) startPTYSession(session ptySession, magicTypeLabel string, cmd
 	session.DisablePTYEmulation()
 
 	if !isQuietLogin(session.RawCommand()) {
+		err := s.showServiceBanner(ctx, session)
+		if err != nil {
+			s.logger.Error(ctx, "agent failed to show service banner", slog.Error(err))
+			s.metrics.sessionErrors.WithLabelValues(magicTypeLabel, "yes", "service_banner").Add(1)
+		}
 		manifest := s.Manifest.Load()
 		if manifest != nil {
 			err := showMOTD(session, manifest.MOTDFile)
@@ -428,6 +435,19 @@ func (s *Server) startPTYSession(session ptySession, magicTypeLabel string, cmd
 	return nil
 }
 
+func (s *Server) showServiceBanner(ctx context.Context, session io.Writer) error {
+	banner, err := s.GetServiceBanner(ctx)
+	if err != nil {
+		return err
+	}
+	if banner.Enabled && banner.Message != "" {
+		// The banner supports Markdown so we might want to parse it but Markdown is
+		// still fairly readable in its raw form.
+		_, err = io.WriteString(session, banner.Message+"\n\n\r")
+	}
+	return err
+}
+
 func (s *Server) sftpHandler(session ssh.Session) {
 	s.metrics.sftpConnectionsTotal.Add(1)