1package auth
2
3import (
4 "context"
5 "crypto/sha256"
6 "encoding/json"
7 "errors"
8 "fmt"
9 "io"
10 "os/exec"
11 "strings"
12 "sync"
13 "time"
14)
15
16const (
17 defaultAPIKeyCommandTimeout = 120 * time.Second
18 maxAPIKeyCommandOutput = 64 << 10
19)
20
21// APIKeyCommand describes a program that prints an API key to stdout.
22// Program is executed directly, without a shell.
23type APIKeyCommand struct {
24 Program string `json:"program"`
25 Args []string `json:"args,omitempty"`
26 TimeoutMS int `json:"timeout_ms,omitempty"`
27}
28
29type apiKeyCommandCall struct {
30 done chan struct{}
31 value string
32 err error
33}
34
35var apiKeyCommandCache = struct {
36 sync.Mutex
37 values map[[sha256.Size]byte]string
38 inflight map[[sha256.Size]byte]*apiKeyCommandCall
39}{
40 values: make(map[[sha256.Size]byte]string),
41 inflight: make(map[[sha256.Size]byte]*apiKeyCommandCall),
42}
43
44// ResolveAPIKeyCommand executes command and returns its API key. Successful
45// results are cached in memory for the lifetime of the process. Concurrent
46// requests for the same command share one execution.
47func ResolveAPIKeyCommand(ctx context.Context, command APIKeyCommand) (string, error) {
48 if strings.TrimSpace(command.Program) == "" {
49 return "", errors.New("api key command has an empty program")
50 }
51 if command.TimeoutMS < 0 {
52 return "", errors.New("api key command timeout_ms must not be negative")
53 }
54
55 encoded, err := json.Marshal(command)
56 if err != nil {
57 return "", errors.New("encode api key command")
58 }
59 cacheKey := sha256.Sum256(encoded)
60 apiKeyCommandCache.Lock()
61 if value, ok := apiKeyCommandCache.values[cacheKey]; ok {
62 apiKeyCommandCache.Unlock()
63 return value, nil
64 }
65 if call, ok := apiKeyCommandCache.inflight[cacheKey]; ok {
66 apiKeyCommandCache.Unlock()
67 select {
68 case <-call.done:
69 return call.value, call.err
70 case <-ctx.Done():
71 return "", errors.New("api key command was canceled")
72 }
73 }
74 call := &apiKeyCommandCall{done: make(chan struct{})}
75 apiKeyCommandCache.inflight[cacheKey] = call
76 apiKeyCommandCache.Unlock()
77
78 value, resolveErr := runAPIKeyCommand(ctx, command)
79 apiKeyCommandCache.Lock()
80 delete(apiKeyCommandCache.inflight, cacheKey)
81 if resolveErr == nil {
82 apiKeyCommandCache.values[cacheKey] = value
83 }
84 call.value = value
85 call.err = resolveErr
86 close(call.done)
87 apiKeyCommandCache.Unlock()
88 return value, resolveErr
89}
90
91func runAPIKeyCommand(ctx context.Context, command APIKeyCommand) (string, error) {
92 timeout := defaultAPIKeyCommandTimeout
93 if command.TimeoutMS > 0 {
94 timeout = time.Duration(command.TimeoutMS) * time.Millisecond
95 }
96 commandCtx, cancel := context.WithTimeout(ctx, timeout)
97 defer cancel()
98
99 cmd := exec.CommandContext(commandCtx, command.Program, command.Args...)
100 configureAPIKeyCommandProcess(cmd)
101 stdout, err := cmd.StdoutPipe()
102 if err != nil {
103 return "", errors.New("prepare api key command output")
104 }
105 if err := cmd.Start(); err != nil {
106 return "", errors.New("start api key command")
107 }
108
109 output, readErr := io.ReadAll(io.LimitReader(stdout, maxAPIKeyCommandOutput+1))
110 if len(output) > maxAPIKeyCommandOutput {
111 cancelAPIKeyCommandProcess(cmd)
112 _ = cmd.Wait()
113 return "", fmt.Errorf("api key command output exceeds %d bytes", maxAPIKeyCommandOutput)
114 }
115 waitErr := cmd.Wait()
116 if commandCtx.Err() != nil {
117 return "", errors.New("api key command timed out or was canceled")
118 }
119 if readErr != nil {
120 return "", errors.New("read api key command output")
121 }
122 if waitErr != nil {
123 return "", errors.New("api key command exited unsuccessfully")
124 }
125
126 value := strings.TrimRight(string(output), "\r\n")
127 if value == "" {
128 return "", errors.New("api key command returned empty output")
129 }
130 if strings.ContainsAny(value, "\r\n") {
131 return "", errors.New("api key command returned multiple lines")
132 }
133
134 return value, nil
135}