Skip to content

Commit

Permalink
cordination between agent and foreground cmd to keep tokens updated i…
Browse files Browse the repository at this point in the history
…n both
  • Loading branch information
btoews committed Apr 8, 2024
1 parent 7bb4ca8 commit d0f22ee
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 73 deletions.
94 changes: 35 additions & 59 deletions agent/server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down Expand Up @@ -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))
}
}

Expand Down
12 changes: 4 additions & 8 deletions internal/command/agent/run.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"}
Expand All @@ -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
Expand All @@ -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),
Expand Down
32 changes: 26 additions & 6 deletions internal/command/command.go
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -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
}

Expand Down

0 comments on commit d0f22ee

Please sign in to comment.