Address PR feedback

mafredri · mafredri · commit 94d759396ad8 · 2023-04-06T09:50:51.000Z
diff --git a/agent/agent.go b/agent/agent.go
@@ -167,6 +167,7 @@ func (a *agent) init(ctx context.Context) {
 	}
 	sshSrv.Env = a.envVars
 	sshSrv.AgentToken = func() string { return *a.sessionToken.Load() }
+	sshSrv.Manifest = &a.manifest
 	a.sshServer = sshSrv
 
 	go a.runLoop(ctx)
@@ -478,7 +479,6 @@ func (a *agent) run(ctx context.Context) error {
 	}
 
 	oldManifest := a.manifest.Swap(&manifest)
-	a.sshServer.SetManifest(&manifest)
 
 	// The startup script should only execute on the first run!
 	if oldManifest == nil {
diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
@@ -16,11 +16,11 @@ import (
 	"runtime"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/gliderlabs/ssh"
 	"github.com/pkg/sftp"
+	"go.uber.org/atomic"
 	gossh "golang.org/x/crypto/ssh"
 	"golang.org/x/xerrors"
 
@@ -54,8 +54,7 @@ type Server struct {
 
 	Env        map[string]string
 	AgentToken func() string
-
-	manifest atomic.Pointer[agentsdk.Manifest]
+	Manifest   *atomic.Pointer[agentsdk.Manifest]
 
 	connCountVSCode     atomic.Int64
 	connCountJetBrains  atomic.Int64
@@ -130,11 +129,6 @@ func NewServer(ctx context.Context, logger slog.Logger, maxTimeout time.Duration
 	return s, nil
 }
 
-// SetManifest sets the manifest used for starting commands.
-func (s *Server) SetManifest(m *agentsdk.Manifest) {
-	s.manifest.Store(m)
-}
-
 type ConnStats struct {
 	Sessions  int64
 	VSCode    int64
@@ -215,7 +209,7 @@ func (s *Server) sessionStart(session ssh.Session) (retErr error) {
 		session.DisablePTYEmulation()
 
 		if !isQuietLogin(session.RawCommand()) {
-			manifest := s.manifest.Load()
+			manifest := s.Manifest.Load()
 			if manifest != nil {
 				err = showMOTD(session, manifest.MOTDFile)
 				if err != nil {
@@ -388,7 +382,7 @@ func (s *Server) CreateCommand(ctx context.Context, script string, env []string)
 		return nil, xerrors.Errorf("get user shell: %w", err)
 	}
 
-	manifest := s.manifest.Load()
+	manifest := s.Manifest.Load()
 	if manifest == nil {
 		return nil, xerrors.Errorf("no metadata was provided")
 	}
diff --git a/agent/agentssh/agentssh_test.go b/agent/agentssh/agentssh_test.go
@@ -1,3 +1,5 @@
+// Package agentssh_test provides tests for basic functinoality of the agentssh
+// package, more test coverage can be found in the `agent` and `cli` package(s).
 package agentssh_test
 
 import (
@@ -10,6 +12,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
 	"go.uber.org/goleak"
 	"golang.org/x/crypto/ssh"
 
@@ -34,7 +37,7 @@ func TestNewServer_ServeClient(t *testing.T) {
 
 	// The assumption is that these are set before serving SSH connections.
 	s.AgentToken = func() string { return "" }
-	s.SetManifest(&agentsdk.Manifest{})
+	s.Manifest = atomic.NewPointer(&agentsdk.Manifest{})
 
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
@@ -74,7 +77,7 @@ func TestNewServer_CloseActiveConnections(t *testing.T) {
 
 	// The assumption is that these are set before serving SSH connections.
 	s.AgentToken = func() string { return "" }
-	s.SetManifest(&agentsdk.Manifest{})
+	s.Manifest = atomic.NewPointer(&agentsdk.Manifest{})
 
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)

Original file line number	Diff line number	Diff line change
`@@ -167,6 +167,7 @@ func (a *agent) init(ctx context.Context) {`
`167`	`167`	`}`
`168`	`168`	`sshSrv.Env = a.envVars`
`169`	`169`	`sshSrv.AgentToken = func() string { return *a.sessionToken.Load() }`
	`170`	`+ sshSrv.Manifest = &a.manifest`
`170`	`171`	`a.sshServer = sshSrv`
`171`	`172`
`172`	`173`	`go a.runLoop(ctx)`
`@@ -478,7 +479,6 @@ func (a *agent) run(ctx context.Context) error {`
`478`	`479`	`}`
`479`	`480`
`480`	`481`	`oldManifest := a.manifest.Swap(&manifest)`
`481`		`- a.sshServer.SetManifest(&manifest)`
`482`	`482`
`483`	`483`	`// The startup script should only execute on the first run!`
`484`	`484`	`if oldManifest == nil {`