Landscape mock: fix issue were distros would receive a new UID every …

…time

mocks/landscape/landscapemockservice/landscape_mock_service.go

@@ -82,38 +82,56 @@ func New() *Service {
 // Connect implements the Connect API call.
 // Upon first contact ever, a UID is randombly assigned to the host and sent to it.
 // In subsequent contacts, this UID will be its unique identifier.
-func (s *Service) Connect(stream landscapeapi.LandscapeHostAgent_ConnectServer) error {
+func (s *Service) Connect(stream landscapeapi.LandscapeHostAgent_ConnectServer) (err error) {
 	ctx, cancel := context.WithCancel(stream.Context())
 	defer cancel()
 
 	firstContact := true
-	ch := make(chan HostInfo)
+
+	type msg struct {
+		info HostInfo
+		err  error
+	}
+
+	ch := make(chan msg)
 	defer close(ch)
 
+	var hostInfo HostInfo
 	for {
 		go func() {
 			recv, err := stream.Recv()
-			if err != nil {
-				return
-			}
 
 			select {
 			case <-ctx.Done():
-				return
-			default:
+			case ch <- msg{newHostInfo(recv), err}:
 			}
 
-			ch <- newHostInfo(recv)
+			return
 		}()
 
-		var hostInfo HostInfo
+		var m msg
+		select {
+		case m = <-ch:
+		case <-ctx.Done():
+			slog.Info(fmt.Sprintf("Landscape: %s: terminated connection: %v", hostInfo.Hostname, ctx.Err()))
+			return nil
+		}
+
+		// Avoid races when ctx.Done and ch were both ready
 		select {
-		case hostInfo = <-ch:
 		case <-ctx.Done():
 			slog.Info(fmt.Sprintf("Landscape: %s: terminated connection: %v", hostInfo.Hostname, ctx.Err()))
 			return nil
+		default:
 		}
 
+		if m.err != nil {
+			slog.Info(fmt.Sprintf("Landscape: %s: terminated connection: %v", hostInfo.Hostname, err))
+			return err
+		}
+
+		hostInfo = m.info
+
 		if hostInfo.Token == "" {
 			// No token: no Landscape
 			err := errors.New("connection rejected: no Ubuntu Pro token provided")