cordination between agent and foreground cmd to keep tokens updated i…

…n both
superfly · Apr 8, 2024 · d0f22ee · d0f22ee
1 parent 7bb4ca8
commit d0f22ee
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 73 deletions.
diff --git a/agent/server/server.go b/agent/server/server.go
@@ -137,17 +137,19 @@ func (s *server) serve(parent context.Context, l net.Listener) (err error) {
 		return nil
 	})
 
-	eg.Go(func() error {
-		if f := config.Tokens(ctx).FromConfigFile; f == "" {
-			s.print("monitoring for token expiration")
-			s.updateMacaroonsInMemory(ctx)
-		} else {
-			s.print("monitoring for token changes and expiration")
-			s.updateMacaroonsInFile(ctx, f)
-		}
+	if toks := config.Tokens(ctx); len(toks.MacaroonTokens) != 0 {
+		eg.Go(func() error {
+			if f := toks.FromConfigFile; f == "" {
+				s.print("monitoring for token expiration")
+				s.updateMacaroonsInMemory(ctx)
+			} else {
+				s.print("monitoring for token changes and expiration")
+				s.updateMacaroonsInFile(ctx, f)
+			}
 
-		return nil
-	})
+			return nil
+		})
+	}
 
 	eg.Go(func() (err error) {
 		s.printf("OK %d", os.Getpid())
@@ -399,75 +401,49 @@ func (s *server) updateMacaroonsInMemory(ctx context.Context) {
 	}
 }
 
-// updateMacaroons updates the agent's tokens as the config file changes. It
-// also prunes expired tokens and fetches discharge tokens as necessary. Those
-// updates are written back to the config file.
+// updateMacaroons prunes expired tokens and fetches discharge tokens as
+// necessary. Those updates are written back to the config file.
 func (s *server) updateMacaroonsInFile(ctx context.Context, path string) {
-	toks := config.Tokens(ctx)
+	configToks := config.Tokens(ctx)
 
 	ticker := time.NewTicker(time.Minute)
 	defer ticker.Stop()
 
 	var lastErr error
 
-	fiBefore, err := os.Stat(path)
-	if err != nil {
-		s.print("failed stating config file:", err)
-		s.updateMacaroonsInMemory(ctx)
-		return
-	}
-
 	for {
-		updated, err := toks.Update(ctx, tokens.WithDebugger(s))
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+		}
+
+		// the tokens in the config are continually updated as the config file
+		// changes. We do our updates on a copy of the tokens so we can still
+		// tell if the tokens in the config changed out from under us.
+		configToksBefore := configToks.All()
+		localToks := tokens.Parse(configToksBefore)
+
+		updated, err := localToks.Update(ctx, tokens.WithDebugger(s))
 		if err != nil && err != lastErr {
 			s.print("failed upgrading authentication tokens:", err)
 			lastErr = err
 
 			// Don't continue loop here! It might only be partial failure
 		}
 
-		if updated {
-			fiAfter, err := os.Stat(path)
-			if err != nil {
-				s.print("failed stating config file:", err)
+		// the consequences of a race here (agent and foreground command both
+		// fetching updates simultaneously) are low, so don't bother with a lock
+		// file.
+		if updated && configToks.All() == configToksBefore {
+			if err := config.SetAccessToken(path, localToks.All()); err != nil {
+				s.print("Failed to persist authentication token:", err)
 				s.updateMacaroonsInMemory(ctx)
 				return
 			}
 
-			// Don't write updates if the file changed out from under us. This
-			// isn't as strong of an assurance as a lockfile would be, but a
-			// race isn't that consequential.
-			if fiBefore.ModTime() == fiAfter.ModTime() {
-				if err := config.SetAccessToken(path, toks.All()); err != nil {
-					s.print("Failed to persist authentication token:", err)
-					s.updateMacaroonsInMemory(ctx)
-					return
-				}
-
-				s.print("Authentication tokens upgraded")
-			}
-		}
-
-		select {
-		case <-ctx.Done():
-			return
-		case <-ticker.C:
+			s.print("Authentication tokens upgraded")
 		}
-
-		if fiBefore, err = os.Stat(path); err != nil {
-			s.print("failed stating config file:", err)
-			s.updateMacaroonsInMemory(ctx)
-			return
-		}
-
-		tok, err := config.ReadAccessToken(path)
-		if err != nil {
-			s.print("failed reading config file:", err)
-			s.updateMacaroonsInMemory(ctx)
-			return
-		}
-
-		toks.Replace(tokens.Parse(tok))
 	}
 }
 

diff --git a/internal/command/agent/run.go b/internal/command/agent/run.go
@@ -27,7 +27,9 @@ func newRun() (cmd *cobra.Command) {
 		long  = short + "\n"
 	)
 
-	cmd = command.New("run", short, long, run)
+	cmd = command.New("run", short, long, run,
+		command.RequireSession,
+	)
 
 	cmd.Args = cobra.MaximumNArgs(1)
 	cmd.Aliases = []string{"daemon-start"}
@@ -46,12 +48,6 @@ func run(ctx context.Context) error {
 	}
 	defer closeLogger()
 
-	apiClient := fly.ClientFromContext(ctx)
-	if !apiClient.Authenticated() {
-		logger.Println(fly.ErrNoAuthToken)
-		return fly.ErrNoAuthToken
-	}
-
 	unlock, err := lock(ctx, logger)
 	if err != nil {
 		return err
@@ -61,7 +57,7 @@ func run(ctx context.Context) error {
 	opt := server.Options{
 		Socket:           socketPath(ctx),
 		Logger:           logger,
-		Client:           apiClient,
+		Client:           fly.ClientFromContext(ctx),
 		Background:       logPath != "",
 		ConfigFile:       state.ConfigFile(ctx),
 		ConfigWebsockets: viper.GetBool(flyctl.ConfigWireGuardWebsockets),

diff --git a/internal/command/command.go b/internal/command/command.go
@@ -557,9 +557,11 @@ func RequireSession(ctx context.Context) (context.Context, error) {
 // updateMacaroons prune any invalid/expired macaroons and fetch needed third
 // party discharges
 func updateMacaroons(ctx context.Context) (context.Context, error) {
-	log := logger.FromContext(ctx)
-
-	toks := config.Tokens(ctx)
+	var (
+		log  = logger.FromContext(ctx)
+		cfg  = config.FromContext(ctx)
+		toks = cfg.Tokens
+	)
 
 	updated, err := toks.Update(ctx,
 		tokens.WithUserURLCallback(tryOpenUserURL),
@@ -570,15 +572,33 @@ func updateMacaroons(ctx context.Context) (context.Context, error) {
 		log.Debug(err)
 	}
 
-	if !updated || toks.FromConfigFile == "" {
+	if toks.FromConfigFile == "" {
 		return ctx, nil
 	}
 
-	if err := config.SetAccessToken(toks.FromConfigFile, toks.All()); err != nil {
-		log.Warn("Failed to persist authentication token.")
+	if updated {
+		if err := config.SetAccessToken(toks.FromConfigFile, toks.All()); err != nil {
+			log.Warn("Failed to persist authentication token.")
+			log.Debug(err)
+		}
+	}
+
+	sub, err := cfg.Watch(ctx)
+	if err != nil {
+		log.Warn("Failed to watch config file for changes.")
 		log.Debug(err)
+		return ctx, nil
 	}
 
+	go func() {
+		for newCfg := range sub {
+			if cfg.Tokens.All() != newCfg.Tokens.All() {
+				log.Debug("Authentication tokens updated from config file.")
+				cfg.Tokens.Replace(newCfg.Tokens)
+			}
+		}
+	}()
+
 	return ctx, nil
 }