diff --git a/pkg/server/admin/admin.go b/pkg/server/admin/admin.go index db0c78a..3e552c6 100644 --- a/pkg/server/admin/admin.go +++ b/pkg/server/admin/admin.go @@ -96,7 +96,7 @@ func (admin *Admin) getFreeId(publicId models.RendezVousId) (models.RendezVousId if admin.agents[publicId] == nil { return publicId, nil } - for i := 0; i < 100; i++ { + for i := range 100 { candidate := models.RendezVousId(string(publicId) + "-" + strconv.Itoa(i)) if admin.agents[candidate] == nil { return candidate, nil @@ -130,6 +130,8 @@ func (admin *Admin) AddAgent(hostKey []byte, publicId models.RendezVousId, agent Ok: false, Message: err.Error(), }) + return nil, fmt.Errorf( + "Agent requested id '%s' which is already taken anc could not allocate a new unique id", publicId) } agentCheck := admin.agents[publicId] if agentCheck != nil { diff --git a/pkg/server/admin/admin_test.go b/pkg/server/admin/admin_test.go index 2035484..e2f3c38 100644 --- a/pkg/server/admin/admin_test.go +++ b/pkg/server/admin/admin_test.go @@ -10,6 +10,7 @@ import ( "github.com/stretchr/testify/suite" "io" "net/http" + "strings" "testing" "time" ) @@ -62,7 +63,7 @@ func (s *AdminTestSuite) SetupTest() { s.admin = NewAdmin() s.hostKey = make([]byte, 100) - rand.Read(s.hostKey) + s.NotNil(rand.Read(s.hostKey)) } func (suite *AdminTestSuite) TearDownTest() { @@ -72,21 +73,43 @@ func TestAdminTestSuite(t *testing.T) { suite.Run(t, &AdminTestSuite{}) } -func (s *AdminTestSuite) agentRegisters(requestedPublicId, assignedPublicId string) *agentConnection { +type AddAgentResult struct { + agentConn *agentConnection + err error +} + +type AgentRegisterResult struct { + registration comms.AgentRegistration + err error +} + +func (s *AdminTestSuite) agentRegisters(requestedPublicId, assignedPublicId string) (AddAgentResult, AgentRegisterResult) { agentRW, serverRW := s.createPipe() res := testsupport.RunAndWait( &s.Suite, func() any { - return s.addAgent(requestedPublicId, assignedPublicId, serverRW) + agentConn, err := s.addAgent(requestedPublicId, assignedPublicId, serverRW) + return AddAgentResult{ + agentConn: agentConn, + err: err, + } }, func() any { - return s.agentRegistration(assignedPublicId, agentRW) + res := s.agentRegistration(assignedPublicId, agentRW) + if assignedPublicId != "" { + s.Nil(res.err) + s.True(res.registration.Ok) + s.Equal(s.hostKey, res.registration.HostPrivateKey) + } + return res }) - return res[0].(*agentConnection) + return res[0].(AddAgentResult), res[1].(AgentRegisterResult) } func (s *AdminTestSuite) Test_AgentRegisters() { - agentConn := s.agentRegisters("abc", "abc") + res, _ := s.agentRegisters("abc", "abc") + s.Nil(res.err) + agentConn := res.agentConn state := s.admin.CreateNotifification() s.Equal(1, len(state.Agents)) s.Equal(0, len(state.Clients)) @@ -99,7 +122,9 @@ func (s *AdminTestSuite) Test_ManyAgentsRegister() { for i := range N { publicId := fmt.Sprintf("abc%d", i) agentRegistrations[i] = func() any { - return s.agentRegisters(publicId, publicId) + res, _ := s.agentRegisters(publicId, publicId) + s.Nil(res.err) + return res.agentConn } } res := testsupport.RunAndWait( @@ -115,30 +140,40 @@ func (s *AdminTestSuite) Test_ManyAgentsRegister() { } func (s *AdminTestSuite) Test_agentDuplicateId() { - s.agentRegisters("abc", "abc") - s.agentRegisters("abc", "abc-0") + res, _ := s.agentRegisters("abc", "abc") + s.Nil(res.err) + for i := range 100 { + res, _ = s.agentRegisters("abc", fmt.Sprintf("abc-%d", i)) + s.Nil(res.err) + } + res, agentSideResult := s.agentRegisters("abc", "") + s.NotNil(res.err) + // verify it is the correct error an dnot an id mismatch. + s.True(strings.Contains(res.err.Error(), "could not allocate a new unique id")) + s.False(agentSideResult.registration.Ok) } -func (s *AdminTestSuite) agentRegistration(expectedPublicId string, agentRW io.ReadWriteCloser) any { +func (s *AdminTestSuite) agentRegistration(expectedPublicId string, agentRW io.ReadWriteCloser) AgentRegisterResult { // verify registration message received agentRegistration, err := comms.ReceiveRegistrationMessage(agentRW) - s.Nil(err) - s.True(agentRegistration.Ok) - s.Equal(expectedPublicId, agentRegistration.Id) - s.Equal(s.hostKey, agentRegistration.HostPrivateKey) - + if err != nil { + return AgentRegisterResult{err: err} + } commChannel, err := comms.NewCommChannel(comms.Agent, agentRW) - s.Nil(err) + if err != nil { + return AgentRegisterResult{registration: agentRegistration, err: err} + } s.NotNil(commChannel) - return nil + return AgentRegisterResult{registration: agentRegistration, err: nil} } -func (s *AdminTestSuite) addAgent(publicId string, assignedPublicId string, serverRW io.ReadWriteCloser) any { +func (s *AdminTestSuite) addAgent(publicId string, assignedPublicId string, serverRW io.ReadWriteCloser) (*agentConnection, error) { agentConn, err := s.admin.AddAgent( s.hostKey, models.RendezVousId(publicId), comms.EnvironmentInfo{}, serverRW) - s.Nil(err) + if err != nil { + return nil, err + } s.Equal(assignedPublicId, string(agentConn.Info.PublicId)) - - return agentConn + return agentConn, nil }