Merge pull request armon#5 from svent/add-logger-to-config

Add configuration option for the log target

Author: armon

request.go

@@ -155,8 +155,8 @@ func (s *Server) handleConnect(conn conn, bufConn io.Reader, dest, realDest *Add
 
 	// Start proxying
 	errCh := make(chan error, 2)
-	go proxy("target", target, bufConn, errCh)
-	go proxy("client", conn, target, errCh)
+	go proxy("target", target, bufConn, errCh, s.config.Logger)
+	go proxy("client", conn, target, errCh, s.config.Logger)
 
 	// Wait
 	select {
@@ -301,13 +301,13 @@ func sendReply(w io.Writer, resp uint8, addr *AddrSpec) error {
 
 // proxy is used to suffle data from src to destination, and sends errors
 // down a dedicated channel
-func proxy(name string, dst io.Writer, src io.Reader, errCh chan error) {
+func proxy(name string, dst io.Writer, src io.Reader, errCh chan error, logger *log.Logger) {
 	// Copy
 	n, err := io.Copy(dst, src)
 
 	// Log, and sleep. This is jank but allows the otherside
 	// to finish a pending copy
-	log.Printf("[DEBUG] socks: Copied %d bytes to %s", n, name)
+	logger.Printf("[DEBUG] socks: Copied %d bytes to %s", n, name)
 	time.Sleep(10 * time.Millisecond)
 
 	// Send any errors

request_test.go

@@ -4,7 +4,9 @@ import (
 	"bytes"
 	"encoding/binary"
 	"io"
+	"log"
 	"net"
+	"os"
 	"strings"
 	"testing"
 )
@@ -50,6 +52,7 @@ func TestRequest_Connect(t *testing.T) {
 	s := &Server{config: &Config{
 		Rules:    PermitAll(),
 		Resolver: DNSResolver{},
+		Logger:   log.New(os.Stdout, "", log.LstdFlags),
 	}}
 
 	// Create the connect request
@@ -119,6 +122,7 @@ func TestRequest_Connect_RuleFail(t *testing.T) {
 	s := &Server{config: &Config{
 		Rules:    PermitNone(),
 		Resolver: DNSResolver{},
+		Logger:   log.New(os.Stdout, "", log.LstdFlags),
 	}}
 
 	// Create the connect request

socks5.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log"
 	"net"
+	"os"
 )
 
 const (
@@ -38,6 +39,10 @@ type Config struct {
 
 	// BindIP is used for bind or udp associate
 	BindIP net.IP
+
+	// Logger can be used to provide a custom log target.
+	// Defaults to stdout.
+	Logger *log.Logger
 }
 
 // Server is reponsible for accepting connections and handling
@@ -68,6 +73,11 @@ func New(conf *Config) (*Server, error) {
 		conf.Rules = PermitAll()
 	}
 
+	// Ensure we have a log target
+	if conf.Logger == nil {
+		conf.Logger = log.New(os.Stdout, "", log.LstdFlags)
+	}
+
 	server := &Server{
 		config: conf,
 	}
@@ -110,28 +120,28 @@ func (s *Server) ServeConn(conn net.Conn) error {
 	// Read the version byte
 	version := []byte{0}
 	if _, err := bufConn.Read(version); err != nil {
-		log.Printf("[ERR] socks: Failed to get version byte: %v", err)
+		s.config.Logger.Printf("[ERR] socks: Failed to get version byte: %v", err)
 		return err
 	}
 
 	// Ensure we are compatible
 	if version[0] != socks5Version {
 		err := fmt.Errorf("Unsupported SOCKS version: %v", version)
-		log.Printf("[ERR] socks: %v", err)
+		s.config.Logger.Printf("[ERR] socks: %v", err)
 		return err
 	}
 
 	// Authenticate the connection
 	if err := s.authenticate(conn, bufConn); err != nil {
 		err = fmt.Errorf("Failed to authenticate: %v", err)
-		log.Printf("[ERR] socks: %v", err)
+		s.config.Logger.Printf("[ERR] socks: %v", err)
 		return err
 	}
 
 	// Process the client request
 	if err := s.handleRequest(conn, bufConn); err != nil {
 		err = fmt.Errorf("Failed to handle request: %v", err)
-		log.Printf("[ERR] socks: %v", err)
+		s.config.Logger.Printf("[ERR] socks: %v", err)
 		return err
 	}
 

socks5_test.go

@@ -4,7 +4,9 @@ import (
 	"bytes"
 	"encoding/binary"
 	"io"
+	"log"
 	"net"
+	"os"
 	"testing"
 	"time"
 )
@@ -38,9 +40,10 @@ func TestSOCKS5_Connect(t *testing.T) {
 	creds := StaticCredentials{
 		"foo": "bar",
 	}
-	cator := UserPassAuthenticator{Credentials : creds}
+	cator := UserPassAuthenticator{Credentials: creds}
 	conf := &Config{
-		AuthMethods : []Authenticator{cator},
+		AuthMethods: []Authenticator{cator},
+		Logger:      log.New(os.Stdout, "", log.LstdFlags),
 	}
 	serv, err := New(conf)
 	if err != nil {