Refactor

.github/workflows/test-and-lint.yaml

@@ -4,6 +4,7 @@ on:
   push:
     branches:
       - main
+
     paths:
       - "**/*.go"
       - "go.mod"

auth/cache.go

@@ -6,7 +6,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"log"
 	"os"
 	"path/filepath"
 	"regexp"
@@ -25,7 +24,7 @@ type (
 	}
 
 	Cacher struct {
-		logger   *log.Logger
+		logger   Logger
 		cacheDir string
 		cipher   Cipher
 	}
@@ -85,7 +84,7 @@ func DecodeFromFileName(roleArn, fileName string) (ts time.Time, err error) {
 }
 
 // Non-nil returned error wraps [ErrCacheInit].
-func NewCacher(logger *log.Logger, cipher Cipher) (*Cacher, error) {
+func NewCacher(logger Logger, cipher Cipher) (*Cacher, error) {
 	home, err := os.UserHomeDir()
 	if err != nil {
 		return nil, fmt.Errorf("%w: could not locate user home directory", ErrCacheInit)
@@ -137,6 +136,8 @@ func (c *Cacher) Save(roleArn string, output *CredentialProcessOutput) (contents
 	return contents, nil
 }
 
+// Retrieve tries to retrieve AWS credentials from cache files.
+// It succeeded if and only if the returned byte slice is not nil.
 func (c *Cacher) Retrieve(roleArn string) (contents []byte) {
 	max := time.Now().Add(time.Minute * 10)
 	actives := make(cacheFileSlice, 0)

auth/command.go

@@ -7,7 +7,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"log"
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -16,8 +15,14 @@ import (
 )
 
 type (
+	Logger interface {
+		Printf(string, ...any)
+		Print(...any)
+	}
+
 	AssumeRoleCmd struct {
-		prompter        *Prompter
+		logger          Logger
+		prompter        Prompter
 		cacher          *Cacher
 		client          *sts.Client
 		RoleArn         string
@@ -29,7 +34,6 @@ type (
 	}
 
 	Prompter struct {
-		*log.Logger
 		io.ReadWriter
 	}
 )
@@ -38,11 +42,7 @@ var (
 	ErrInvalidInput = errors.New("invalid CLI input")
 )
 
-func NewPrompter(tty io.ReadWriter, prefix string, flag int) *Prompter {
-	return &Prompter{ReadWriter: tty, Logger: log.New(tty, prefix, flag)}
-}
-
-func (c *Prompter) MFAToken() (code string, err error) {
+func (c Prompter) MFAToken() (code string, err error) {
 	_, err = io.WriteString(c, "MFA code: ")
 	if err != nil {
 		return "", fmt.Errorf("failed to prompt for MFA code: %w", err)
@@ -82,8 +82,10 @@ func (a *AssumeRoleCmd) ValidateInputs(args []string) error {
 }
 
 // Non-nil returned error wraps [ErrCacheInit].
-func (a *AssumeRoleCmd) InitCache(tty io.ReadWriter, cfg aws.Config) error {
-	a.prompter = NewPrompter(tty, "toolkit-assume-role: ", 0)
+func (a *AssumeRoleCmd) Init(tty *TTY, cfg aws.Config) error {
+	a.prompter = Prompter{ReadWriter: tty}
+
+	a.logger = tty
 
 	a.client = sts.NewFromConfig(cfg)
 
@@ -92,11 +94,12 @@ func (a *AssumeRoleCmd) InitCache(tty io.ReadWriter, cfg aws.Config) error {
 		return fmt.Errorf("%w: failed to create cache cipher: %s", ErrCacheInit, err.Error())
 	}
 
-	a.cacher, err = NewCacher(a.prompter.Logger, cipher)
+	a.cacher, err = NewCacher(a.logger, cipher)
 
 	return err
 }
 
+// Non-nil returned error means failure.
 func (a *AssumeRoleCmd) Run(ctx context.Context, dest io.Writer) (err error) {
 	// Output of the AWS CLI credential process.
 	var output []byte
@@ -139,7 +142,7 @@ func (a *AssumeRoleCmd) Run(ctx context.Context, dest io.Writer) (err error) {
 		if errors.Is(err, ErrInvalidCredential) {
 			return err
 		} else if err != nil {
-			a.prompter.Print(err.Error())
+			a.logger.Print(err.Error())
 		}
 
 		_, err = dest.Write(output)

auth/command_test.go

@@ -75,9 +75,6 @@ func TestAssumeRoleCmdRun(t *testing.T) {
 
 	defer func() { fromKeyring = keyringGet }()
 
-	tty := MockFileDescriptor{}
-	dest := MockFileDescriptor{}
-
 	hdm := HomeDirMocker{}
 
 	hdm.SetUp(t)
@@ -101,6 +98,15 @@ func TestAssumeRoleCmdRun(t *testing.T) {
 
 	duration := int32(cmd.DurationSeconds)
 
+	mockedTtyDevice := &MockFileDescriptor{}
+
+	_, err := mockedTtyDevice.r.WriteString(token + "\n")
+	require.NoError(t, err, "should be able to write token to mocked TTY file descriptor")
+
+	tty := NewTTY(mockedTtyDevice, "toolkit-assume-role: ", 0)
+
+	dest := MockFileDescriptor{}
+
 	soutput := CredentialProcessOutput{
 		AccessKeyId:     "access-key-id",
 		SecretAccessKey: "secret-access-key",
@@ -131,9 +137,6 @@ func TestAssumeRoleCmdRun(t *testing.T) {
 		Error: nil,
 	})
 
-	_, err := tty.r.WriteString(token + "\n")
-	require.NoError(t, err, "should be able to write token to mocked TTY file descriptor")
-
 	ctx := context.Background()
 
 	err = cmd.ValidateInputs([]string{roleArn})
@@ -145,7 +148,7 @@ func TestAssumeRoleCmdRun(t *testing.T) {
 	cfg, err := config.LoadDefaultConfig(ctx, config.WithCredentialsProvider(mockCredsProvider), config.WithRegion(cmd.Region))
 	require.NoError(t, err, "should be able to load config using a mocked credentials provider")
 
-	err = cmd.InitCache(&tty, cfg)
+	err = cmd.Init(tty, cfg)
 	require.NoError(t, err, "should be able to init caching without error")
 
 	// Stub the STS client object.

auth/doc.go

@@ -0,0 +1,4 @@
+// Package auth implements AWS credential process with caching.
+// Cache files are saved on disk and encrypted via AES-GCM with the encryption key stored in the operating system's "native" credentials store.
+// For example, Keychain is used on macOS.
+package auth

auth/tty.go

@@ -0,0 +1,62 @@
+package auth
+
+import (
+	"bytes"
+	"io"
+	"log"
+	"sync"
+)
+
+type (
+	TTY struct {
+		dest   io.ReadWriter
+		logger *log.Logger
+		buf    bytes.Buffer
+		mux    sync.Mutex // Guards the whole struct
+	}
+)
+
+func NewTTY(dest io.ReadWriter, prefix string, flag int) *TTY {
+	tty := TTY{dest: dest}
+
+	tty.logger = log.New(&tty.buf, prefix, flag)
+
+	return &tty
+}
+
+func (t *TTY) Read(p []byte) (n int, err error) {
+	t.mux.Lock()
+	defer t.mux.Unlock()
+
+	return t.dest.Read(p)
+}
+
+func (t *TTY) Write(p []byte) (n int, err error) {
+	t.mux.Lock()
+	defer t.mux.Unlock()
+
+	return t.dest.Write(p)
+}
+
+func (t *TTY) Printf(format string, v ...any) {
+	t.mux.Lock()
+	defer t.mux.Unlock()
+
+	t.logger.Printf(format, v...)
+}
+
+func (t *TTY) Print(v ...any) {
+	t.mux.Lock()
+	defer t.mux.Unlock()
+
+	t.logger.Print(v...)
+}
+
+func (t *TTY) FlushLogs() error {
+	t.mux.Lock()
+	defer t.mux.Unlock()
+
+	_, err := t.dest.Write(t.buf.Bytes())
+
+	return err
+}

cmd/toolkit-assume-role/main.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"flag"
 	"fmt"
-	"log"
 	"os"
 
 	"github.com/aws/aws-sdk-go-v2/config"
@@ -37,53 +36,65 @@ func registerFlagsAndHelp() {
 		_, _ = fmt.Fprintf(flag.CommandLine.Output(), helpMsg, os.Args[0])
 
 		flag.PrintDefaults()
-
-		tty, err := os.OpenFile("/dev/tty", os.O_RDWR|os.O_SYNC, 0600)
-		if err != nil {
-			_, _ = fmt.Fprintf(flag.CommandLine.Output(), "\nNot ready to run as credential process: cannot open /dev/tty: %s\n", err)
-
-			os.Exit(1)
-		}
-
-		defer func() { _ = tty.Close() }()
-
-		_, _ = fmt.Fprint(flag.CommandLine.Output(), "\nReady to run as credential process!\n")
 	}
 }
 
 func main() {
+	exitCode := 0
+
+	defer func() { os.Exit(exitCode) }()
+
 	registerFlagsAndHelp()
 
 	flag.Parse()
 
-	tty, err := os.OpenFile("/dev/tty", os.O_RDWR|os.O_SYNC, 0600)
+	ttyDevice, err := os.OpenFile("/dev/tty", os.O_RDWR|os.O_SYNC, 0600)
 	if err != nil {
-		os.Exit(1)
+		exitCode = 1
+
+		return
 	}
 
-	defer func() { _ = tty.Close() }()
+	defer func() { _ = ttyDevice.Close() }()
 
-	ctx := context.Background()
+	tty := auth.NewTTY(ttyDevice, "toolkit-assume-role: ", 0)
+	defer func() {
+		if tty.FlushLogs() != nil {
+			exitCode = 1
+		}
+	}()
 
-	logger := log.New(tty, "toolkit-assume-role: ", 0)
+	ctx := context.Background()
 
 	err = cmd.ValidateInputs(flag.Args())
 	if err != nil {
-		logger.Fatal(err.Error())
+		tty.Print(err.Error())
+
+		exitCode = 1
+
+		return
 	}
 
 	cfg, err := config.LoadDefaultConfig(ctx, config.WithSharedConfigProfile(cmd.Profile), config.WithRegion(cmd.Region))
 	if err != nil {
-		logger.Fatalf("failed to load AWS SDK configuration: %s\n", err)
+		tty.Printf("failed to load AWS SDK configuration: %s\n", err)
+
+		exitCode = 1
+
+		return
 	}
 
-	err = cmd.InitCache(tty, cfg)
+	err = cmd.Init(tty, cfg)
 	if err != nil {
-		logger.Print(err.Error())
+		tty.Print(err.Error())
 	}
 
 	err = cmd.Run(ctx, os.Stdout)
 	if err != nil {
-		logger.Fatal(err.Error())
+		tty.Print(err.Error())
+
+		exitCode = 1
+
+		return
 	}
 }