Use configurable io.Writer for user facing log messages instead of hardcoding stdout

Signed-off-by: Chance Zibolski <chance.zibolski@gmail.com>

Author: chancez

src/datachannel/streaming.go

@@ -21,8 +21,8 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"math"
-	"os"
 	"reflect"
 	"sync"
 	"time"
@@ -115,6 +115,9 @@ type DataChannel struct {
 
 	// AgentVersion received during handshake
 	agentVersion string
+
+	// Out is where user ssm plugin logs go
+	Out io.Writer
 }
 
 type ListMessageBuffer struct {
@@ -510,7 +513,7 @@ func (dataChannel *DataChannel) handleHandshakeComplete(log log.T, clientMessage
 		handshakeComplete.HandshakeTimeToComplete.Seconds())
 
 	if handshakeComplete.CustomerMessage != "" {
-		fmt.Fprintln(os.Stdout, handshakeComplete.CustomerMessage)
+		fmt.Fprintln(dataChannel.Out, handshakeComplete.CustomerMessage)
 	}
 
 	return err
@@ -783,9 +786,9 @@ func (dataChannel DataChannel) HandleChannelClosedMessage(log log.T, stopHandler
 
 	log.Infof("Exiting session with sessionId: %s with output: %s", sessionId, channelClosedMessage.Output)
 	if channelClosedMessage.Output == "" {
-		fmt.Fprintf(os.Stdout, "\n\nExiting session with sessionId: %s.\n\n", sessionId)
+		fmt.Fprintf(dataChannel.Out, "\n\nExiting session with sessionId: %s.\n\n", sessionId)
 	} else {
-		fmt.Fprintf(os.Stdout, "\n\nSessionId: %s : %s\n\n", sessionId, channelClosedMessage.Output)
+		fmt.Fprintf(dataChannel.Out, "\n\nSessionId: %s : %s\n\n", sessionId, channelClosedMessage.Output)
 	}
 
 	stopHandler()

src/sessionmanagerplugin/session/portsession/basicportforwarding.go

@@ -16,6 +16,7 @@ package portsession
 
 import (
 	"fmt"
+	"io"
 	"net"
 	"os"
 	"os/signal"
@@ -39,6 +40,7 @@ type BasicPortForwarding struct {
 	sessionId      string
 	portParameters PortParameters
 	session        session.Session
+	out            io.Writer
 }
 
 // getNewListener returns a new listener to given address and type like tcp, unix etc.
@@ -132,7 +134,7 @@ func (p *BasicPortForwarding) startLocalConn(log log.T) (err error) {
 		return err
 	}
 	log.Infof("Connection accepted for session %s.", p.sessionId)
-	fmt.Printf("Connection accepted for session %s.\n", p.sessionId)
+	fmt.Fprintf(p.out, "Connection accepted for session %s.\n", p.sessionId)
 
 	p.listener = &listener
 	p.stream = &tcpConn
@@ -159,7 +161,7 @@ func (p *BasicPortForwarding) startLocalListener(log log.T, portNumber string) (
 	}
 
 	log.Info(displayMessage)
-	fmt.Println(displayMessage)
+	fmt.Fprintln(p.out, displayMessage)
 	return
 }
 
@@ -169,13 +171,13 @@ func (p *BasicPortForwarding) handleControlSignals(log log.T) {
 	signal.Notify(c, sessionutil.ControlSignals...)
 	go func() {
 		<-c
-		fmt.Println("Terminate signal received, exiting.")
+		fmt.Fprintln(p.out, "Terminate signal received, exiting.")
 
 		if version.DoesAgentSupportTerminateSessionFlag(log, p.session.DataChannel.GetAgentVersion()) {
 			if err := p.session.DataChannel.SendFlag(log, message.TerminateSession); err != nil {
 				log.Errorf("Failed to send TerminateSession flag: %v", err)
 			}
-			fmt.Fprintf(os.Stdout, "\n\nExiting session with sessionId: %s.\n\n", p.sessionId)
+			fmt.Fprintf(p.out, "\n\nExiting session with sessionId: %s.\n\n", p.sessionId)
 			p.Stop()
 		} else {
 			p.session.TerminateSession(log)

src/sessionmanagerplugin/session/portsession/basicportforwarding_test.go

@@ -48,6 +48,7 @@ func TestSetSessionHandlers(t *testing.T) {
 		Session:        mockSession,
 		portParameters: PortParameters{PortNumber: "22", Type: "LocalPortForwarding"},
 		portSessionType: &BasicPortForwarding{
+			out:            os.Stdout,
 			session:        mockSession,
 			portParameters: PortParameters{PortNumber: "22", Type: "LocalPortForwarding"},
 		},
@@ -84,6 +85,7 @@ func TestStartSessionTCPLocalPortFromDocument(t *testing.T) {
 		Session:        getSessionMock(),
 		portParameters: PortParameters{PortNumber: "22", Type: "LocalPortForwarding", LocalPortNumber: "54321"},
 		portSessionType: &BasicPortForwarding{
+			out:            os.Stdout,
 			session:        getSessionMock(),
 			portParameters: PortParameters{PortNumber: "22", Type: "LocalPortForwarding"},
 		},
@@ -101,6 +103,7 @@ func TestStartSessionTCPAcceptFailed(t *testing.T) {
 		Session:        getSessionMock(),
 		portParameters: PortParameters{PortNumber: "22", Type: "LocalPortForwarding"},
 		portSessionType: &BasicPortForwarding{
+			out:            os.Stdout,
 			session:        getSessionMock(),
 			portParameters: PortParameters{PortNumber: "22", Type: "LocalPortForwarding"},
 		},
@@ -117,6 +120,7 @@ func TestStartSessionTCPConnectFailed(t *testing.T) {
 		Session:        getSessionMock(),
 		portParameters: PortParameters{PortNumber: "22", Type: "LocalPortForwarding"},
 		portSessionType: &BasicPortForwarding{
+			out:            os.Stdout,
 			session:        getSessionMock(),
 			portParameters: PortParameters{PortNumber: "22", Type: "LocalPortForwarding"},
 		},

src/sessionmanagerplugin/session/portsession/muxportforwarding.go

@@ -61,6 +61,7 @@ type MuxPortForwarding struct {
 	session        session.Session
 	muxClient      *MuxClient
 	mgsConn        *MgsConn
+	out            io.Writer
 }
 
 func (c *MgsConn) close() {
@@ -131,7 +132,7 @@ func (p *MuxPortForwarding) WriteStream(outputMessage message.ClientMessage) err
 		binary.Read(buf, binary.BigEndian, &flag)
 
 		if message.ConnectToPortError == flag {
-			fmt.Printf("\nConnection to destination port failed, check SSM Agent logs.\n")
+			fmt.Fprintf(p.out, "\nConnection to destination port failed, check SSM Agent logs.\n")
 		}
 	}
 	return nil
@@ -190,12 +191,12 @@ func (p *MuxPortForwarding) handleControlSignals(log log.T) {
 	signal.Notify(c, sessionutil.ControlSignals...)
 	go func() {
 		<-c
-		fmt.Println("Terminate signal received, exiting.")
+		fmt.Fprintln(p.out, "Terminate signal received, exiting.")
 
 		if err := p.session.DataChannel.SendFlag(log, message.TerminateSession); err != nil {
 			log.Errorf("Failed to send TerminateSession flag: %v", err)
 		}
-		fmt.Fprintf(os.Stdout, "\n\nExiting session with sessionId: %s.\n\n", p.sessionId)
+		fmt.Fprintf(p.out, "\n\nExiting session with sessionId: %s.\n\n", p.sessionId)
 		p.Stop()
 	}()
 }
@@ -252,10 +253,10 @@ func (p *MuxPortForwarding) handleClientConnections(log log.T, ctx context.Conte
 	defer listener.Close()
 
 	log.Infof(displayMsg)
-	fmt.Printf(displayMsg)
+	fmt.Fprintf(p.out, displayMsg)
 
 	log.Infof("Waiting for connections...\n")
-	fmt.Printf("\nWaiting for connections...\n")
+	fmt.Fprintf(p.out, "\nWaiting for connections...\n")
 
 	var once sync.Once
 	for {
@@ -269,7 +270,7 @@ func (p *MuxPortForwarding) handleClientConnections(log log.T, ctx context.Conte
 				log.Infof("Connection accepted from %s\n for session [%s]", conn.RemoteAddr(), p.sessionId)
 
 				once.Do(func() {
-					fmt.Printf("\nConnection accepted for session [%s]\n", p.sessionId)
+					fmt.Fprintf(p.out, "\nConnection accepted for session [%s]\n", p.sessionId)
 				})
 
 				stream, err := p.muxClient.session.OpenStream()

src/sessionmanagerplugin/session/portsession/portsession.go

@@ -15,6 +15,8 @@
 package portsession
 
 import (
+	"os"
+
 	"github.com/aws/session-manager-plugin/src/config"
 	"github.com/aws/session-manager-plugin/src/jsonutil"
 	"github.com/aws/session-manager-plugin/src/log"
@@ -70,12 +72,14 @@ func (s *PortSession) Initialize(log log.T, sessionVar *session.Session) {
 				sessionId:      s.SessionId,
 				portParameters: s.portParameters,
 				session:        s.Session,
+				out:            os.Stdout,
 			}
 		} else {
 			s.portSessionType = &BasicPortForwarding{
 				sessionId:      s.SessionId,
 				portParameters: s.portParameters,
 				session:        s.Session,
+				out:            os.Stdout,
 			}
 		}
 	} else {

src/sessionmanagerplugin/session/portsession/test_portsession.go

@@ -15,6 +15,8 @@
 package portsession
 
 import (
+	"os"
+
 	"github.com/aws/session-manager-plugin/src/communicator/mocks"
 	"github.com/aws/session-manager-plugin/src/datachannel"
 	"github.com/aws/session-manager-plugin/src/log"
@@ -41,11 +43,13 @@ func getSessionMock() session.Session {
 }
 
 func getSessionMockWithParams(properties interface{}, agentVersion string) session.Session {
-	datachannel := &datachannel.DataChannel{}
+	out := os.Stdout
+	datachannel := &datachannel.DataChannel{Out: out}
 	datachannel.SetAgentVersion(agentVersion)
 
 	var mockSession = session.Session{
 		DataChannel: datachannel,
+		Out:         out,
 	}
 
 	mockSession.DataChannel.Initialize(mockLog, "clientId", "sessionId", "targetId", false)

src/sessionmanagerplugin/session/session.go

@@ -85,14 +85,15 @@ type Session struct {
 	SessionType           string
 	SessionProperties     interface{}
 	DisplayMode           sessionutil.DisplayMode
+	Out                   io.Writer
 }
 
-//startSession create the datachannel for session
+// startSession create the datachannel for session
 var startSession = func(session *Session, log log.T) error {
 	return session.Execute(log)
 }
 
-//setSessionHandlersWithSessionType set session handlers based on session subtype
+// setSessionHandlersWithSessionType set session handlers based on session subtype
 var setSessionHandlersWithSessionType = func(session *Session, log log.T) error {
 	// SessionType is set inside DataChannel
 	sessionSubType := SessionRegistry[session.SessionType]
@@ -203,7 +204,8 @@ func ValidateInputAndStartSession(args []string, out io.Writer) {
 		session.Endpoint = ssmEndpoint
 		session.ClientId = clientId
 		session.TargetId = target
-		session.DataChannel = &datachannel.DataChannel{}
+		session.DataChannel = &datachannel.DataChannel{Out: out}
+		session.Out = out
 
 	default:
 		fmt.Fprint(out, "Invalid Operation")
@@ -217,12 +219,12 @@ func ValidateInputAndStartSession(args []string, out io.Writer) {
 	}
 }
 
-//Execute create data channel and start the session
+// Execute create data channel and start the session
 func (s *Session) Execute(log log.T) (err error) {
-	fmt.Fprintf(os.Stdout, "\nStarting session with SessionId: %s\n", s.SessionId)
+	fmt.Fprintf(s.Out, "\nStarting session with SessionId: %s\n", s.SessionId)
 
 	// sets the display mode
-	s.DisplayMode = sessionutil.NewDisplayMode(log)
+	s.DisplayMode = sessionutil.NewDisplayMode(log, s.Out)
 
 	if err = s.OpenDataChannel(log); err != nil {
 		log.Errorf("Error in Opening data channel: %v", err)

src/sessionmanagerplugin/session/session_test.go

@@ -104,7 +104,7 @@ func TestValidateInputAndStartSessionWithWrongEnvVariableName(t *testing.T) {
 }
 
 func TestExecute(t *testing.T) {
-	sessionMock := &Session{}
+	sessionMock := &Session{Out: os.Stdout}
 	sessionMock.DataChannel = mockDataChannel
 	SetupMockActions()
 	mockDataChannel.On("Open", mock.Anything).Return(nil)
@@ -128,7 +128,7 @@ func TestExecute(t *testing.T) {
 }
 
 func TestExecuteAndStreamMessageResendTimesOut(t *testing.T) {
-	sessionMock := &Session{}
+	sessionMock := &Session{Out: os.Stdout}
 	sessionMock.DataChannel = mockDataChannel
 	SetupMockActions()
 	mockDataChannel.On("Open", mock.Anything).Return(nil)

src/sessionmanagerplugin/session/sessionhandler.go

@@ -129,7 +129,7 @@ func (s *Session) ResumeSessionHandler(log log.T) (err error) {
 		return
 	} else if s.TokenValue == "" {
 		log.Debugf("Session: %s timed out", s.SessionId)
-		fmt.Fprintf(os.Stdout, "Session: %s timed out.\n", s.SessionId)
+		fmt.Fprintf(s.Out, "Session: %s timed out.\n", s.SessionId)
 		os.Exit(0)
 	}
 	s.DataChannel.GetWsChannel().SetChannelToken(s.TokenValue)

src/sessionmanagerplugin/session/sessionhandler_test.go

@@ -16,13 +16,15 @@ package session
 
 import (
 	"fmt"
+	"os"
 	"testing"
 
 	wsChannelMock "github.com/aws/session-manager-plugin/src/communicator/mocks"
 	"github.com/aws/session-manager-plugin/src/config"
 	"github.com/aws/session-manager-plugin/src/datachannel"
 	dataChannelMock "github.com/aws/session-manager-plugin/src/datachannel/mocks"
 	"github.com/aws/session-manager-plugin/src/message"
+	"github.com/aws/session-manager-plugin/src/sessionmanagerplugin/session/sessionutil"
 	"github.com/stretchr/testify/mock"
 
 	"github.com/stretchr/testify/assert"
@@ -38,7 +40,7 @@ func TestOpenDataChannel(t *testing.T) {
 	mockDataChannel = &dataChannelMock.IDataChannel{}
 	mockWsChannel = &wsChannelMock.IWebSocketChannel{}
 
-	sessionMock := &Session{}
+	sessionMock := &Session{Out: os.Stdout}
 	sessionMock.DataChannel = mockDataChannel
 	SetupMockActions()
 	mockDataChannel.On("Open", mock.Anything).Return(nil)
@@ -51,7 +53,7 @@ func TestOpenDataChannelWithError(t *testing.T) {
 	mockDataChannel = &dataChannelMock.IDataChannel{}
 	mockWsChannel = &wsChannelMock.IWebSocketChannel{}
 
-	sessionMock := &Session{}
+	sessionMock := &Session{Out: os.Stdout}
 	sessionMock.DataChannel = mockDataChannel
 	SetupMockActions()
 
@@ -69,10 +71,12 @@ func TestProcessFirstMessageOutputMessageFirst(t *testing.T) {
 		Payload:     []byte("testing"),
 	}
 
-	dataChannel := &datachannel.DataChannel{}
+	dataChannel := &datachannel.DataChannel{Out: os.Stdout}
 	dataChannel.Initialize(logger, clientId, sessionId, instanceId, false)
 	session := Session{
+		Out:         os.Stdout,
 		DataChannel: dataChannel,
+		DisplayMode: sessionutil.NewDisplayMode(logger, os.Stdout),
 	}
 
 	session.ProcessFirstMessage(logger, outputMessage)