Browse Source

Merge pull request #68 from hashicorp/f-logger

Allow passing in a logger
Alex Dadgar 6 years ago
parent
commit
7221087c3d
2 changed files with 18 additions and 2 deletions
  1. 12 1
      mux.go
  2. 6 1
      session.go

+ 12 - 1
mux.go

@@ -3,6 +3,7 @@ package yamux
 import (
 	"fmt"
 	"io"
+	"log"
 	"os"
 	"time"
 )
@@ -30,8 +31,13 @@ type Config struct {
 	// window size that we allow for a stream.
 	MaxStreamWindowSize uint32
 
-	// LogOutput is used to control the log destination
+	// LogOutput is used to control the log destination. Either Logger or
+	// LogOutput can be set, not both.
 	LogOutput io.Writer
+
+	// Logger is used to pass in the logger to be used. Either Logger or
+	// LogOutput can be set, not both.
+	Logger *log.Logger
 }
 
 // DefaultConfig is used to return a default configuration
@@ -57,6 +63,11 @@ func VerifyConfig(config *Config) error {
 	if config.MaxStreamWindowSize < initialStreamWindow {
 		return fmt.Errorf("MaxStreamWindowSize must be larger than %d", initialStreamWindow)
 	}
+	if config.LogOutput != nil && config.Logger != nil {
+		return fmt.Errorf("both Logger and LogOutput may not be set, select one")
+	} else if config.LogOutput == nil && config.Logger == nil {
+		return fmt.Errorf("one of Logger or LogOutput must be set, select one")
+	}
 	return nil
 }
 

+ 6 - 1
session.go

@@ -86,9 +86,14 @@ type sendReady struct {
 
 // newSession is used to construct a new session
 func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
+	logger := config.Logger
+	if logger == nil {
+		logger = log.New(config.LogOutput, "", log.LstdFlags)
+	}
+
 	s := &Session{
 		config:     config,
-		logger:     log.New(config.LogOutput, "", log.LstdFlags),
+		logger:     logger,
 		conn:       conn,
 		bufRead:    bufio.NewReader(conn),
 		pings:      make(map[uint32]chan struct{}),