From d86385b825cc0a16c066bb29d626294d21fe24ae Mon Sep 17 00:00:00 2001
From: mo <thinkgo@aliyun.com>
Date: Sun, 19 Apr 2020 17:47:39 +0800
Subject: [PATCH] use logger interface

---
 logger.go       | 22 ++++++++++++++++++++++
 request_test.go |  4 ++--
 socks5.go       | 17 ++++++++---------
 socks5_test.go  |  2 +-
 4 files changed, 33 insertions(+), 12 deletions(-)
 create mode 100644 logger.go

diff --git a/logger.go b/logger.go
new file mode 100644
index 0000000..9bda141
--- /dev/null
+++ b/logger.go
@@ -0,0 +1,22 @@
+package socks5
+
+import (
+	"log"
+)
+
+type Logger interface {
+	Errorf(format string, arg ...interface{})
+}
+
+// 标准输出
+type Std struct {
+	*log.Logger
+}
+
+func NewLogger(l *log.Logger) *Std {
+	return &Std{l}
+}
+
+func (sf Std) Errorf(format string, args ...interface{}) {
+	sf.Logger.Printf("[E]: "+format, args...)
+}
diff --git a/request_test.go b/request_test.go
index 5465113..3b66147 100644
--- a/request_test.go
+++ b/request_test.go
@@ -52,7 +52,7 @@ func TestRequest_Connect(t *testing.T) {
 	s := &Server{config: &Config{
 		Rules:    PermitAll(),
 		Resolver: DNSResolver{},
-		Logger:   log.New(os.Stdout, "", log.LstdFlags),
+		Logger:   NewLogger(log.New(os.Stdout, "socks5: ", log.LstdFlags)),
 	}}
 
 	// Create the connect request
@@ -127,7 +127,7 @@ func TestRequest_Connect_RuleFail(t *testing.T) {
 	s := &Server{config: &Config{
 		Rules:    PermitNone(),
 		Resolver: DNSResolver{},
-		Logger:   log.New(os.Stdout, "", log.LstdFlags),
+		Logger:   NewLogger(log.New(os.Stdout, "socks5: ", log.LstdFlags)),
 	}}
 
 	// Create the connect request
diff --git a/socks5.go b/socks5.go
index 3d3808c..bd53b0f 100644
--- a/socks5.go
+++ b/socks5.go
@@ -4,9 +4,9 @@ import (
 	"bufio"
 	"context"
 	"fmt"
+	"io/ioutil"
 	"log"
 	"net"
-	"os"
 )
 
 const (
@@ -42,8 +42,8 @@ type Config struct {
 	BindIP net.IP
 
 	// Logger can be used to provide a custom log target.
-	// Defaults to stdout.
-	Logger *log.Logger
+	// Defaults to ioutil.Discard.
+	Logger Logger
 
 	// Optional function for dialing out
 	Dial func(ctx context.Context, network, addr string) (net.Conn, error)
@@ -79,7 +79,7 @@ func New(conf *Config) (*Server, error) {
 
 	// Ensure we have a log target
 	if conf.Logger == nil {
-		conf.Logger = log.New(os.Stdout, "", log.LstdFlags)
+		conf.Logger = NewLogger(log.New(ioutil.Discard, "socks5: ", log.LstdFlags))
 	}
 
 	server := &Server{
@@ -113,7 +113,6 @@ func (s *Server) Serve(l net.Listener) error {
 		}
 		go s.ServeConn(conn)
 	}
-	return nil
 }
 
 // ServeConn is used to serve a single connection.
@@ -124,14 +123,14 @@ func (s *Server) ServeConn(conn net.Conn) error {
 	// Read the version byte
 	version := []byte{0}
 	if _, err := bufConn.Read(version); err != nil {
-		s.config.Logger.Printf("[ERR] socks: Failed to get version byte: %v", err)
+		s.config.Logger.Errorf("[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)
-		s.config.Logger.Printf("[ERR] socks: %v", err)
+		s.config.Logger.Errorf("[ERR] socks: %v", err)
 		return err
 	}
 
@@ -139,7 +138,7 @@ func (s *Server) ServeConn(conn net.Conn) error {
 	authContext, err := s.authenticate(conn, bufConn)
 	if err != nil {
 		err = fmt.Errorf("Failed to authenticate: %v", err)
-		s.config.Logger.Printf("[ERR] socks: %v", err)
+		s.config.Logger.Errorf("[ERR] socks: %v", err)
 		return err
 	}
 
@@ -160,7 +159,7 @@ func (s *Server) ServeConn(conn net.Conn) error {
 	// Process the client request
 	if err := s.handleRequest(request, conn); err != nil {
 		err = fmt.Errorf("Failed to handle request: %v", err)
-		s.config.Logger.Printf("[ERR] socks: %v", err)
+		s.config.Logger.Errorf("[ERR] socks: %v", err)
 		return err
 	}
 
diff --git a/socks5_test.go b/socks5_test.go
index 8cfbee0..d1bf37d 100644
--- a/socks5_test.go
+++ b/socks5_test.go
@@ -43,7 +43,7 @@ func TestSOCKS5_Connect(t *testing.T) {
 	cator := UserPassAuthenticator{Credentials: creds}
 	conf := &Config{
 		AuthMethods: []Authenticator{cator},
-		Logger:      log.New(os.Stdout, "", log.LstdFlags),
+		Logger:      NewLogger(log.New(os.Stdout, "socks5: ", log.LstdFlags)),
 	}
 	serv, err := New(conf)
 	if err != nil {