1 // Copyright 2023 The Sigstore Authors. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 package ui 15 16 import ( 17 "bytes" 18 "context" 19 "io" 20 "os" 21 ) 22 23 // An Env is the environment that the CLI exists in. 24 // 25 // It contains handles to STDERR and STDIN. Eventually, it will contain 26 // configuration pertaining to the current invocation (e.g., is this a terminal 27 // or not). 28 // 29 // UI methods should be defined on an Env. Then, the Env can be 30 // changed for easy testing. The Env will be retrieved from the current 31 // application context. 32 type Env struct { 33 Stderr io.Writer 34 Stdin io.Reader 35 } 36 37 // defaultEnv returns the default environment (writing to os.Stderr and 38 // reading from os.Stdin). 39 func defaultEnv() *Env { 40 return &Env{ 41 Stderr: os.Stderr, 42 Stdin: os.Stdin, 43 } 44 } 45 46 type ctxKey struct{} 47 48 func (c ctxKey) String() string { 49 return "cosign/ui:env" 50 } 51 52 var ctxKeyEnv = ctxKey{} 53 54 // getEnv gets the environment from ctx. 55 // 56 // If ctx does not contain an environment, getEnv returns the default 57 // environment (see defaultEnvironment). 58 func getEnv(ctx context.Context) *Env { 59 e, ok := ctx.Value(ctxKeyEnv).(*Env) 60 if !ok { 61 return defaultEnv() 62 } 63 return e 64 } 65 66 // WithEnv adds the environment to the context. 67 func WithEnv(ctx context.Context, e *Env) context.Context { 68 return context.WithValue(ctx, ctxKeyEnv, e) 69 } 70 71 type WriteFunc func(string) 72 type callbackFunc func(context.Context, WriteFunc) 73 74 // RunWithTestCtx runs the provided callback in a context with the UI 75 // environment swapped out for one that allows for easy testing and captures 76 // STDOUT. 77 // 78 // The callback has access to a function that writes to the test STDIN. 79 func RunWithTestCtx(callback callbackFunc) string { 80 var stdin bytes.Buffer 81 var stderr bytes.Buffer 82 e := Env{&stderr, &stdin} 83 84 ctx := WithEnv(context.Background(), &e) 85 write := func(msg string) { stdin.WriteString(msg) } 86 callback(ctx, write) 87 88 return stderr.String() 89 } 90